Skip to content

DESIGN.md

Latest commit

 

History

History
433 lines (337 loc) · 19.9 KB

DESIGN.md

File metadata and controls

433 lines (337 loc) · 19.9 KB
Raw

Inference-time chemical steering: how it works, and how to add it to OpenFold3

Reference notes from reading three codebases at the commits below, plus a working prototype (foldsteer/) that implements the design.

source commit / version license
jwohlwend/boltz b1ebfc4, v2.2.1 MIT
bytedance/Protenix 4c355be Apache-2.0
aqlaboratory/openfold-3 f5df7b8, v0.4.x Apache-2.0

1. The problem

AF3-style models place atoms with a diffusion process trained on a coordinate loss. Nothing in that objective knows what a chemical bond is. The result is a characteristic failure mode: a ligand can land within 2 Å RMSD of the reference pose and still be chemically impossible — an inverted stereocenter, a twisted amide, a distorted aromatic ring, atoms overlapping across a chain interface.

The fix is not retraining. It is to inject chemical knowledge at sampling time, using the fact that the diffusion sampler predicts a clean structure x0 at every step and there is nothing stopping you from editing it before the sampler takes its next step.

Two mechanisms do this, and both Boltz-1x and Protenix-v2 implement both.

2. Mechanism 1 — physical guidance

Define an energy E(x) that is zero for chemically valid geometry and positive otherwise. After the denoiser predicts x0, run a few steps of gradient descent on E and add the correction:

x0  ->  x0 - Σ_k  w_k(t) · ∇E_k(x0)

In Boltz (model/modules/diffusion.py, ~line 608) this is an inner loop of num_gd_steps (default 20) applied to atom_coords_denoised, after which the sampler's ordinary Euler step proceeds unchanged. Guidance edits the denoised prediction, never the noisy state — that is what keeps it compatible with the sampler's update rule.

Every term is a flat-bottom restraint: zero energy inside a permitted interval, linear penalty outside it. This choice does real work:

  • the gradient magnitude is bounded, so a badly violated constraint cannot blow up the update the way a harmonic term would;
  • a structure that is already valid gets exactly zero energy and zero gradient, so guidance is inert on predictions the model already got right.

The second property is the one that makes this safe to leave on by default. The prototype asserts it as a test: a valid conformer produces max_energy = 0.0 and max_update = 0.0.

The terms

term geometric variable catches
bounds matrix interatomic distance bond lengths, 1-3 angles, internal clashes
VDW overlap interatomic distance steric clash across chains
chirality improper dihedral inverted stereocenters
stereo bond |dihedral| E/Z flips
planarity (sp2) |improper| distorted aromatic rings
non-planarity (sp3) |improper| flattened sp3 centers
conjugated torsion |sin(dihedral)| twisted amides
connections distance covalently linked chains drifting apart

Bounds come from RDKit's GetMoleculeBoundsMatrix with triangle smoothing — the same distance-geometry machinery used for conformer embedding, so the restraint targets a real conformer ensemble rather than one idealized geometry. Buffers (12.5% on bonds, 30° on stereocenters) loosen the intervals so normal conformational spread is unpenalized.

Chirality deserves special mention: it is the one property here that cannot be repaired by post-hoc relaxation, because inverting a center requires passing through a planar transition state. It has to be enforced during sampling.

3. Mechanism 2 — Feynman-Kac steering

Guidance is local: it descends into the nearest valid geometry. If the model committed to a bad binding mode early, no amount of gradient descent recovers.

FK steering runs num_particles trajectories at once and periodically culls them. At each checkpoint, particles are scored and resampled with replacement in proportion to softmax(λ · log G).

The important detail is the choice of potential:

log_G = energy_traj[:, -2] - energy_traj[:, -1]     # improvement, not level

log G is the energy decrease since the last checkpoint, not the raw energy. Rewarding improvement rather than absolute score keeps the population from collapsing onto whichever particle started in an easy basin.

When guidance is also active, the resampling weights include a correction:

ll_difference = (eps² − (eps + scaled_guidance_update)²).sum() / (2·noise_var)

This accounts for guidance having already biased the transition kernel away from the model's own distribution — without it, the particle weights would double-count the steering that guidance already applied.

Cost is linear in particle count: 3 particles ≈ 3× the denoiser calls. This is the expensive half of steering, and the reason Boltz gates it behind a separate flag from guidance.

4. What Protenix-v2 adds

