ProteinJEPA: Latent prediction improves protein language model pretraining
cs.LG, cs.AI, q-bio.BM, stat.ML
Submitted: 2026-05-08
Updated: 2026-09-23
Code: https://github.com/nadavbra/protein_bert
License: http://creativecommons.org/licenses/by-nc-sa/4.0/
The gist: Protein language models are trained primarily with masked language modeling (MLM), which predicts masked amino-acid identities.
Terminology
Abstract
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
Sources
- Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture
- LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristics
- Revisiting Feature Prediction for Learning Visual Representations from Video
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
- LLM-JEPA: Large Language Models Meet Joint Embedding Predictive Architectures
- JEPA-DNA: Grounding Genomic Foundation Models through Joint-Embedding Predictive Architectures
- LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels
- ProGen2: Exploring the Boundaries of Protein Language Models
- Evaluating Protein Transfer Learning with TAPE
- Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: The Predictive Cost of ELBO Binding
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks