Skip to content

ESMFold2: expose MSA inference-diversity options#3

Merged
faustomilletari merged 1 commit into
mainfrom
fm/esmfold2-msa-inference-diversity
Jun 8, 2026
Merged

ESMFold2: expose MSA inference-diversity options#3
faustomilletari merged 1 commit into
mainfrom
fm/esmfold2-msa-inference-diversity

Conversation

@faustomilletari

Copy link
Copy Markdown
Collaborator

What does this PR do?

Port the MSA inference-time diversity feature to the esmfold2 model:

  • modeling_esmfold2_common.py: add maybe_subsample_msa and maybe_apply_msa_column_masking helpers (operate on [B, depth, L] encoder inputs, always keep query row 0).
  • modeling_esmfold2.py + modeling_esmfold2_experimental.py: new forward kwargs msa_max_depth (1024) / msa_column_mask_rate (0.1) / msa_subsample_at_inference (True). Column masking is applied once before the recycling loop (shared across loops); row subsampling is drawn fresh per loop inside the loop, matching the pro-msa setup.

Defaults are hardcoded in the public signatures (not model config), so these are pure inference-time knobs. The diffusion sampler / num_sampling_steps schedule is untouched.

Fixes #
absence of inference-time diversity knobs in implementation

Before submitting

  • This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
  • Did you read the contributor guideline,
    Pull Request section?
  • Was this discussed/approved via a Github issue or the forum? Please add a link
    to it if that's the case.
  • Did you make sure to update the documentation with your changes? Here are the
    documentation guidelines, and
    here are tips on formatting docstrings.
  • Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

…mask)

Port the MSA inference-time diversity feature to the esmfold2 model:

- modeling_esmfold2_common.py: add maybe_subsample_msa and
  maybe_apply_msa_column_masking helpers (operate on [B, depth, L] encoder
  inputs, always keep query row 0).
- modeling_esmfold2.py + modeling_esmfold2_experimental.py: new forward kwargs
  msa_max_depth (1024) / msa_column_mask_rate (0.1) / msa_subsample_at_inference
  (True). Column masking is applied once before the recycling loop (shared
  across loops); row subsampling is drawn fresh per loop inside the loop,
  matching the pro-msa setup.

Defaults are hardcoded in the public signatures (not model config), so these
are pure inference-time knobs. The diffusion sampler / num_sampling_steps
schedule is untouched.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@faustomilletari faustomilletari merged commit ef32577 into main Jun 8, 2026
13 of 14 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants