recAMPLIFY: Large-Scale Pretraining of Recursive Protein Language Models

AI summary

Overview

Research area: Protein language models (pLMs), efficient deep learning, and adaptive/test-time compute scaling. Technical level: Advanced — assumes familiarity with Transformer encoders, masked language modeling, and the standard pLM evaluation suite (ProteinGym, CASP, FLIP-style probes). Scope: The paper adapts depth-recurrent ("looped") pretraining to a bidirectional masked protein language model and introduces a gating mechanism, Gated Mixture Recurrence (GMR), that lets each protein decide how many recurrent cycles to use at inference.

What This Paper Is About

Protein language models have improved mainly by getting bigger, which makes them expensive and prone to memorizing training sequences. The authors ask whether a much smaller model can recover that performance by reusing the same weights repeatedly ("thinking longer" at test time) rather than adding parameters. Because masked protein encoders have no generation loop and amino acids cannot be "verbalized" as reasoning text, they build recursion directly into the latent representations and then add a learned mechanism that allocates different amounts of recursion to different proteins.

Key Contributions

1. A recursive pretraining framework for masked pLMs. The authors adapt a minimal depth-recurrent design (a one-time prelude, a weight-tied Transformer core applied for L cycles, and a one-time coda) from causal decoders to the bidirectional AMPLIFY encoder, keeping hidden states one-to-one aligned with residues. Effective depth becomes p + (r · L) + c while the unique parameter count stays fixed.
2. recAMPLIFY, a pLM with 155M unique parameters pretrained on the same 2T tokens as AMPLIFY-350M, which is competitive with the larger model despite 56% fewer unique parameters.
3. Gated Mixture Recurrence (GMR), a framework unifying halting and gating under one learned survival signal: predictions are decoded from a probability-weighted mixture over recurrent states, and each cycle's residual update is gated at scalar, low-rank (ch = 16, D = 960), or per-dimension resolution. Ungated updates with an independent halt head recover PonderNet-style halting as a special case.
4. Mechanistic characterization of the recurrence, showing that residual gating (not halting alone) drives the state toward an empirical fixed point beyond the trained depth, and that controlled halt depth tracks protein compositional properties.

Main Findings