Protenix ships steering as its own subpackage (protenix/tfg/, 2360 lines: config.py, engine.py, potentials.py) with a class registry, config-driven term construction, and a single integration point in protenix/model/generator.py. That is the strongest available evidence that the "separate activatable package" you are asking for is the right shape — a second team building this independently converged on it.

Three substantive additions beyond Boltz:

Projection (SHAKE-like). Instead of only descending the gradient, solve a linearized constraint projection for the minimum-norm correction:

dx = −Jᵀ(JJᵀ)⁻¹ v

for the active constraint set (_solve_constraint_projection, potentials.py:443). This converges in far fewer iterations than gradient descent for tight constraints like bond lengths, at the cost of a linear solve per batch element. Protenix runs it as nested loops (projection_outer=2, projection_inner=10) ordered chirality → distances.

Denoiser-path guidance (rho). Optionally backpropagate ∇log p(x0) through the denoiser to get a gradient on x_t rather than x0. Principled, but it requires a differentiable forward pass through the network. It defaults to rho=0.0 — off — and I would leave it off: the memory cost of taping an AF3-scale denoiser is severe and the reported gains come from the other terms.

A stricter validity definition. This is the finding most worth acting on. Protenix-v2 reports that models can score well on standard PoseBusters checks while remaining chemically implausible, and extends PXMeter (v1.1.0) with checks on sp2-center planarity, amide planarity, and sp3 non-planarity. Their Table 1 shows Protenix-v2-TFG at 60.46% joint RMSD-and-validity success under the revised criterion versus Boltz-1x at 53.96%, and Figure 7 identifies sp2 planarity and sp3 non-planarity as the checks with the largest spread between steered and unsteered models.

Boltz's PlanarBondPotential keys off BondType.DOUBLE. RDKit types aromatic bonds as AROMATIC, and an amide C–N as SINGLE. So a direct port of Boltz's term covers neither aromatic rings nor amides — exactly the two failure classes Protenix-v2 calls out. The prototype adds three terms for this (Sp2CenterPotential, Sp3CenterPotential, ConjugatedTorsionPotential) and verifies each against a deliberately constructed distortion.

Note the amide case needs a different geometric variable. Planarity there is a torsion, not an improper — after hydrogens are stripped a secondary amide N has only two heavy neighbours, so no improper exists — and it is satisfied at both 0 and π. A flat-bottom interval on the torsion cannot express "near 0 or near π"; taking |sin(φ)| collapses both into one restraint near zero.

5. OpenFold3: where steering attaches

5.1 The sampling loop

SampleDiffusion._sample_rollout in openfold3/core/model/structure/diffusion_module.py (~line 289):

for tau, c_tau in enumerate(noise_schedule[1:]):
    xl = centre_random_augmentation(xl=xl, atom_mask=atom_mask)
    gamma = self.gamma_0 if c_tau > self.gamma_min else 0
    t = noise_schedule[tau] * (gamma + 1)
    noise = self.noise_scale * sqrt(t² − noise_schedule[tau]²) * randn_like(xl)
    xl_noisy = xl + noise
    xl_denoised = self.diffusion_module(...)      # <-- guidance goes here
    delta = (xl_noisy − xl_denoised) / t
    xl = xl_noisy + self.step_scale * (c_tau − t) * delta

Structurally identical to Boltz's loop. Guidance inserts between xl_denoised and delta; FK resampling inserts immediately after, reindexing xl, xl_noisy, and xl_denoised together.

Two OF3-specific details:

  • Coordinate layout is [B, S, N, 3] — batch, rollout samples, atoms. The engine works on [*B, N, 3], so flatten the leading dims and restore after. The rollout-sample axis doubles as the FK particle axis; the final resampling step collapses each group to one survivor.
  • centre_random_augmentation runs at the top of every step. A guidance update computed at step n is expressed in a frame that no longer applies at step n+1. Boltz handles this by explicitly rotating its stored scaled_guidance_update by the same random rotation. The simplest correct approach is to recompute each step and use the stored update only within the step that produced it.

5.2 Where the chemistry comes from — the part that makes this feasible

Steering needs bonds, elements, and stereo assignments. In OF3 these are already present at sampling time:

  • single_datasets/inference.py:318 attaches the biotite AtomArray to the feature dict as features["atom_array"], with a comment calling it a pseudo-feature pending a cleaner mechanism.
  • core/utils/tensor_utils.py:66 — dict_multimap special-cases AtomArray | str and passes them through as a plain list instead of trying to pad-stack them.

