Mini-Fold — Protein Structure Diffusion Model

Why v13: The Beta-Sheet Problem

v12b achieved record structural metrics (FAPE 1.584, frame_rot 0.790) but visual inspection of generated structures revealed a systematic failure: the model learns alpha-helices well but cannot form beta-sheets. The root cause is the pair stack’s limited receptive field — with max_rel_pos=32, residue pairs separated by more than 32 positions in sequence have zero relative position signal. Beta-sheet hydrogen bonds typically connect residues 20–100+ positions apart, making them invisible to v12b’s pair representation.

v13 Overview (~35–38M params)

v13 makes three high-impact changes to solve the beta-sheet problem while preserving v12b’s proven IPA denoiser. The denoiser, aux pair stack, Rg predictor, and distance head are initialized from v12b b30 EMA weights (exact match via strict=False). The pair stack is entirely new architecture and trains from scratch.

1. Deeper Pair Stack with Triangle Multiplicative Updates (8 blocks)

The pair stack doubles from 4 to 8 EnhancedPairBlock blocks. Each block now applies Evoformer-style triangle multiplicative updates (outgoing + incoming) before the existing row/column axial attention + FFN. The triangle updates implement transitivity: “if residue i contacts residue k, and residue j contacts residue k, then i and j are structurally related.” This is exactly how beta-sheet topology is encoded — two strands share contacts through loop residues.

Implementation: each triangle update projects the pair representation to gate/value tensors (tri_mul_dim=64), computes einsum('bikd,bjkd->bijd') (outgoing) or 'bkid,bkjd->bijd' (incoming), then projects back to d_pair=128. The einsum is forced to fp32 to prevent fp16 overflow from the L-dimensional accumulation. Output is clamped to [-1e4, 1e4] and returned as a residual delta to avoid catastrophic cancellation.

2. Log-Scaled Relative Position Encoding (128 bins, max_sep=512)

Replaces the linear-clipped RPE (max_rel_pos=32, 65 bins) with 128 log-spaced bins covering separations from 0 to 512 residues. The encoding is sign-aware (separate embeddings for upstream and downstream). Bin spacing:

• Bins 0–8: Linear spacing (1-residue resolution) for helix contacts (i+3, i+4)
• Bins 8–128: Log-spaced for long-range sheet contacts (i+20 to i+512)

Additionally, 32-dimensional sinusoidal continuous RPE features (sin/cos encoding projected to d_pair) provide smooth interpolation between discrete bins.

3. Classifier-Free Guidance Training

10% of training batches (p_uncond=0.1) replace residue tokens with MASK tokens (token ID 1), preserving CLS/EOS/PAD structure so attention masks remain valid. This trains the model for both conditional and unconditional generation. At sampling time, CFG enables guided generation: ε = εuncond + w · (εcond − εuncond).

Critical bug found and fixed: Initially used PAD (token 0) as the null token. Since the attention mask is computed as ids.ne(PAD), this produced an all-False mask, creating degenerate pair representations that caused NaN. Using MASK (token 1) instead preserves valid masks.

v13 LR Schedule: Warmup → Constant → Cosine Decay

Three-phase schedule motivated by v12b’s finding that breakthroughs happen in narrow LR windows:

1. Warmup (3 epochs): Linear 0.01x → 1x base LR
2. Constant (20 epochs): Hold at peak LR=2e-5 (pair_stack at 3x = 6e-5)
3. Cosine decay (~37 epochs): Slow decay to eta_min=1e-6

Per-module gradient clipping carried from v12b: denoiser max_norm=1.0, pair_stack=0.5.

v13 Numerical Stability Fixes

v13 development uncovered three classes of fp16 instability present since v12 (see NaN Debugging Story in the Structure Folding tab). All fixes are applied in both v12 and v13 codebases:

• fp16 epsilon underflow: clamp(min=1e-6) underflows to 0 in fp16 (min positive ~6e-5). Fixed by forcing fp32 in Gram-Schmidt, slerp, IPA point attention, and all loss functions.
• Triangle einsum overflow: L-dimensional accumulation exceeds fp16 max (65504). Forced fp32 with torch.amp.autocast("cuda", enabled=False).
• Unsafe torch.cdist backward: Produces NaN when distance is exactly 0. Replaced with manual (diff.pow(2).sum(-1) + 1e-10).sqrt().

IPA Denoiser (8 layers in v12/v13, 6 in v11)

The core structure module, adapted from AlphaFold2. Each IPA block performs three operations:

1. Invariant Point Attention — multi-head attention on the single representation, augmented with pair bias and 3D point attention. Each head generates query/key/value points in \(\mathbb{R}^3\) transformed into each residue’s local frame. Attention weights depend on geometric distances between learned points — invariant to global rotation/translation. Why it matters: EGNN (v8–v10) had no concept of local frames and could not optimize frame-aligned metrics like FAPE. IPA was the architectural pivot that broke through the topology ceiling.
2. Transition MLP — 2-layer feedforward on the single representation.
3. Frame update — predicts a small quaternion + translation update per residue, composed in the local frame (right-multiplication for SE(3) equivariance). Initialized near-zero so frames are approximately preserved in early training.

SNR-Gated Frame Initialization

Per-residue rigid frames are built from noised C\(\alpha\) coordinates via Gram-Schmidt orthogonalization on consecutive backbone triplets. At high noise (SNR < 1.0, roughly t > 700), the noised coordinates are near-isotropic and Gram-Schmidt is numerically unstable. The frame confidence mechanism smoothly blends toward identity frames:

