dn6

/

RFDiffusion-3

Like 0

Builds on #1. That PR made the diffusion path produce real backbones; this one makes the sequence
design path usable. It contains #1's changes, so merging this alone lands both.

MPNN never ran

MPNNModel.forward built a flat dict and handed it to foundry, which reads everything from
network_input["input_features"]. Every call died with KeyError: 'input_features'. Nesting it
surfaced the rest of the contract: the network also validates decode_type, causality_pattern,
mask_for_loss, bias, pair_bias, symmetry_equivalence_group, symmetry_weight and
repeat_sample_num, and it wants a per-residue temperature rather than a scalar. The scalar
settings now match MPNN_PER_INPUT_INFERENCE_DEFAULTS in foundry's mpnn/utils/inference.py;
the optional conditioning keys are present as None, which is what the network checks for.

MPNN was designing against invented geometry

Even once it ran, the caller synthesised the backbone by adding three fixed offsets to every alpha
carbon in the global frame:

n_offset = torch.tensor([-1.458, 0.0, 0.0])
c_offset = torch.tensor([0.550, 1.424, 0.0])
o_offset = torch.tensor([0.550, 2.500, 0.0])

Every nitrogen was displaced in the same global direction regardless of backbone orientation, so
the graph features MPNN depends on were meaningless. After #1 the sampler has real atom-level
coordinates, so the block now selects N, CA, C and O through the is_backbone mask the feature
pipeline already emits, and the offsets are deleted.

Also fixed

• A latent shape bug from #1: the sequence-design block read motif_mask, which became atom-level,
  where it needs residue-level. It now reads a separate motif_token_mask.
• MPNN is attached with update_components and commonly left on CPU while the sampler ran on GPU.
  The block follows the component's device for the call and keeps everything after it on the
  sampler's device.

Verification

Run on an A10G. Bond lengths confirm the selected atoms really are N, CA, C, O in that order.

The CA slot matches the pipeline's own reported coordinates to 0.000000, which is what pins the
ordering rather than assuming it. A 100-residue design returns a full-length sequence over the
canonical alphabet. Sampling at 0.1 and 1.0 gives 5% sequence identity with composition shifting
from alanine and valine rich to diverse, confirming temperature reaches the sampler.

Known gap

Designs are not reproducible. RFDiffusionScheduler.add_noise uses torch.randn_like with no
generator, so passing generator= seeds only the initial noise and not the trajectory. Pre-existing
and untouched here.

No weight files are touched.