So batch["atom_array"] survives collation into the model, carrying the BondList (with bond orders), elements, and chain ids. That is enough to rebuild an RDKit Mol per ligand chain and derive every constraint. No new feature has to be threaded through the data pipeline, which is what would otherwise make this a large change.

The prototype's _mols_from_atom_array does this reconstruction. It restricts to ligand chains — polymer geometry is already well handled by the network, and Boltz likewise scopes its bounds-matrix term to non-polymer entities.

Three details in that reconstruction only became apparent against real OF3 input, and each failed silently rather than raising:

  1. molecule_type_id holds MoleculeType integers (LIGAND == 3), not names. Selecting ligands by string dropped every one of them, leaving an empty constraint set and steering that did nothing.
  2. biotite bond orders 5 and 6 are AROMATIC_SINGLE and AROMATIC_DOUBLE, not two spellings of AROMATIC. Mapping both onto RDKit's AROMATIC throws the Kekulé structure away and sanitization then fails on nearly every aromatic ligand — which the extractor logs and skips, so again: no constraints.
  3. Stereo must be perceived from ref_pos, OF3's reference-conformer feature (RDKit's own generated conformer per component). atom_array.coord is the obvious candidate and the wrong one: at inference it holds no meaningful coordinates, because it is what the model is about to predict.

The lesson generalizes past OF3 — every one of these turns steering into a no-op without an error — so the adapter counts what it found (STATS) and tests/test_of3_extraction.py pins all three.

5.3 The precedent to imitate

OF3 already has a feature shaped exactly like steering: pocket constraints.

  • data/pipelines/featurization/pocket_constraints.py builds sampler-only tensors (pocket_sampling_vdw_radii, ligand masks, RDKit conformers) and returns {} when the feature is not requested;
  • model/structure/pocket_constraints.py consumes them;
  • SampleDiffusion.forward calls _pocket_sampling_enabled(batch) and runs a second partial-diffusion rollout when set.

It even computes VDW radii and generates RDKit conformers already. Steering is the same shape of change, which is good evidence the seam is real rather than being forced.

6. Package design

The requirement was that steering work independently of OpenFold3 — installable and activatable on its own. The design that achieves this is a two-sided contract with a single data structure in the middle.

      host model                foldsteer                    host model
  ┌────────────────┐      ┌──────────────────────┐      ┌───────────────┐
  │ batch /        │─────▶│  ChemicalContext     │─────▶│ sampling loop │
  │ AtomArray /    │ adapt│  (atom-indexed       │ hook │  x0 += guide  │
  │ RDKit mols     │      │   constraint tensors)│      │  idx = resamp │
  └────────────────┘      └──────────────────────┘      └───────────────┘

ChemicalContext holds only atom-indexed tensors in the host's atom ordering. It knows nothing about MSAs, trunk embeddings, tokens, or checkpoints. The engine consumes it and exposes two pure functions:

update = engine.guide(x0, t)                  # tensor -> tensor
idx    = engine.resample_indices(x0, t, ...)  # -> integer indices

The engine never sees the denoiser, the noise schedule, or the model. That is what makes it host-agnostic: adding a second host means writing one adapter function, not touching the physics.

Concretely, the package imports nothing from OpenFold3 at module scope — the adapter imports it lazily inside the patch function. Verified:

openfold3 imported at module scope: False
9 passed

The whole test suite runs with no folding model installed and no GPU.

Activation

from foldsteer.adapters.openfold3 import patch_sample_diffusion
from foldsteer import default_config

unpatch = patch_sample_diffusion(model, default_config(num_particles=3))
...                        # run inference as normal
unpatch()                  # restore

Monkey-patching the bound method rather than subclassing is deliberate: OF3 builds SampleDiffusion inside the model's __init__ from a config dict, so substituting the class means rebuilding the model. Patching the method leaves checkpoint loading untouched and makes the change reversible within a process.

That form needs the model object, which under OF3's own runner does not exist yet: Lightning constructs the model inside InferenceExperimentRunner, so there is no seam to reach the sampler beforehand. patch_sample_diffusion_class(config) closes the gap by wrapping SampleDiffusion.__init__, so every sampler Lightning builds afterwards is handed to the same instance-level patch above. The cost is that the CLI must be invoked in-process (cli.main(["predict", ...], standalone_mode=False)) — a subprocess builds its model in a different interpreter and runs unsteered, with no error. examples/run_of3_steered.py is the whole activation path, and the STATS counters exist so a silently unsteered run is visible.