• Parameter efficiency at full scale: With 155M unique parameters versus AMPLIFY-350M's 354M, and identical 2T-token pretraining, fixed-L recAMPLIFY reaches fold accuracy 0.696 versus 0.568 (the reported 22.5% improvement, from 56.8% to 69.6%) and localization accuracy 0.741 versus 0.716 (a 3.5% improvement, from 71.6% to 74.1%). It also improves CASP contact recovery (0.179 versus 0.146).
• Small trade-offs elsewhere: The same model scores lower on KL calibration (59.32 versus 69.31 for target KL; 37.70 versus 38.79 for non-target KL) and zero-shot ProteinGym DMS Spearman (0.348 versus 0.360). Secondary-structure prediction ties at 0.742.
• Gains over the similarly sized baseline: recAMPLIFY outperforms AMPLIFY-120M (118M parameters) across supervised probes by up to 23%, and beats AMPLIFY-120M on fold (0.696 versus 0.566), localization (0.741 versus 0.651), stability (0.703 versus 0.591), and CASP (0.179 versus 0.126).
• More depth helps up to a point: Additional cycles improve held-out perplexity and downstream separability. On subcellular localization, deeper recursion yields macro-F1 gains of 20–31 pp across FLIP's train/test cluster splits. But the model trained with L = 6 degrades beyond that horizon — stability falls from 0.644 to 0.465 — while variable-depth pretraining (cycles sampled from {1, …, 6}) remains largely stable.
• No single optimal depth: At 100K pretraining steps, KL calibration, CASP and SSP peak at L_eval = L = 6, whereas fold, localization and stability peak slightly earlier at L_eval ∈ {4, 5}. Per-assay optima on ProteinGym DMS also differ.
• Adaptive recursion cuts compute: Every GMR variant lowers expected test-time depth from 52 to 33–37 layer applications, while improving over fixed-L recAMPLIFY on CASP (0.135–0.161 versus 0.117), fold, localization and SSP, and on stability for three of four variants. The cost is a small drop in DMS (0.346–0.347 versus 0.359) and at most 0.32 NT points of KL calibration. Halt-only GMR reaches the shallowest expected depth, E[k] = 3.57 (2 + 3.57 × 8 + 2 ≈ 33 layers).
• Gating, not halting, produces fixed points: Unrolled to L_eval = 50, scalar and per-dimension gates drive the per-layer step norm toward zero and show rapid decay of the normalized residual to the deepest state within a few cycles of the trained horizon; halt-only GMR decays only linearly, and the rank-16 gate drifts radially. Cosine alignment is near-perfect for all variants and is only an indirect diagnostic.
• Gating preserves layerwise mixing structure: Measured by ColSum concentration, only per-dimension gating approaches the feedforward reference in both shape and amplitude at extrapolated depth; within the trained horizon all variants retain a non-flat profile across the eight tied core layers.
• Halt depth mostly reflects confidence, then composition: The raw correlation between per-protein halt depth and KL is r = −0.39, but after controlling for sequence length and the model's prediction sharpness it becomes +0.11, a reduction of roughly 72%. The controlled halt depth, computed on 10K held-out UniProt sequences, most robustly tracks compositional diversity, with over-pondered proteins showing higher sequence entropy and hydrophobic content.
• Context against ESM-2 (caveated): Larger ESM-2 checkpoints (4–18× more parameters, different UniRef release and clustering) outperform recAMPLIFY on three of four supervised probes, but recAMPLIFY exceeds ESM-2-650M and ESM-2-3B on stability prediction and achieves better target KL calibration than ESM-2-150M and ESM-2-650M.
• Cost scales with depth, not parameters: A 133M recAMPLIFY at L = 3 achieves measured effective throughput comparable to AMPLIFY-350M, which has approximately 2.7× as many unique parameters.

Methodology in Plain English

The authors start from a published, compact architecture (AMPLIFY) and split its encoder into three parts: a short prelude that turns amino acids into an initial latent representation, a core of tied Transformer layers that is applied over and over, and a short coda that turns the final state into predictions. The input representation is re-injected at every cycle so the repeated application stays stable, and the latent state always starts from the same frozen random vector (σ = 0.02), which keeps one state vector per residue and preserves the alignment needed for masked prediction. Crucially, the number of cycles L is a dial that can be turned up at inference without adding weights.

At full scale they use (p, r, c) = (2, 10, 2) with 15 attention heads and hidden dimension 960, pretraining on UniRef100, SCOP and OAS under AMPLIFY's recipe for 1M steps. Because full-scale ablations are expensive, they first verify that a smaller (2, 8, 2) configuration trained for 100K steps shows a representative performance profile, then run all ablations there.

Their adaptive variant, GMR, adds a small head that reads the current state (with gradients stopped so it does not distort the language-modeling features) and outputs a halt probability per cycle. These probabilities define a distribution over cycles; predictions are decoded from the probability-weighted average of the visited states, and the training loss adds an entropy term to discourage halting too early. The same head can instead emit a survival gate that controls how much of each cycle's proposed update is accepted, at three resolutions: one value for the whole sequence, 16 shared trajectories projected onto the 960 features, or one value per feature. They evaluate with zero-shot probes (KL calibration, ProteinGym DMS, CASP16 contact recovery) and four supervised probes from an 18-task benchmark, reporting means over 5 seeds after a 50-trial Bayesian TPE search. Mechanistic analyses then unroll models to 50 cycles and track step norms, distance to the deepest state, attention geometry and ColSum concentration.

Why This Matters

