LLM-JEPA
Free while signed in. Answers cite the passages they came from.
A JEPA-style training objective is adapted to LLMs by treating paired views of the same underlying content (for example, text and code) as prediction targets in embedding space, added on top of the usual next-token loss. The result consistently improves fine-tuning and shows promising pretraining gains, while being more resistant to overfitting.
Idea in one line: Keep the standard next-token objective and add a JEPA term that predicts the embedding of one view from another using tied LLM weights with special predictor tokens k, optimized with a cosine metric and weight λ. This preserves generation while improving abstraction.
Why it helps: Minimizing next-token loss alone does not reduce the JEPA prediction error; adding the JEPA term closes this gap and explains the accuracy lift.
Main results: Across Llama, Gemma, OpenELM and OLMo, LLM-JEPA improves exact-match accuracy on NL-RX (SYNTH and TURK), GSM8K, and Spider.
Representation effects: t-SNE plots show clearer structure when using LLM-JEPA, and a near-linear mapping from Enc(Text) to Enc(Code) is supported by low regression error and compressed singular values.
Pretraining signal and costs: Adding JEPA during pretraining improves downstream sentiment classification after standard fine-tuning, while keeping generative quality. Current limitation is extra compute from separate forward passes for each view, plus nontrivial hyperparameter sweeps over k and λ.
Get next week’s papers.
The same picks and the same summaries, in your inbox. Free, and 176 issues deep.
Subscribe on Substack