For upstreaming rather than bolting on, the same engine calls would go directly into _sample_rollout behind a steering_config argument defaulting to None — about 15 lines at two call sites.

7. Validation

Everything below is measured by the prototype, not asserted.

Gradients are correct. All eight potentials check against autograd across four molecules (alanine, a trans-amide, benzanilide, a dipeptide). Max absolute error ≈ 1e-15 — machine precision.

chiral_drug  BoundsMatrixPotential    K=15    err=8.9e-16
chiral_drug  ChiralAtomPotential      K=1     err=2.3e-15
aromatic     BoundsMatrixPotential    K=105   err=8.9e-16
complex      BoundsMatrixPotential    K=136   err=4.4e-16
complex      Sp2CenterPotential       K=3     err=7.2e-16
amide        ConjugatedTorsionPotential K=1   err=0.0e+00

Writing analytic gradients rather than using autograd was the single largest source of bugs — two broadcasting errors in the chain rule, both caught only by this comparison. If you implement this, write the autograd check first.

Both bugs were the same mistake: inserting the arity axis at -3 instead of -2 when scattering dE/dv · dv/dx back to atom space. It is worth understanding why this is dangerous rather than merely wrong. de_dv is [*B, K] and dv_dx is [*B, A, K, 3], where the arity A is 2 for distance terms and 4 for dihedrals. unsqueeze(-3) on a single-batch tensor yields (1, B, K, 1) instead of the required (B, 1, K, 1). When B == A this still broadcasts — silently, with wrong values and no exception. When B != A it raises, and for unbatched input -3 is outside torch's valid range entirely.

So the failure is load-bearing on batch size, which is exactly the kind of bug that survives a smoke test and reaches production. tests/test_broadcast_regression.py sweeps batch sizes that both do and do not collide with the arity, plus the unbatched and two-leading-dim ([B, S, N, 3], OF3's layout) cases. Reintroducing the bad axis fails 7 of its 8 cases — and passes the 8th, which is the collision case, confirming that a single-batch-size test would have missed it.

Valid geometry is untouched. An MMFF-optimized conformer gives max_energy = 0.0, max_update = 0.0.

The documented failure modes are caught. Each distortion built explicitly from a valid conformer:

distortion E(valid) E(distorted)
twisted amide (90°) 0.0000 0.6754
puckered aromatic ring 0.0000 1.2202
flattened sp3 center 0.0000 0.3487

End-to-end against a mock OF3 sampler replicating _sample_rollout's exact contract — same signature, same [B,S,N,3] layout, same augmentation and noise schedule — with a denoiser that returns a deliberately distorted ligand:

variant mean final strain (6 seeds)
unsteered 11.30
guidance only 5.74
FK only (3 particles) 9.28
guidance + FK (3 particles) 5.66

Guidance accounts for most of the reduction. Caveat: the mock denoiser adds independent Gaussian noise each step, so its particles differ by noise rather than occupying distinct binding modes — the situation FK steering exists to exploit. This benchmark therefore understates FK and should not be read as evidence against it; the published Boltz-1x and Protenix-v2 results, which use real denoisers, are the relevant evidence there.

8. Recommendations

  1. Start with guidance, not FK. Most of the validity gain, ~1.2× cost instead of 3×, and no change to the sample-count semantics.
  2. Port Boltz's default weights verbatim to begin with. They are the only published settings validated against PoseBusters at scale. Retune against a benchmark, not intuition.
  3. Include the sp2/sp3/amide terms from the start. They are cheap, and Protenix-v2's analysis says they are where steered and unsteered models differ most.
  4. Write the autograd gradient check before the potentials. See §7.
  5. Evaluate under the revised validity criterion, not just standard PoseBusters — otherwise you cannot see the failures §4 is about.
  6. Watch the augmentation frame. The most likely silent bug in an OF3 port is applying a stale guidance update after centre_random_augmentation has rotated the coordinates.

Not yet implemented

The prototype covers guidance, FK steering, and the eight potentials. Left out: Protenix's constraint projection (worth adding — better convergence on bond lengths for a linear solve per step), denoiser-path guidance (rho; skip), symmetric-chain COM and template restraints (in the context schema, no potentials written), and the real _mols_from_atom_array path, which is written against the OF3 AtomArray API but has been exercised only on synthetic arrays — it needs a run against genuine OF3 inference output before it can be trusted.