The work reframes protein model scaling: rather than only adding parameters and data, a model can be given more computation per input at inference. That matters given published evidence that larger pLMs can memorize training data, and it offers a route to strong performance under a fixed parameter budget.

Real-world applications:

• Protein function annotation: predicting fold, subcellular localization and enzyme class from sequence alone, where recAMPLIFY's largest gains appear.
• Variant-effect and mutation screening: zero-shot scoring of mutations for clinical and research prioritization, though the paper reports a small DMS cost relative to AMPLIFY-350M.
• Protein and enzyme engineering: guiding stability-improving mutations, where recAMPLIFY outperforms AMPLIFY-350M and the larger ESM-2 checkpoints.
• Structure-adjacent modeling: contact recovery from sequence, evaluated on CASP16, as a cheap first pass before expensive structure prediction.

Industry relevance is direct: a 155M-parameter checkpoint is far cheaper to store, serve and fine-tune than multi-billion-parameter models, and input-dependent depth offers a knob for trading latency against accuracy per request. The author list includes an industry affiliation (Amgen), underscoring the therapeutic-protein interest. The authors are explicit that the reported depth reductions characterize the learned allocation policy, not measured wall-clock savings, since evaluation decodes the full mixture over the cycle budget.

Future Directions

• Turning adaptive depth into real speedups: since evaluation decodes the whole mixture, actual latency and energy savings from halting are not reported and need hardware-aware measurement.
• Better halting objectives for masked pLMs: the paper notes no verifier exists for masked language modeling, motivating the self-supervised entropy-regularized head; whether alternative self-supervised signals produce better depth allocation is open.
• Extending recurrence beyond the studied regime: the full-scale model was trained at fixed L, and the ablations cap at 6 cycles and 100K steps; scaling cycle counts, and combining recurrence with structure-aware recycling as in AlphaFold/ESMFold, remains unexplored.
• Explaining per-assay optima: why fold, localization and stability peak at L_eval ∈ {4, 5} while KL, CASP and SSP peak at L = 6 is observed but not mechanistically resolved, and the gate resolution (scalar versus low-rank versus per-dimension) shows no consistent winner across tasks.

Target Audience

Researchers and practitioners working on protein language models and computational biology who want a concrete recipe for test-time compute scaling; efficient-ML researchers interested in depth recurrence, adaptive computation and halting; and industry teams in biotech or pharma who need strong protein representations at a fraction of the parameter cost. Readers without background in masked language modeling or the ProteinGym/CASP/FLIP benchmark ecosystem will find the evaluation section dense.

Authors’ abstract

Protein language models (pLMs) are powerful tools for protein annotation and mutation-effect prediction, but progress has largely depended on scaling model size and data, making them expensive and prone to memorization. Test-time scaling offers a complementary path, but methods developed for autoregressive LLMs do not transfer directly: masked encoders lack a generation loop, and amino acids do not verbalize in a reasoning language. Instead, we adapt recursive pretraining to masked pLMs, where a weight-tied Transformer core refines residue-aligned hidden states over an adjustable number of cycles. We instantiate this recipe as recAMPLIFY, a pLM with 155M unique parameters trained on the same 2T tokens as AMPLIFY-350M. Despite 56% fewer parameters, it improves CASP contact recovery and three out of four supervised tasks, notably raising fold accuracy by 22.5%. Because optimal depth for recAMPLIFY varies across tasks and assays, we then introduce Gated Mixture Recurrence (GMR), a framework that decodes each protein from a probability-weighted mixture of hidden states visited during recursion. GMR employs a learned residual gate at scalar, low-rank, or per-dimension resolution to control how much of each cycle's proposed update is accepted into the state. This unifies halting and gating into a single adaptive computation objective while recovering PonderNet-style halting as a special case without gates. Layered on top of recAMPLIFY, GMR further reduces expected layer depth by up to 37%, although the accuracy trade-off varies by task. Mechanistic analyses show that residual gating controls recurrence stability beyond the trained depth and that halt depth tracks biological properties.