Protein language model improves prediction with fewer parameters
recAMPLIFY: Large-Scale Pretraining of Recursive Protein Language Models
Machine LearningArtificial IntelligenceComputation and Language
Summary
Predicting how proteins fold and behave is important but challenging, and current models need to be very large to work well. The authors present recAMPLIFY, a new protein language model that uses a clever way to repeatedly refine its understanding of amino acid sequences without needing as many parameters. Their approach improves key prediction tasks like protein folding accuracy and contact recovery while being smaller than comparable models. They also introduce a gating mechanism that adapts how much each refinement step influences the model, making it more efficient and stable, and capable of adjusting for different tasks.
What this means in practice
- •For biotechnology developers: Use the recAMPLIFY model to improve protein structure and function predictions with smaller models, reducing computational costs in biotech pipelines.
- •For drug discovery teams: Employ the adaptive gating approach to balance accuracy and efficiency in predicting mutation effects during drug candidate evaluation.
Authors
Saurav Jha, Qianli Liao, Robert Vernon, Sarath Chandar, Quentin Fournier
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.