ProteinJEPA: Latent prediction improves protein language model pretraining

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

Understanding language model scaling for protein fitness prediction
Chao Hou, Ди Лю, Aziz Zafar +1
Nature Computational Science13 Jul 2026Peer-reviewedLicence not statedFull text22 citedMolecular biology & gene regulationOpenAlex
Protein Language Model-Boosted Prediction of Isoelectric Point
Fangfang Luo, Xiangxiang Lu, Zhitao Cai +2
ChemRxiv5 Oct 2026PreprintCC BY-NC-NDFull textMolecular biology & gene regulationChemRxiv
Protein–Protein Interface Residue Prediction by Integrating Protein Language Representations and Structural Graph Attention
Zhen Hou, Wei He, Weinan Cao +4
ACS Omega7 Oct 2026Peer-reviewedCC BY-NC-ND 4.0Full textMolecular biology & gene regulationCrossrefEurope PMC
AIPFusion-Net: an interaction-fusion framework for anti-inflammatory peptide prediction using protein language model embeddings, local sequence patterns, and handcrafted descriptors.
Wentao Wang, Weidong Yang, Hongmei Li +2
Frontiers in molecular biosciences5 Oct 2026Peer-reviewedCC BYFull textMolecular biology & gene regulationEurope PMC
SaltyMeta: a curated benchmark and protein language model-informed web tool for salty peptide prediction
Wanchao Chen, Wen Li, Yanan He +1
arXiv15 Sep 2026PreprintCC BY-SA 4.0Full textBiophysicsarXiv

Cite

BibTeX RIS

BibTeX

@misc{ofer2026proteinjepa,
  title = {{ProteinJEPA: Latent prediction improves protein language model pretraining}},
  author = {Dan Ofer and Dafna Shahaf and Michal Linial},
  howpublished = {arXiv},
  year = {2026},
  month = sep,
  doi = {10.48550/arxiv.2605.07554},
  eprint = {2605.07554},
  archivePrefix = {arXiv},
  url = {https://arxiv.org/abs/2605.07554},
  note = {CC BY-NC-SA 4.0},
}

RIS

TY  - UNPB
TI  - ProteinJEPA: Latent prediction improves protein language model pretraining
AU  - Dan Ofer
AU  - Dafna Shahaf
AU  - Michal Linial
PY  - 2026
DA  - 2026/09/23
PB  - arXiv
DO  - 10.48550/arxiv.2605.07554
UR  - https://arxiv.org/abs/2605.07554
AB  - 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
ER  -

Abstract

Related

Cite