Source-linked AI summary
LLM-JEPA: Large Language Models Meet Joint Embedding Predictive Architectures
Hai Huang, Yann LeCun, Randall Balestriero
TL;DR
LLMs are largely trained and judged through input-space reconstruction and generation, while adapting JEPA-style embedding objectives to language remains challenging. LLM-JEPA combines reconstruction with a JEPA objective over paired views, and the paper reports improvements over standard fine-tuning across models and datasets. Its current limitations include costly hyperparameter tuning.
Problem
LLMs are commonly judged through generated text, making it challenging to leverage JEPA objectives for their perception and reasoning tasks.
Source: §1 · p. 2Method
LLM-JEPA combines the standard reconstruction loss with a JEPA objective that uses text and code as views of the same underlying knowledge.
Source: §1 · p. 2Results
Takeaways & Limitations
The study presents a first step toward more JEPA-centric LLM pretraining and fine-tuning, with preliminary pretraining experiments described as promising.
Source: §1 · p. 2, §6 · p. 9LLM-JEPA introduces two extra hyperparameters, and its optimal configuration may lie anywhere in a grid, making tuning costly.
Source: §6 · p. 10
Abstract
from arXiv · showhide
Large Language Model (LLM) pretraining, finetuning, and evaluation rely on input-space reconstruction and generative capabilities. Yet, it has been observed in vision that embedding-space training objectives, e.g., with Joint Embedding Predictive Architectures (JEPAs), are far superior to their input-space counterpart. That mismatch in how training is achieved between language and vision opens up a natural question: {\em can language training methods learn a few tricks from the vision ones?} The lack of JEPA-style LLM is a testimony of the challenge in designing such objectives for language. In this work, we propose a first step in that direction where we develop LLM-JEPA, a JEPA based solution for LLMs applicable both to finetuning and pretraining. Thus far, LLM-JEPA is able to outperform the standard LLM training objectives by a significant margin across models, all while being robust to overfiting. Those findings are observed across numerous datasets (NL-RX, GSM8K, Spider, RottenTomatoes) and various models from the Llama3, OpenELM, Gemma2 and Olmo families. Code: https://github.com/rbalestr-lab/llm-jepa.
1 INTRODUCTION
Language modeling remains dominated by reconstruction and generation, while JEPA-style objectives remain challenging to leverage in LLMs evaluated through generated text. LLM-JEPA uses paired text/code views to add an embedding-space objective while retaining generative capabilities.
1 INTRODUCTION
- Reconstruction-based methods predominate in NLP, although LLM tasks also involve perception and reasoning where JEPAs are known to be preferable.Source: §1 · pp. 1–2
- LLM-JEPA adds a JEPA objective to text→code generation by treating text and code as views of the same underlying knowledge.The paper emphasizes that non-trivial views are crucial to the approach.Source: §1 · p. 2
- The approach is restricted to datasets with non-trivial views; a vision-like data-augmentation mechanism could extend it to any dataset.Source: §1 · p. 2
- The proposed objective operates in embedding space with different views while preserving generation, and reports improvements over standard LLM fine-tuning.Source: §1 · p. 2
- Evaluation spans the Llama, Gemma, OpenELM, and Olmo families, NL-RX, GSM8K, Spider, RottenTomatoes, and multiple model sizes.Source: §1 · p. 2
2 BACKGROUND
Contemporary LLMs use Transformer architectures and are commonly trained with autoregressive token-space reconstruction. Alternative embedding-space objectives exist, but the approaches discussed here do not retain generative capabilities.
2 BACKGROUND
- LLMs are commonly trained with autoregressive token-space reconstruction, usually cross-entropy between predicted and ground-truth tokens.The classifier predicts the next token from preceding tokens, with computation performed through causal autoregression.Source: §2 · pp. 2–3
- Contemporary LLMs stack nonlinear operations and skip connections in Transformers, while high-quality pretraining data remains a main performance driver.Source: §2 · p. 2
- SimCSE uses dropout-induced sentence views for contrastive learning (Gao et al. (2021)), while weakly supervised text pairs support joint embeddings (Wang et al. (2022)); these approaches lack generation.Source: §2 · p. 3
- A BERT semantic loss improved semantic-task performance over masked language modeling alone but prevented generative evaluation and use (Reimers and Gurevych (2019)).Source: §2 · p. 3
3 JEPA-LLM: IMPROVING LLMS’ REASONING AND GENERATIVE CAPABILITIES
This section introduces LLM-JEPA, which combines the LLM objective with JEPA prediction across different views, and describes its implementation. It tests whether next-token training alone minimizes the JEPA objective.
3 JEPA-LLM: IMPROVING LLMS’ REASONING AND GENERATIVE CAPABILITIES
- The section presents LLM-JEPA as an objective for improving LLM reasoning and generative capabilities and evaluates its finetuning and pretraining benefits.Text and Code serve as examples of different views of the same underlying knowledge.Source: §3 · p. 3, §3.1 · p. 3
3.1 THE LLM-JEPA OBJECTIVE
- LLM-JEPA adds a JEPA embedding-prediction term to the LLM objective, combining generative and abstraction objectives with λ balancing their contributions.The predictor compares the predicted Text embedding with the Code embedding using a chosen metric, such as ℓ2 distance.Source: §3.1 · p. 3
3.2 IMPLEMENTATION WITH CUSTOM ATTENTION MASK
- A custom self-attention mask obtains separate embeddings for the two views, enabling the JEPA loss in two forward passes instead of three, though training remains slower.The mask is causal per block so the representation of one view does not rely on the other.Source: §3.2 · p. 4
3.3 A GOOD NEXT TOKEN PREDICTOR IS NOT A GOOD JEPA
- 71.10% accuracy for LLLM-JEPA versus 51.95% for LLLM accompanies similar next-token loss, while only LLLM-JEPA minimizes JEPA prediction loss.The authors report that next-token loss cannot explain the accuracy gap and suggest the JEPA prediction loss as its main reason.Source: Figure 3 · p. 5
4 EMPIRICAL VALIDATION: LLM-JEPAS OUTPERFORM LLMS
Across the evaluated model families and datasets, LLM-JEPA improves fine-tuning results and resists overfitting. The section also examines pretraining, representation structure, and alternative design choices.
4.1 FINE-TUNING AND PRETRAINING STRONGER GENERATIVE MODELS VIA JEPA
- LLM-JEPA improves performance across four model families, four datasets, six training epochs, and four model sizes, and resists overfitting, including with LoRA.The LoRA setting also retains performance gains across ranks.Source: §4.1 · p. 6
- On NL-RX-SYNTH pretraining, LLM-JEPA improves learned representation quality, though the small dataset prevents reliable generation termination and prompts a prefix-match validity criterion.The evaluation counts a generated solution as valid if it begins with the ground-truth sequence.Source: §4.1 · p. 6
- Pretraining on groups of five paraphrases yields statistically significant Rotten Tomatoes and Yelp gains after fine-tuning without JEPA loss, despite no direct dataset link.Each paraphrase version predicts the next version during pretraining.Source: §4.1 · pp. 6–7
4.2 STRUCTURED REPRESENTATIONS INDUCED BY LLM-JEPA
- LLM-JEPA yields structured Text and Code representations and a near-linear mapping between them, supported by smaller singular values and least-squares regression results.The authors present the near-linear mapping as a possible basis for accuracy improvements.Source: §4.2 · p. 7, Figure 4 · p. 7
4.3 ABLATION STUDY ON DESIGN CHOICES
- Cosine similarity, appending [PRED], and Text-to-Code prediction outperform the tested alternatives; InfoNCE has lower accuracy than the baseline and higher standard deviation.The alternatives include ℓ2-norm and MSE, prepending [PRED], and Code-to-Text prediction.Source: §4.3 · pp. 7–8
5 TOWARDS FASTER AND MORE GENERAL LLM-JEPAS
The section tests LLM-JEPA beyond the usual Text-and-Code setup, including question answering and reasoning benchmarks, and introduces loss dropout to reduce compute. LLM-JEPA improves accuracy in the reported evaluations, while dropout supports faster convergence and higher accuracy at the same compute budget.
5.1 BEYOND CODE: APPLYING LLM-JEPA TO Q&A DATASETS AND REASONING MODELS
- LLM-JEPA significantly improves accuracy on NQ-Open and HellaSwag, and both Qwen3-1.7B and DeepSeek-R1-Distill-Qwen-1.5B improve on GSM8K.The GSM8K improvements are statistically significant.Source: §5.1 · p. 8
- NQ-Open pairs questions with short answer spans, while HellaSwag pairs context with the correct continuation rather than a single-token answer label.The HellaSwag context and continuation are both integral parts of the question.Source: §5.1 · p. 8
5.2 FASTER LLM-JEPAS VIA LOSS DROPOUT
- Random JEPA-loss dropout helps LLM-JEPA converge faster and improve accuracy at the same compute; batch-level dropout avoids an extra forward pass for encoding both views.With dropout rate α, per-epoch cost is 2−α times standard fine-tuning; λ × (1−α) approximately constant is a co-tuning guideline.Source: §5.2 · pp. 8–9, Figure 5 · p. 9
6 CONCLUSION AND FUTURE WORK
LLM-JEPA improves abstract prompt representations while retaining generative capabilities, and preliminary pretraining results are promising. Training costs twice as much, with random loss dropout offered as a mitigation; two added hyperparameters also make tuning costly.
6 CONCLUSION AND FUTURE WORK
- LLM-JEPA improves abstract prompt representations while maintaining LLM generative capabilities across datasets and models.The experiments mostly focus on fine-tuning; preliminary pretraining experiments are described as promising.Source: §6 · p. 9
- Random JEPA-loss dropout saves PFLOPs and boosts accuracy relative to regular LLM-JEPA, according to the Table 6 caption.The conclusion identifies a twofold increase in training compute as the primary bottleneck and says random loss dropout mitigates it.Source: §6 · p. 9, Table 6 · p. 10
- Two additional hyperparameters make tuning costly because the best configuration may occur anywhere in the (λ, k) grid.Adjacent grid points often yield similar accuracy, suggesting a possible route to more efficient tuning.Source: §6 · p. 10
A.1 FASTER LORA CONVERGENCE
LoRA fine-tuning with LLM-JEPA achieves higher accuracy and converges faster than LoRA with the LLM loss alone. The supplied material also describes the comparison’s rank behavior, evaluation setup, and representation plots.
A.1 FASTER LORA CONVERGENCE
- LLM-JEPA LoRA fine-tuning achieves substantially higher accuracy and converges more quickly than using the LLM loss alone.At LoRA rank 512, LLM-JEPA reaches accuracy comparable to full fine-tuning, while LLM-loss-only LoRA retains a clear gap.Source: §A.1 · p. 15
- The Table 9 setup pretrains Llama-3.2-1B-Instruct on paraphrases, then fine-tunes and evaluates on Rotten Tomatoes and Yelp.LLM-JEPA is applied only during pretraining; the caption does not report outcome values.Source: Table 9 · p. 16
- Figure 6 shows LLM-JEPA inducing structure in Text and Code representations, while NTP fine-tuning disrupts the base-model structure.The panels compare the base model, NTP fine-tuning, and LLM-JEPA at k = 0 and k = 1.Source: Figure 6 · p. 16
A.2 LLM-JEPA INDUCES STRUCTURED REPRESENTATION
LLM-JEPA fine-tuning yields structured Text and Code representations, with similar patterns across the tested k values; standard fine-tuning appears to disrupt the base model’s structure. Accuracy across the (k, λ) sweep has no consistent best location.
A.2 LLM-JEPA INDUCES STRUCTURED REPRESENTATION
- Different k values yield similar Text and Code representation patterns under LLM-JEPA, while standard fine-tuning appears to further disrupt base-model structure.These observations are based on t-SNE plots.Source: §A.2 · p. 16
- With γ = 0.0, the NL-RX-SYNTH configuration in Table 10 generates only empty output.The caption specifies λ = 1, k = 0, and a fixed learning rate; it reports averages and standard deviations over five runs.Source: Table 10 · p. 17
- Accuracy has no consistent best location across λ or k, though adjacent grid cells usually have close accuracy.The sweep covers k ∈ {0, 1, 2, 3, 4} and λ ∈ {0.5, 1, 2, 4}; each cell averages five runs at epoch 4.Source: Figure 7 · p. 17
A.3 ABLATION STUDY ON THE ROLE OF LLLM
The ablation introduces γ to control the generative loss’s relative strength, alongside a JEPA abstraction term. When γ is zero, the model produces only empty outputs, indicating that the JEPA term complements rather than replaces the generative loss.
A.3 ABLATION STUDY ON THE ROLE OF LLLM
- γ explicitly controls the relative strength of the LLM loss because the prior formulation could not reduce its contribution to zero.The equation labels this term as serving generative capabilities.Source: §A.3 · p. 17
- The objective pairs the generative LLM loss with an abstraction-capability term weighted by λ.The abstraction term compares predicted Text embeddings with Code embeddings.Source: §A.3 · p. 17
- When γ = 0, the fine-tuned model produces only empty outputs, showing that the LLM loss remains essential for generative performance.The authors characterize the JEPA component as a regularization term that complements the generative loss.Source: §A.3 · p. 17
A.4 ADDITIONAL GENERATION EXAMPLES
The appendix adds generation examples, examines LoRA overfitting and learned representations, and reports performance across model sizes. Its results include continued LLM-JEPA improvement during LoRA fine-tuning, near-linear Text-to-Code representations, and significant accuracy gains across scales.
A.4 ADDITIONAL GENERATION EXAMPLES
- Table 11 provides additional regular-expression generations from Llama-3.2-1B-Instruct fine-tuned with LLLM and LLM-JEPA losses.Its color coding marks wrong, extra, and missing content.Source: §A.4 · p. 17, Table 11 · p. 18
A.5 OVERFITTING BEHAVIOR IN LORA FINE-TUNING
- During LoRA fine-tuning, LLM-JEPA accuracy generally keeps improving across epochs while standard LLLM shows clear overfitting.LoRA has higher standard deviations than full fine-tuning, likely reflecting its lower capacity; standard-deviation fluctuations may be unreliable indicators of generalization quality.Source: §A.5 · p. 18, Figure 8 · p. 18
A.6 STRUCTURED REPRESENTATIONS INDUCED BY LLM-JEPA
- LLM-JEPA promotes a near-linear transformation from Text to Code representations, alongside emerging structure and smaller singular values in their difference.The authors suggest this representation structure may underlie accuracy improvements.Source: §A.6 · pp. 18–19, Table 14 · p. 19
A.7 PERFORMANCE ACROSS MODEL SIZES
- LLM-JEPA shows statistically significant accuracy improvements at all evaluated scales, including Llama-3.1-8B under a startswith criterion.The 8B model’s exact-termination performance collapsed because of difficulty terminating regular expressions; Figure 9 also reports continued HellaSwag performance improvement as λ scales to 1024 for Llama-3.2-1B.Source: §A.7 · p. 19, Table 15 · p. 20, Figure 9 · p. 20