cs.LGMay 8, 2026

ProteinJEPA: Latent prediction improves protein language model pretraining

Authors: Dan OferDafna ShahafMichal Linial

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+0.0106 on structure- and homology-sensitive tasks versus +0.0041+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+0.0068) without improving MLM loss. In random initialization the gain is smaller and replicates inconsistently across seeds (p=0.059p{=}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

Explore similar work

CardsList
  1. Task- and dataset-specific information in protein language models

    Aug 12, 2026Roman Joeres, Ilya Senatorov, Anastasia Kolchina +2