Mini-Fold — Protein Structure Diffusion Model

The Topology Ceiling: Why We Abandoned Equivariant Graph Networks

We started with an 8-layer SE(3)-equivariant graph neural network (EGNN) as the denoiser. The appeal was obvious: EGNN is equivariant by construction, meaning it respects physical symmetries without needing data augmentation. After months of training across several iterations, we hit a hard wall. The model learned that “proteins are compact blobs of the right size” — distance MSE dropped 60% below random baseline, bond lengths converged — but FAPE stayed stuck at ~1.31 (the random baseline) and TM-scores never exceeded 0.131.

The root cause was architectural: EGNN passes messages based on pairwise distances and updates coordinates through distance-weighted vectors. It has no concept of local reference frames. FAPE measures frame-aligned point error — how well the structure looks from each residue’s perspective — and EGNN has no inductive bias to optimize this. The model could get global shape roughly right but could never learn which helix packs against which sheet, which loop connects where.

The IPA Pivot: Learning to Think in Local Frames

The breakthrough came when we replaced EGNN with Invariant Point Attention (IPA), the structure module from AlphaFold2 (Jumper et al., Nature 2021). Each of the 8 IPA layers 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.
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.

This was the single most important architectural decision in Mini-Fold’s development. IPA gave the model the ability to reason about local geometry — the exact thing FAPE demands. Within a few epochs, FAPE began dropping below the random baseline for the first time.

The Hidden Bug That Blocked Topology Learning

But IPA alone was not enough. Our first IPA-based model improved distances and bonds, then collapsed around epoch 14. After extensive debugging, we found a single-line code error: x0_pred = t_vec was silently discarding the learned rotation matrices from IPA. The FAPE loss was being computed using Gram-Schmidt-rebuilt frames from predicted coordinates, so the rotation parameters received zero gradient. The model could learn where atoms should be, but never learned how residues should be oriented.

The fix was twofold: use the IPA-predicted rotations directly in FAPE computation, and add an explicit frame rotation loss (\(\mathcal{L}_{\text{rot}} = 1 - \cos\theta\)) that provides direct angular supervision. This single change — adding rotation loss — was the most important insight in Mini-Fold’s development. It unlocked real topology learning: FAPE dropped to 1.655 and frame_rot to 0.830 within 8 epochs.

Scaling Up: From 8M to 31M Parameters

With the rotation bug fixed, the next bottleneck was capacity. The model showed gradient competition in the shared 128-dimensional single representation — the denoiser, pair stack, and auxiliary heads were all fighting for the same limited bandwidth. We scaled the model roughly 4x (8.4M to 31.1M parameters): d_ipa_hidden from 256 to 512, 8 attention heads, 8 query points, and a 2-layer FrameUpdate MLP. This resolved the gradient competition and set a new record: validation total loss of 2.413, FAPE 1.584, frame_rot 0.790.

Getting there required two training innovations. First, per-module gradient clipping: the denoiser clips at max_norm=1.0, the pair stack at 0.5. Without this, the pair stack’s gradients starved as the larger denoiser dominated updates. Second, a 3x learning rate multiplier for the pair stack, ensuring this smaller module could keep pace with the denoiser.

Why Alpha-Helices Were Easy and Beta-Sheets Were Impossible

After the scaling breakthrough, we had record structural metrics — but visual inspection of generated structures revealed a systematic failure. The model produced beautiful alpha-helices but could not form beta-sheets. The diagnosis pointed to one number: max_rel_pos=32.

Our pair stack encoded relative position as a linear-clipped embedding with 65 bins covering separations from -32 to +32 residues. Alpha-helix hydrogen bonds connect residues i and i+4 — well within this window. But beta-sheet hydrogen bonds typically connect residues 20 to 100+ positions apart in sequence. Any pair separated by more than 32 positions had zero relative position signal. The model was literally blind to the long-range contacts that define sheet topology.

Triangle Multiplicative Updates: Teaching the Model Transitivity