Why it matters: Without this, the first IPA layer receives garbage frames at high noise levels, producing cascading errors through all 8 layers. This was a major source of training instability in early v11.

Frame-Aware Self-Conditioning (SC=0.25 in v13)

25% of training steps (50% in v11): run a no-grad forward pass to get x\(_{0}^{\text{prev}}\), build clean frames from it (treated as t=0), and use those as the initial frames for the second pass. At high noise where x\(_t\) frames are identity, self-conditioning provides the model’s best guess at clean local geometry — the IPA layers refine good frames instead of building them from scratch.

Why reduced to 0.25 in v13: Higher SC probability means fewer “cold start” training steps. The model needs enough cold-start experience to generalize at inference when no previous prediction exists.

Fixed-Scale Coordinates (introduced v11)

All coordinates divided by a fixed constant (10\(\text{\AA}\)) instead of per-protein R\(_g\). Why: R\(_g\) normalization made the noise schedule protein-size-dependent — a protein with R\(_g\)=5\(\text{\AA}\) had coordinate values ~1.0 while R\(_g\)=25\(\text{\AA}\) gave ~0.2–0.5, meaning the same noise level destroyed more signal for larger proteins. This silently capped TM-scores and looked like a “plateau” rather than a systematic bias. All successful protein diffusion models (FrameDiff, RFDiffusion, Genie) use fixed-scale coordinates.

1. FAPE (Frame-Aligned Point Error)  \(w = 1.0\)

The primary metric. Measures local structural consistency by computing point error in each residue’s local coordinate frame:

Why it’s the hardest loss: Requires global structural correctness, not just local geometry. Random baseline ~1.31; drops below 1.0 only when the model learns correct fold topology. v11b’s key fix: use learned R\(_{\text{pred}}\) from IPA layers instead of Gram-Schmidt rebuilt frames, giving rotations direct gradient for the first time.

2. Frame Rotation Loss  \(w = 0.5\) (introduced v11b)

Direct angular distance between learned rotation matrices and Gram-Schmidt ground truth: \(\mathcal{L}_{\text{rot}} = 1 - \cos\theta\), where \(\theta\) is the rotation angle between R\(_{\text{pred}}\) and R\(_{\text{GT}}\). Random baseline ~1.0 (~90\(^\circ\)), target <0.5.

Why it was added: v11 had a critical bug where x0_pred = t_vec discarded learned R. FAPE rebuilt frames via Gram-Schmidt, so rotations got no direct gradient. Adding this loss was the single change that enabled topology learning.

3. Distance MSE  \(w = 1.0\)

MSE on all pairwise C\(\alpha\) distances. Random baseline ~0.54. The easiest structural loss — even v8’s EGNN could reduce this 60% below random. But distance alone cannot encode topology (many different folds have similar distance distributions).

4. Bond Geometry  \(\beta(e) = \min(3.0,\; 1.0 + 2.0 \cdot \min(e/10, 1))\)

MSE on consecutive C\(\alpha\)–C\(\alpha\) distances vs ideal 3.8\(\text{\AA}\). Weight is annealed from 1.0 to 3.0 over 10 epochs. Random baseline ~0.17; below 0.02 = bonds within 0.1\(\text{\AA}\) of ideal.

Why annealed: Starting high prevents the model from learning global structure (it just chains beads at 3.8\(\text{\AA}\) apart). Starting low lets the model explore, then gradually enforces physical backbone geometry.

5. Chirality  \(w = 0.1\)

MSE on normalized signed volumes (scalar triple products) of C\(\alpha\) quartets. Ensures correct backbone handedness — natural proteins are L-amino acids with consistent chirality. Without this loss the model can generate mirror-image structures that score well on all other metrics.

6. Bond Angle  \(w = 0.5\)

MSE on cosines of C\(\alpha\)–C\(\alpha\)–C\(\alpha\) bond angles. Ideal angle ~120\(^\circ\) (\(\cos\theta \approx -0.5\)). Working in cosine space avoids discontinuities at 0\(^\circ\)/360\(^\circ\). Random baseline ~0.70; below 0.1 = correct backbone geometry.

7. Radius of Gyration  \(w = 0.5\)

MSE on log-transformed R\(_g\) predictions. A separate MLP predicts absolute R\(_g\) from sequence embeddings to recover real-space coordinates at inference. Converges below 0.05 by E2 and stays solved.

8. Auxiliary Distance CE  \(w = 0.03\)

Ordinal regression on binned pairwise distances (32 bins, 2–40\(\text{\AA}\)) from an independent lightweight pair stack (64-dim, ~500K params). The pair representation is detached — aux gradients train only the distance head, not the main pair stack.

Lesson learned: In v10 the distance head read from the main pair stack through detach(), causing feature drift divergence (aux_dist_ce: 3.95→38.8 over 3 epochs). v11+ uses a completely independent pair stack to avoid this.

DDIM Evaluation Metrics

Every 5 epochs (v13) or every epoch (v12b), we generate structures via 50-step DDIM sampling using EMA weights and evaluate against ground truth:

• TM-score: Global fold similarity [0, 1]. >0.17 = same fold, >0.5 = same topology. Random ~0.10.
• RMSD: Average atomic displacement after superposition. Random ~15–16\(\text{\AA}\). <5\(\text{\AA}\) = high quality.
• GDT: Fraction of residues within 1–8\(\text{\AA}\) of truth. Random ~3–4%.