Skip to content

Fix chain center-of-mass loss weighting#571

Open
taivu1998 wants to merge 1 commit into
aqlaboratory:mainfrom
taivu1998:tdv/issue-540-chain-com-loss
Open

Fix chain center-of-mass loss weighting#571
taivu1998 wants to merge 1 commit into
aqlaboratory:mainfrom
taivu1998:tdv/issue-540-chain-com-loss

Conversation

@taivu1998
Copy link
Copy Markdown

Summary

Fixes #540 by removing the local weight multiplier from chain_center_of_mass_loss.

The chain center-of-mass loss was applying weight inside the squared clamp expression, and then AlphaFoldLoss applied loss.chain_center_of_mass.weight again during top-level aggregation. Because the local multiplier was inside the square, the default multimer contribution was effectively scaled by 0.05 ** 3 instead of the intended single top-level 0.05 coefficient.

Changes

  • Compute the raw chain center-of-mass penalty in chain_center_of_mass_loss.
  • Keep the weight argument in the function signature for config/backward compatibility, but document that top-level weighting is handled by AlphaFoldLoss.
  • Add deterministic CPU regression coverage that:
    • verifies the direct loss is independent of the weight argument
    • checks the expected raw loss value for a two-chain fixture
    • verifies gradients still flow through predicted CA coordinates
  • Add basic finite/nonnegative assertions to the existing optional chain COM smoke test.

Notes

This intentionally leaves config defaults and AlphaFoldLoss aggregation unchanged. It also leaves diagonal chain-pair normalization untouched, since that is a separate paper-faithfulness question from the weighting bug reported in #540.

Validation

  • PYTHONPATH=/private/tmp/openfold-issue540-stubs:/private/tmp/openfold-issue540-test-deps-py310 python3 -m pytest tests/test_loss.py -k chain_center_of_mass -q
    • 2 passed, 1 skipped, 21 deselected
  • PYTHONPATH=/private/tmp/openfold-issue540-stubs:/private/tmp/openfold-issue540-test-deps-py310 python3 -m py_compile openfold/utils/loss.py tests/test_loss.py
    • passed
  • PYTHONPATH=/private/tmp/openfold-issue540-stubs:/private/tmp/openfold-issue540-test-deps-py310 python3 -m pytest tests/test_loss.py -q
    • 1 failed, 6 passed, 17 skipped
    • failure is an existing local setup/resource issue: missing openfold/resources/stereo_chemical_props.txt for test_find_structural_violations
  • PYTHONPATH=/private/tmp/openfold-issue540-stubs:/private/tmp/openfold-issue540-test-deps-py310 python3 -m pytest tests/test_loss.py -k "not find_structural_violations" -q
    • 6 passed, 16 skipped, 2 deselected

@taivu1998 taivu1998 marked this pull request as ready for review May 11, 2026 03:44
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.

Potential bug in chain_center_of_mass_loss?

1 participant