The current architecture (~35–38M params) makes three targeted changes to solve the beta-sheet problem. The proven IPA denoiser, auxiliary pair stack, R\(_g\) predictor, and distance head are all warm-started from the previous best checkpoint. The main pair stack is entirely new and trains from scratch.

The first change doubles the pair stack from 4 to 8 EnhancedPairBlock blocks, each now equipped with Evoformer-style triangle multiplicative updates (outgoing + incoming) before the existing row/column axial attention + FFN. The intuition behind triangle updates is 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 works — two strands share contacts through connecting 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.

Log-Scaled Position Encoding: Seeing the Whole Protein

The second change replaces the linear-clipped relative position encoding (65 bins, max separation 32) with 128 log-spaced bins covering separations from 0 to 512 residues. The encoding is sign-aware (separate embeddings for upstream and downstream). The bin spacing reflects a simple biological insight:

• 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)

Local structure needs fine-grained resolution; long-range contacts just need to be distinguishable from each other. Additionally, 32-dimensional sinusoidal continuous RPE features (sin/cos encoding projected to d_pair) provide smooth interpolation between discrete bins.

Classifier-Free Guidance: Conditional and Unconditional Generation

The third change adds classifier-free guidance (CFG) 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).

A subtle bug we caught early: our first implementation 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.

SNR-Gated Frame Initialization

Per-residue rigid frames are built from noised C\(\alpha\) coordinates via Gram-Schmidt orthogonalization on consecutive backbone triplets. We discovered early on that at high noise (SNR < 1.0, roughly t > 700), the noised coordinates are near-isotropic and Gram-Schmidt is numerically unstable. Without intervention, the first IPA layer receives garbage frames, producing cascading errors through all 8 layers. Our solution smoothly blends toward identity frames as noise increases:

This was a major source of training instability before we identified and fixed it.

Frame-Aware Self-Conditioning

25% of training steps (reduced from 50% in earlier iterations): 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.

We reduced SC probability from 50% to 25% after observing that higher rates meant fewer “cold start” training steps. The model needs enough cold-start experience to generalize at inference when no previous prediction exists.

Fixed-Scale Coordinates

All coordinates are divided by a fixed constant (10\(\text{\AA}\)) instead of per-protein R\(_g\). This seems like a minor choice, but it was one of the most impactful changes. 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 — once we switched, the plateau disappeared.

Learning Rate Schedule and Numerical Stability

The current learning rate schedule uses three phases, motivated by the observation that breakthroughs happen in narrow LR windows:

1. Warmup (3 epochs): Linear 0.01x to 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

During development we also uncovered three classes of fp16 instability (detailed in the NaN Debugging Story under the Structure Folding tab):

• 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().

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

Introduced by Jumper et al. in AlphaFold2 (Nature, 2021), FAPE is the gold standard loss for protein structure models. It measures how well predicted coordinates match ground truth in each residue’s local reference frame:

Unlike RMSD, which is a single global average, FAPE evaluates structure from every residue’s perspective. A model can have low RMSD by getting the overall shape roughly right, but FAPE requires getting the local neighborhoods correct — which helix packs against which sheet, which loop connects where.

Random baseline: ~1.31. Drops below 1.0 only when the model learns correct fold topology. This is the loss that distinguishes a blob from a protein.

Frame Rotation Loss  \(w = 0.5\)

Direct supervision on per-residue backbone orientations, measuring the angular distance between predicted and ground-truth rotation matrices: \(\mathcal{L}_{\text{rot}} = 1 - \cos\theta\). This is conceptually similar to the auxiliary heads in AlphaFold2’s structure module, but applied directly to the IPA-predicted frames.

Why this loss exists (the v11 rotation bug): In v11, a code error (x0_pred = t_vec) silently discarded the learned rotations from IPA. FAPE was computed using Gram-Schmidt-rebuilt frames, so rotations received zero gradient. The model could reduce distance-based losses but never learned orientations. Adding frame_rot was the single fix that unlocked topology learning — the most important architectural insight in Mini-Fold’s development.

Random baseline: ~1.0 (frames pointing ~90° off). Target: <0.5. Currently at 0.86–0.88 — the model is learning orientations but still has significant room to improve.

