[RFD3] Build real conditioning features instead of zero tensors

#1
by dn6 HF Staff - opened
Owner

The pipeline ran to completion and exited zero while producing coordinates that were not a
protein. Backbones came out with a mean CA-CA distance of 13.20 A against an expected 3.80,
and 112 pairs closer than 3 A.

Root cause

RFDiffusionTransformerModel.forward had a "simplified path", taken whenever the caller passed
no feature dict, that fed all-zero tensors for Q_L_init, C_L, P_LL, S_I and Z_II. The
denoise block never passed a feature dict, so that path always ran. Those tensors are learned
embeddings of real chemical and positional features, not optional inputs, so the model was being
evaluated on empty conditioning at every step.

Three loading bugs hid this. The scheduler could never load, because RFDiffusionScheduler
inherited only ConfigMixin and so had no from_pretrained. Both the noise-schedule block and
the denoise block silently substituted their own simplified samplers when it came back None.
And bfloat16, which the model card recommended, crashed inside foundry.

Changes

  • Build the genuine feature dict by driving foundry's own DesignInputSpecification and
    build_atom14_base_pipeline, and embed it through the token initializer. The zero-conditioning
    path is deleted rather than kept as a fallback.
  • Run the token initializer once per design via encode_conditioning, not once per step. It does
    not depend on noise level, and foundry calls it once.
  • RFDiffusionScheduler inherits SchedulerMixin, and strips the loader-only kwargs that
    ComponentSpec.load forwards because the index records it with type_hint=AutoModel.
  • Both blocks raise when a component is missing instead of degrading silently.
  • Half-precision weights raise with the remedy. foundry builds the per-atom noise level with a
    hardcoded .float(), so the scaled coordinates are float32 regardless of the caller's dtype and
    no cast on this side can fix it. README snippets now use float32.
  • Motif contigs and input_xyz raise. Fixed residues need per-atom element, atom-name and
    occupancy annotations that a coordinate tensor cannot supply. Documented as unsupported.

The sampler now runs on padded atom-level coordinates, 14 slots per token, and collapses to alpha
carbons at the end. state.output.xyz keeps its documented [1, 100, 3] shape.

Verification

Run on an A10G, calling the model card example verbatim.

Metric Before After Target
CA-CA mean 13.20 A 3.81 A 3.80 A
CA-CA range 1.14-30.66 3.77-3.89 tight
Bonds in 3.6-4.0 A 2% 100% > 95%
Clashes under 3 A 112 0 0
Radius of gyration 10.34 A 11.89 A ~12.7 A

The result is a mixed alpha-beta fold, 48% helical and 23% extended, 10.5 contacts per residue
within 8 A. A 60-residue design passes the same checks.

Not covered

MPNN sequence design still fails with KeyError: 'input_features'; foundry expects the tensors
nested under that key and a per-residue temperature. Separately, the caller fabricates N, C and O
positions by adding fixed global offsets to each alpha carbon, which ignores backbone orientation.
Real backbone atoms now exist internally, so that can be replaced properly in a follow-up.

No weight files are touched. The safetensors are byte-identical to main.

Owner

Verified on an A10G: mean CA-CA 3.81 A, 100% of bonds in range, zero clashes.

dn6 changed pull request status to open
dn6 changed pull request status to closed

Sign up or log in to comment