# ProteinJEPA: Latent prediction improves protein language model pretraining

arXiv: https://arxiv.org/abs/2605.07554 · HTML: https://mlsanity.com/paper/2605.07554/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 18 of its 58 references that are in the map.
- First posted on arXiv 8 May 2026.

## Builds on (most cited first, 18 of 18)

- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Don't Stop Pretraining: Adapt Language Models to Domains and Tasks](https://mlsanity.com/paper/2004.10964/) (2020)
- [VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning](https://mlsanity.com/paper/2105.04906/) (2021)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [data2vec: A General Framework for Self-supervised Learning in Speech, Vision and Language](https://mlsanity.com/paper/2202.03555/) (2022)
- [Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference](https://mlsanity.com/paper/2412.13663/) (2024)
- [Revisiting Feature Prediction for Learning Visual Representations from Video](https://mlsanity.com/paper/2404.08471/) (2024)
- [LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels](https://mlsanity.com/paper/2603.19312/) (2026)
- [LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics](https://mlsanity.com/paper/2511.08544/) (2025)
- [Evaluating Protein Transfer Learning with TAPE](https://mlsanity.com/paper/1906.08230/) (2019)
- [ProGen2: Exploring the Boundaries of Protein Language Models](https://mlsanity.com/paper/2206.13517/) (2022)
- [The geometry of hidden representations of large transformer models](https://mlsanity.com/paper/2302.00294/) (2023)
- [BERTology Meets Biology: Interpreting Attention in Protein Language Models](https://mlsanity.com/paper/2006.15222/) (2020)
- [Cramming: Training a Language Model on a Single GPU in One Day](https://mlsanity.com/paper/2212.14034/) (2022)
- [PEER: A Comprehensive and Multi-Task Benchmark for Protein Sequence Understanding](https://mlsanity.com/paper/2206.02096/) (2022)
- [LLM-JEPA: Large Language Models Meet Joint Embedding Predictive Architectures](https://mlsanity.com/paper/2509.14252/) (2025)
- [JEPA-DNA: Grounding Genomic Foundation Models through Joint-Embedding Predictive Architectures](https://mlsanity.com/paper/2602.17162/) (2026)
- [From Words to Amino Acids: Does the Curse of Depth Persist?](https://mlsanity.com/paper/2602.21750/) (2026)

## Abstract (arXiv)

Protein language models are trained primarily with masked language modeling (MLM), which predicts masked amino-acid identities. Joint-embedding predictive architectures (JEPA) instead predict latent representations, but have not been applied to proteins. ProteinJEPA supplements MLM with a cosine loss for predicting the half-depth hidden states of a teacher given the unmasked sequence. On 19 tasks, with ESM2 at 35M and 150M parameters and three pretraining seeds, MLM+JEPA outperforms compute-matched and step-matched MLM-only continued training in 78 and 76 of 114 comparisons (14 losses, 22 ties). The median compute-matched gain is +0.0106 on structure- and homology-sensitive tasks versus +0.0041 elsewhere, led by SCOPe-40 retrieval and remote homology with improvements of 6.1 percentage points in Recall@1 and 2.7 points in accuracy, respectively. Gains on these tasks increase with model size from 8M to 150M. Against the off-the-shelf checkpoint, MLM+JEPA wins 81 of 114 comparisons (median +0.0068) without improving MLM loss. In random initialization the gain is smaller and replicates inconsistently across seeds (p=0.059). The same recipe improves the causal ProGen3 model, beating a compute-matched next-token-prediction control on 12 of 16 tasks. Ablations show that cosine loss beats mean squared error, while adding shallower targets removes most of the task gain. JEPA-only training collapses downstream performance: latent prediction complements MLM rather than replacing it. Code: https://anonymous.4open.science/r/protJepa-FF24