Distance MSE  \(w = 1.0\)

MSE on all pairwise C\(\alpha\) distances. Pairwise distance matrices (also called contact maps when thresholded) were the original representation used in early structure prediction methods (e.g., trRosetta, Yang et al., PNAS 2020). They’re easy to predict but fundamentally limited — many different 3D folds can produce similar distance distributions.

Random baseline: ~0.54. The easiest structural loss — even the EGNN models (v8–v10) could reduce this 60% below random. Useful as a stable gradient signal early in training before FAPE gradients become informative.

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

MSE between consecutive C\(\alpha\)–C\(\alpha\) distances and the ideal value of 3.8Å (the standard peptide bond length from crystallography). The weight is annealed from 1.0 to 3.0 over 10 epochs.

Why anneal? This was learned the hard way in v9. Starting with high bond weight causes the model to produce “bead-on-a-string” chains — perfectly spaced at 3.8Å but with no tertiary structure. The model satisfies the easy loss and ignores the hard one (FAPE). Annealing lets the model first explore topology, then gradually tightens the physical constraints.

Random baseline: ~0.17. Converged: <0.02 (bonds within 0.1Å of ideal). Currently at 0.009 — essentially solved.

Chirality  \(w = 0.1\)

MSE on signed volumes (scalar triple products) of C\(\alpha\) quartets. Natural proteins are built exclusively from L-amino acids, which produces a consistent backbone handedness. Without explicit chirality supervision, diffusion models can generate mirror-image (D-amino acid) structures that score well on all other metrics — FAPE, RMSD, and distance MSE are all chirality-agnostic.

This is a known issue in the protein generation literature. FrameDiff (Yim et al., ICLR 2024) addresses it through SE(3) equivariance; we use an explicit loss term instead, which is simpler and equally effective for our architecture.

Bond Angle  \(w = 0.5\)

MSE on cosines of C\(\alpha\)–C\(\alpha\)–C\(\alpha\) bond angles. The ideal angle is ~120° (\(\cos\theta \approx -0.5\)), reflecting the planar geometry of the peptide bond. Working in cosine space avoids the discontinuity at 0°/360° that plagues angular losses.

Random baseline: ~0.70. Converged: <0.1. Currently at 0.038 — solved.

Radius of Gyration  \(w = 0.5\)

A separate MLP predicts the protein’s absolute radius of gyration (R\(_g\)) from sequence embeddings, trained with MSE on log-transformed values. At inference, this recovers the correct physical scale when converting from normalized coordinates back to Angstroms.

Converges below 0.05 by epoch 2 and stays solved. The simplest loss — included for practical utility, not structural learning.

Auxiliary Distance Cross-Entropy  \(w = 0.03\)

Ordinal regression on binned pairwise distances (32 bins, 2–40Å) from an independent lightweight pair stack (64-dim, ~500K params). Inspired by the distogram head in AlphaFold2, which predicts distance distributions as a byproduct of the pair representation.

Key design choice: The auxiliary pair stack is completely independent from the main pair stack — no shared parameters, no gradient flow between them. This was a hard-won lesson from v10, where the distance head read from the main pair stack via detach(). As the main stack’s representations evolved during training, the distance head’s features drifted out of distribution, causing aux_dist_ce to explode (3.95 → 38.8 over 3 epochs). A fully independent stack avoids this coupling entirely.

DDIM Evaluation Metrics

Every 5 epochs, we generate structures via 50-step DDIM sampling (Song et al., ICLR 2021) using EMA weights and evaluate against ground truth. These metrics measure generative quality — how good the model’s samples are — as opposed to the training losses which measure denoising accuracy.

Current status: At E5, v14’s DDIM samples give TM=0.111 (near random). Training losses are improving steadily but generative quality typically lags — expect a nonlinear jump in DDIM metrics around E10–15 once the denoiser learns enough of the score function for iterative refinement to work. This “phase transition” is well-documented in score-based diffusion models (Song et al., ICLR 2021).