fix: correct deletion_value coefficient (was 4x too large) - #347
Conversation
deletion_value is meant to be (2/pi) * arctan(deletion_count / 3), per AF2/AF3. The coefficient was computed as 2.0 / torch.acos(torch.zeros(1, ...)) * 2, which due to operator precedence evaluates left-to-right as (2.0 / (pi/2)) * 2 == 8/pi instead of 2/pi -- exactly 4x too large. Verified numerically (coefficient 2.5465 vs intended 0.6366) and via a hand-computed value: for deletion_count=3, atan(1) = pi/4, so the correct deletion_value is exactly 0.5; the buggy code produced 2.0. Replaced the acos-based pi computation with torch.pi directly (already used elsewhere in this pipeline, e.g. input_embedders.py), removing the precedence trap rather than just parenthesizing around it. Add regression tests: a hand-computed exact value (deletion_count=3), and a comparison against an independently-implemented reference formula across a range of deletion counts. Both fail on the prior code.
|
hi there @Pana-TsK you are correct and thanks for this fix, this has been fixed in the upcoming release (the newly trained model was trained with the fix). I'll merge it here so that anybody training from scratch takes the benefit, for fine-tuning the impact should be minimal. Because this code-branch stays compatible with our P2 weights, I'll delay merging this (which will be a noop) until the upcoming release. |
|
Hey @jandom, thank you for your swift review! I really appreciate your time, and you and your teams' great work on OpenFold-3. It is an amazing piece of sofware. I have a few questions:
|
|
Our thanks are to you and other members of the community that help improve this codebase @Pana-TsK . Please accept our apologies, we should communicate better about the bugs that were already closed. I know Jennifer reached out to you, we'll debating some ways of doing this. Many thanks and we welcome further contributions from you, seems like you got a real one that's not fix – third time is the charm :P |
Summary
deletion_valueis meant to be(2/π)·arctan(deletion_count/3)per AF2/AF3. The coefficient was computed as2.0 / torch.acos(torch.zeros(1, ...)) * 2, which due to left-to-right operator precedence evaluates as(2.0/(π/2))*2 = 8/πinstead of2/π— exactly 4x too large. Verified numerically (coefficient 2.5465 vs. intended 0.6366) and with a hand-computed value: for deletion_count=3,atan(1) = π/4, so the correct output is exactly 0.5; the prior code produced 2.0.Fires on every MSA row with any deletion — real alignments almost always have some — so this affects the
deletion_valuechannel ofmsa_featfor essentially every training example and inference query.Changes
Replaced the
acos-based π computation withtorch.pidirectly (already used elsewhere in this pipeline, e.g.input_embedders.py:602), removing the precedence trap rather than just adding parentheses around it.Testing
Added regression tests in
openfold3/tests/core/data/pipelines/featurization/test_msa.py: a hand-computed exact value (deletion_count=3 → 0.5) that can't pass by coincidence, plus a comparison against an independently-implementedreference formula across a range of deletion counts. Both fail against the prior code.
Other notes
Same caveat as #346: this changes what the model sees as input, and I don't know what code snapshot the released checkpoint was trained from, so it may carry the same compatibility risk as #312.