Structure-Aligned Protein Language Model (SaESM2 / SaAMPLIFY)

Protein language models learn from sequence alone, and that limits them in a specific way: they score a mutation by how plausible the substituted residue is given the evolutionary record, not by whether the resulting protein would still fold. Structure-Aligned Protein Language Model is a post-training recipe that closes part of that gap, transferring structural knowledge out of pretrained protein graph neural networks and into an existing sequence model without retraining it from scratch. Applied to ESM-2 and to AMPLIFY, it yields the released checkpoints SaESM2 and SaAMPLIFY.

The method combines two objectives. A latent-level contrastive task aligns each residue's language-model representation with the corresponding residue embedding from a pretrained protein graph neural network — GearNet pretrained on the AlphaFold database — pulling matched residues together and pushing mismatched ones apart across many proteins, which injects inter-protein structural signal. A physical-level task then trains the language model to predict Foldseek 3Di structure tokens for each position, injecting intra-protein geometry. Because deposited structures vary widely in quality, a residue loss selection module uses a small reference model trained only on high-resolution entries to pick which per-residue losses are reliable enough to learn from.

The work comes from Mila – Quebec AI Institute and the Chandar Research Lab, with Yoshua Bengio and collaborators at Amgen, and was posted as a preprint in May 2025. It sits alongside earlier joint sequence-structure efforts such as ESM-GearNet, but differs in kind: rather than building a fused two-encoder system, it produces drop-in replacements that keep the original language model's architecture and interface.

#Key Features

• Dual alignment objective: Contrastive latent alignment supplies structural context across proteins, while structure-token prediction supplies geometry within a protein; the two losses are weighted to sum to one.
• Quality-aware loss selection: A reference model trained on structures below 2.0 Å resolution with R-free under 0.20 identifies residues whose losses are informative rather than artifacts of poor experimental data, selecting the top 80% by excess loss.
• Drop-in checkpoints: SaESM2 and SaAMPLIFY keep the base models' tokenizers, masked-language-model heads, and Hugging Face interfaces, so existing pipelines need no code changes.
• Consistent scaling: Gains hold from 8M through 650M parameters and across both model families, indicating the effect is not a quirk of one architecture or size.
• Better zero-shot fitness: On the ProteinGym deep mutational scanning substitution benchmark of 217 assays, structure-aligned models outperform their unaligned counterparts at every size, placing SaESM2 650M and 150M on the Pareto front of accuracy against parameter count.

#Technical Details

Structure alignment is trained on 129,732 PDB proteins (116,713 for training, 13,019 held out), using Foldseek's 20-token 3Di alphabet as the structure vocabulary. On supervised contact prediction, SaESM2 raises P@L/5 on a withheld CASP16 set from 0.181 to 0.288 — the 59% improvement quoted in the abstract — and from 0.390 to 0.461 on trRosetta; SaAMPLIFY moves from 0.155 to 0.169 on CASP16 and 0.253 to 0.320 on trRosetta. Secondary structure accuracy rises from 0.845 to 0.865 for the ESM-2 family, and supervised stability prediction improves from 0.744 to 0.820 Spearman correlation. The authors also report a trade-off: alignment raises pseudo-perplexity on a natural-sequence validation set, meaning some pure sequence-modeling capability is sacrificed, though both flagship checkpoints stay competitive. Released sizes span SaESM2 at 8M, 35M, 150M, and 650M parameters, and SaAMPLIFY at 120M and 350M.

#Applications

The checkpoints are useful anywhere a protein language model is already in service but structural realism matters: zero-shot ranking of variants in deep mutational scanning and protein engineering campaigns, residue-residue contact and secondary structure prediction, thermostability and metal-ion-binding property prediction, and as frozen feature extractors for downstream classifiers. Because the alignment is a lightweight post-training step rather than a new pretraining run, the same recipe can in principle be applied to other sequence models as they are released.

#Impact

The result is an argument that structural knowledge can be grafted onto sequence models cheaply, rather than requiring architectures that consume structures at inference time. Code lives in the AMPLIFY repository under an MIT license and weights are on Hugging Face; the checkpoints ship without substantive model cards, so the preprint remains the reference documentation. Evaluation is entirely computational, and the work has not yet been peer reviewed.