Skip to content

tfg: avoid NaN gradients for unbounded distances - #339

Open
dajiaohuang wants to merge 1 commit into
bytedance:mainfrom
dajiaohuang:fix/336-tfg-infinite-bound-gradient
Open

tfg: avoid NaN gradients for unbounded distances#339
dajiaohuang wants to merge 1 commit into
bytedance:mainfrom
dajiaohuang:fix/336-tfg-infinite-bound-gradient

Conversation

@dajiaohuang

Copy link
Copy Markdown

Summary

  • clamp lower- and upper-bound violations before squaring them in the flat-bottom parabolic potential
  • preserve existing finite-bound energies and analytic gradients
  • add regression coverage for finite bounds and the real clash-only PairwiseDistancePotential autograd path

Root cause

Clash-only pairwise constraints use upper = +inf. The inactive upper-bound branch previously computed (value - upper) ** 2 before masking it with torch.where. Although the forward energy remained finite, autograd propagated a NaN gradient from the infinite intermediate.

Clamping the violation before squaring directly implements the helper's documented max(0, violation) formula and avoids constructing that infinite inactive branch. This does not change public APIs or finite-bound behavior.

Fixes #336.

Validation

  • python -m pytest tests/test_tfg_potentials.py -q (2 passed)
  • the same regression file against unmodified main (1 failed, 1 passed; the real caller produced all-NaN autograd gradients)
  • exact comparison with the previous implementation for 101 finite-bound values in float32 and float64 on CPU
  • infinite lower/upper-bound autograd versus analytic gradients in float32 and float64 on CPU and CUDA
  • uvx pre-commit run --files protenix/tfg/potentials.py tests/test_tfg_potentials.py (all hooks passed)
  • python -m compileall -q protenix/tfg/potentials.py tests/test_tfg_potentials.py
  • git diff --check upstream/main...HEAD

The full test suite was not run locally because collection in the isolated Windows test environment stops on missing optional/runtime packages (optree, Biopython, Biotite, and ml_collections). The focused test imports the real RDKit-backed module and passed with RDKit 2025.09.3, PyTorch 2.12.0.dev20260408+cu128, Python 3.13.13, and CUDA available.

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.

PairwiseDistancePotential returns NaN autograd gradients for clash-only bounds

1 participant