ProteinJEPA-ESMC600M-scratch-MLM-JEPA

An ESM-C 600M-sized protein language model trained from random weights with masked language modeling (MLM) plus a JEPA latent-prediction loss.

What it is

• Architecture: ESM-C 600M layout (36 layers, hidden size 1152, about 575M parameters), from the Synthyra/ESMplusplus_large port.
• Training: from scratch (random init, seed 42), not from the released ESM-C weights.
• Data: UniRef50, 68.3M sequences, one epoch. Each sequence is cut to a random 1,024-token window.
• Objective: MLM with 20% masking (80/10/10), plus a JEPA loss. The last layer of the student predicts the hidden states at block 18 of 36 from an EMA teacher, with a cosine loss (weight 1.0) and a VISReg regularizer (weight 0.008). The EMA decay rises from 0.996 to 0.9995.
• Optimizer: AdamW, peak LR 1e-3, 2,000 warmup steps, WSD schedule with a 1-sqrt cooldown over the last 20%.
• Length: 71,000 updates, about 260k tokens per update.

Results

On a 19-task frozen-embedding probe suite, this model beats the step-matched MLM-only model on 16 tasks, loses on 2, and ties on 1. On SCOPe-40 structural retrieval, Recall@1 is 0.439 versus 0.339 for MLM-only. On remote homology, it scores 0.658 versus 0.581.

Against the compute-matched MLM-only model (83,070 steps, 17% more compute), it wins 18 tasks and loses 1; SCOPe-40 0.439 versus 0.351, remote homology 0.658 versus 0.600.

Against the released ESM-C 600M, it scores 8 wins, 10 losses and 1 tie.

How to load

The architecture code is in this repo, so trust_remote_code=True is required. The tokenizer is also in this repo.

from transformers import AutoModel, AutoTokenizer

repo = "GrimSqueaker/ProteinJEPA-ESMC600M-scratch-MLM-JEPA"
tokenizer = AutoTokenizer.from_pretrained(repo)
model = AutoModel.from_pretrained(repo, trust_remote_code=True)

License

The license for this repo is other, not a standard SPDX license. The architecture code (modeling_esm_plusplus.py) comes from Synthyra/ESMplusplus_large, and that repo's license applies. Its own files disagree: its model card says MIT, its top-level LICENSE is MIT, and its NOTICE.txt includes the EvolutionaryScale Cambrian Non-Commercial License for the ESM-C 600M derivative. Read those files before any use of this model, especially commercial use.

Citation

If you use this model, please cite the ProteinJEPA paper (https://arxiv.org/abs/2605.07554):

@misc{ofer2026proteinjepa,
  title         = {ProteinJEPA: Latent Prediction Improves Protein Language Model Pretraining},
  author        = {Ofer, Dan and Linial, Michal and Shahaf, Dafna},
  year          = {2026},
  eprint        = {2605.07554},
  archivePrefix = {arXiv}
}