Organizations: Mila – Quebec AI Institute · Polytechnique Montréal · Center for Brains, Minds, and Machines · MIT · Amgen · Canada CIFAR AI Chair · University of Montréal
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.
Figures & tables
Figure 1: recAMPLIFY architecture. A non-recurrent prelude maps input embeddings to a residue-aligned latent state zk , a weight-tied Transformer core iteratively refines this state, and a non-recurrent coda maps the refined state to the output space. Gated Mixture Recurrence (GMR) decodes from a probability-weighted mixture over the recurrence trajectory, with a per-sequence residual gate at scalar, low-rank, or per-dimension channel resolution controlling how much of each cycle’s proposed update is accepted. Blue circles illustrate these channel-level parameterizations.
Figure 2: recAMPLIFY’s per-assay optimal recurrence depth on ProteinGym DMS.
Figure 3: recAMPLIFY at varying recurrence depths. (a) For 1 M-step recAMPLIFY, the coda embeddings become progressively more class-separable ( 5 -NN accuracy) with larger test-time depth Leval . (b) For 100 K-step recAMPLIFY, as Leval increases past the trained depth ( L=6 ), the fixed- L model (red) degrades while the variable- L variant (blue) remains largely stable.
Figure 4: Effect of recurrence depth on sequence likelihood and downstream separability at 100 K steps. (a) UniProt perplexity for fixed- L pretraining, evaluated at Leval=L . (b) EC level-1 and level-2 5 -NN macro-F1 across Leval for recAMPLIFY trained with L=6 , using 14 K Swiss-Prot enzymes balanced across the 7 top-level classes relabeled into 46 level-2 subclasses ( ≥30 samples each). (c) SCL 10 -way 5 -NN macro-F1 across five FLIP DeepLoc-1.0 splits ( Didi et al., 2026 ) , fitted on 11,181 shared training sequences and evaluated on each held-out test set.
Figure 5: Dynamics of recurrent block layer ℓ on stability prediction upon training with L=6 cycles (vertical dashed lines). (a) The successive difference norm ∥zkℓ−zk−1ℓ∥2 shows gating drives each layer’s step size toward zero beyond the trained depth. (b) The normalized difference to each layer’s deepest recorded state highlights the rate and timing of decay as key convergence diagnostics.
Figure 6: ColSum concentration through the recurrent block at Leval=50 . Cycles colored light (early) to dark (late). Dashed red line is the feedforward reference.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Performance comparison of different representation choices used for training downstream supervised task probes: z uses the output of the recurrent core while coda is the standard pre-LM head output fc(z) . We sweep Leval∈{1,…,9} on the same recAMPLIFY trained with (p,r,c)=(2,8,2) , and L=3 ( ⋆ denotes Leval with matched L ).
Figure 8: Ablation of physical layer counts: we train each variant for 100 K steps and evaluate at its training depth: L=Leval=4 in (a) and L=Leval=3 in (b). Overall, r=10 performs best in (a) on DMS, CASP, fold and SSP, while r=12 is best on KL (NT). The prelude/coda split in (b) changes every probe by less than 1 pp.
Figure 9: Effect of truncated backpropagation (BPTT) steps on supervised probes. We sweep across Leval∈{1,…,9} for four checkpoints that are trained with L=6 but differ in the truncated backprop-through-time window τ : full BPTT ( τ= full, gradients through every cycle) and truncated windows τ∈{3,4,5} . Here, the last τ cycles receive gradients, while the earlier L−τ do not. Full-BPTT holds up across the sweep and remains competitive with AMPLIFY- 350 M (main-paper Table 1 ). Truncated variants lose about 0.2 in fold accuracy and 0.08 in SSP accuracy at every depth, and do not benefit from additional cycles on SSP. Stability is the exception, with τ=3 and τ=5 matching or exceeding full BPTT, while τ=4 becomes unstable across seeds at larger depths.
Figure 10: Varying context type for L=4 . All models are evaluated at the matched Leval=4 depth [ Liao and Poggio, 2026 ] . The simplest context type, ‘zero’, which uses an identity transformation of the previous cycle’s latent zk−1 , performs surprisingly well throughout all tasks.
Figure 11: Input-injection ablation for recAMPLIFY. Blue line shows a normal forward pass while red denotes injection stream removed from the recurrent core. Left: cos(zk,zL) . Right: ∥zk−zk−1∥2 on log- y . With input injection, the state is a bounded iterated map ( ∥Δz∥2 plateaus below 103 ). Without the injection, the state diverges by three orders of magnitude within a few cycles. Input injection thus helps keep the recurrence bounded, matching the finding of Blayney et al. [2026] on causal-LM looped Transformers. We note that this boundedness does not imply convergence. The blue plateau in the right panel stays around ∥Δz∥2≈300 throughout the OOD range, motivating main-paper Figure 5 .
Figure 12: Per-layer dynamics of each recurrent block layer ℓ on the stability prediction task (median over 32 probes). (a) Successive-difference norm: gating drives each layer’s step size toward zero past the L=6 cycles seen during training. (b) Normalized operator residual to each layer’s deepest recorded state zLℓ (taken as an approximate fixed point) over k∈{1,…,L−1} cycles shows the key diagnostic being the rate and timing of decay. Scalar and per-dim gates reach near-zero residual within a few cycles of the trained horizon, whereas halt-only and rank- 16 decay only slowly and remain substantially non-zero at intermediate cycles. (c) Cosine similarity to zLℓ : alignment is near-perfect for every variant, so cosine is an indirect diagnostic and must be read jointly with (a) and (b). Read together, only the scalar and per-dimension gates deliver empirical fixed-point convergence.
Figure 13: Post-halt attention stability. Head-averaged attention Frobenius distance between recurrent cycles for three representative variants at Leval=50 . Each cell (k,k′) is ∥Ak−Ak′∥F averaged over B=4 probe sequences from the stability prediction task, and the dashed line marks the trained halt horizon at k≈7 . For the variant with fixed test-time depth and varying- L pretraining (left), attention keeps drifting through all 50 cycles, despite the recurrence being bounded (Figure 11 ). For halt-only GMR (middle), a discontinuity appears at the halt horizon but off-diagonal cells continue to darken smoothly. For the scalar gate, cycles 8 – 50 collapse into uniform differences. The comparison isolates the role of the residual gate where neither variable- L training nor a halt signal on its own produces a post-halt attention fixed point; only the gate achieves this.
Figure 14: ColSum concentration vs. physical layer position across five recAMPLIFY variants, evaluated at Leval=50 . Solid lines with circles are the trained cycles ( [1,6] ). Dashed lines with diamonds: the six evenly spaced samples from the extrapolation range [7,50] , orange square: prelude, black triangle: coda. The dashed red line is the feedforward reference which is a recAMPLIFY trained with the same (p,r,c)=(2,8,2) physical layer counts without recursion, i.e., L=1 .
Figure 15: Median residual halt depth E[k] binned by composition-diversity × hydrophobicity. Left: Shannon entropy of the AA distribution (rows) against % of FILMVWY (columns). Right: Jensen–Shannon divergence of the AA distribution to the uniform 1/20 prior (rows) against % of FILMVWY (columns). Both panels are computed on the 10 K UniProt held-out sequences of § 5.1.3 and use E[k] residualized against logT and the model’s prediction sharpness −H[pθ] . Diverging color scale is centered at 0 . Positive (red) means the model spends more cycles than a length- and sharpness-matched null.
Figure 16: Median residual halt depth E[k] by amino acid × per-AA composition decile. Rows are grouped into the seven hydrophobic/aromatic residues (bottom) and the remaining residues (top). The horizontal line separates the hydrophobic block (I, L, V, F, M, W, Y) from the non-hydrophobic block above. Columns are within-AA composition deciles (bin 1 : lowest fraction, bin 10 : highest). Hydrophobic AAs show increases over most of the range: heavier hydrophobic content in general drives more GMR cycles. Non-hydrophobic AAs show an inverted U-shape with extreme composition dominance driving fewer cycles.
Figure 17: LM-head top- 1 probability per real token on stability prediction, evaluated with (ON) and without (OFF) input-injection.
Proteins perform diverse cellular functions, and even single amino-acid substitutions can alter stability, activity, or molecular interactions. Protein language models (PLMs) provide a scalable approach for modeling such sequence--function relationships from unlabeled sequences, but increasing the size of dense Transformer backbones often brings substantial computational cost without consistently improving mutation-sensitive prediction. We introduce ProtLingo, an efficient PLM framework that augments a pretrained single-sequence backbone with conditional local memory and sparse expert routing. ProtLingo maps contextual residue representations into route-specific discrete codes, composes centered local windows into latent N-gram addresses, and retrieves reusable residual signals associated with recurring local sequence contexts. In parallel, selected feed-forward blocks are upcycled into sparse Mixture-of-Experts layers with shared and routed experts, enabling residue-dependent computation while activating only a subset of parameters. Experiments on protein fitness prediction, FLIP benchmarks, and supervised contact prediction show that ProtLingo achieves competitive performance with a 150M-scale backbone, including strong parameter efficiency on mutation-effect prediction and preserved long-range structural representations.
Unlike the predictable scaling laws in natural language processing and computer vision, protein language models (PLMs) scale poorly: for many tasks, models within the same family plateau or even decrease in performance, with mid-sized models often outperforming the largest in the family. We introduce Reverse Distillation, a principled framework that decomposes large PLM representations into orthogonal subspaces guided by smaller models of the same family. The resulting embeddings have a nested, Matryoshka-style structure: the first k dimensions of a larger model's embedding are exactly the representation from the smaller model. This ensures that larger reverse-distilled models consistently outperform smaller ones. A motivating intuition is that smaller models, constrained by capacity, preferentially encode broadly-shared protein features. Reverse distillation isolates these shared features and orthogonally extracts additional contributions from larger models, preventing interference between the two. On ProteinGym benchmarks, reverse-distilled ESM-2 variants outperform their respective baselines at the same embedding dimensionality, with the reverse-distilled 15 billion parameter model achieving the strongest performance. Our framework is generalizable to any model family where scaling challenges persist. Code and trained models are available at https://github.com/rohitsinghlab/plm_reverse_distillation.
Darius Catrina, Christian Bepler, Samuel Sledzieski +1
Department of Computer Science, Duke University, Durham, NC, USA · Depts of Biostats. and Bioinformatics; and Cell Biology, Duke University, Durham, NC, USA
Protein language models (PLMs) have become widely adopted as general-purpose models, demonstrating strong performance in protein engineering and de novo design. Like large language models (LLMs), they are typically trained as deep transformers with next-token or masked-token prediction objectives on massive sequence corpora and are scaled by increasing model depth. Recent work on autoregressive LLMs has identified the Curse of Depth: many later layers contribute little to the final output predictions. These findings naturally raise the question of whether a similar depth inefficiency also appears in PLMs, where many widely used models are not autoregressive, and some are multimodal, accepting both protein sequence and structure as input. In this work, we present a depth analysis of seven popular PLM families across model scales, spanning autoregressive, masked, and diffusion objectives, and quantify how layer contributions evolve with depth using a unified set of probing-, perturbation-, and downstream-evaluation measurements. Across models, we observe consistent depth-dependent patterns that extend prior findings on LLMs: a large fraction of task-relevant computation is concentrated in a subset of layers, while the remaining layers mainly provide incremental refinement of the final prediction. These trends persist beyond sequence-only settings and also appear in multimodal PLMs. Taken together, our results suggest that depth inefficiency is a common feature of modern PLMs, motivating future work on more depth-efficient architectures and training methods.
Aleena Siji, Amir Mohammad Karimi Mamaghan, Ferdinand Kapl +9
Technical University of Munich · Helmholtz AI, Munich · KTH Royal Institute of Technology +2