Alessio Borgi

Alessio Borgi

PhD in Graph Neural Networks and Generative AI 👨🏻‍💻🤖

LLM-JEPA: Add Representation Prediction Without Losing Text Generation

3 minute read

Published:

TL;DR: LLM-JEPA adds an embedding-prediction objective to an existing language-model objective. Related sequences provide different views of the same underlying task. The model continues to generate tokens; the additional representation loss changes its training signal.
First page of LLM-JEPA: Large Language Models Meet Joint Embedding Predictive Architectures
Paper preview: LLM-JEPA: Large Language Models Meet Joint Embedding Predictive Architectures.
Hai Huang, Yann LeCun, and Randall Balestriero
Read the paper (PDF)

A description and a program can express the same requirement

Intuition first. “Select the names of customers with overdue invoices” and a corresponding SQL query are different strings that can describe the same operation. Their relationship supplies a training pair at a level above individual next-token choices.

Huang, LeCun, and Balestriero use such related views to add JEPA training to language models. A prediction made from one view is matched to an embedding of the other, alongside the original language-model loss. Paired views are an important part of the setup; arbitrary neighbouring sentences are not guaranteed to express the same information.

Two objectives supervise different outputs

Write \(x\) for the source view and \(y\) for its paired target. A compact form of the objective is

\[ \begin{aligned} \hat z_y&=\operatorname{Pred}(\operatorname{Enc}(x)),\\ z_y&=\operatorname{Enc}(y),\\ \mathcal L&=\mathcal L_{\mathrm{LM}}+\lambda d(\hat z_y,z_y). \end{aligned} \]

The paper uses the last token’s final hidden state as a sequence representation. Appended prediction tokens let the same language-model weights perform the prediction computation. Related views are encoded separately, avoiding direct attention from the source view into its target. The method section and official code specify the construction and cosine-based comparison.

Two complementary training signalsThe embedding prediction is compared with the separately encoded target. Token generation is supervised by an additional language-model objective. Source viewPredicted embeddingTarget-view embedding
The embedding prediction is compared with the separately encoded target. Token generation is supervised by an additional language-model objective.

Worked example: interpret the loss components separately

For an illustrative cosine-distance convention, set \(d(u,v)=1-u^\top v\) for unit vectors. Let the predicted embedding be \((1,0)\) and its target \((0.8,0.6)\). Both have unit norm, so their distance is \(1-0.8=0.2\).

With language loss 1.4 and an illustrative weight \(\lambda=0.3\),

\[ \mathcal L=1.4+0.3(0.2)=1.46. \]

If the embeddings align perfectly, the extra term becomes zero, but the language loss may still be large. Conversely, a model can improve token predictions without satisfying the cross-view relation. Plotting the components separately reveals which task is improving.

Loss scale matters: changing a token loss from a sum to a mean changes its balance with the embedding term. The numerical value of \(\lambda\) is meaningful only together with the reduction and normalisation conventions.

Training-time information is not inference-time input

The paired answer supplies supervision during training. At inference, a user normally supplies only the source prompt and the model generates its answer. A useful evaluation must respect that boundary.

For a text-to-SQL example, providing the correct SQL to the embedding branch during evaluation would test a different system from one that must generate it. Separate representation diagnostics, which may use paired data, from the actual generation task.

The added encoding passes also cost training computation. The paper’s design does not require those paired-view passes for ordinary generation, but that does not imply cost-free training or the elimination of autoregressive decoding.

What evidence would establish a useful representation?

A close embedding pair does not guarantee an executable or correct program. Two SQL statements may have similar wording and different filters. Evaluate the generated operation on appropriate tests or benchmark protocols, and inspect whether the representation preserves those critical differences.

Key Insight: LLM-JEPA combines token-level supervision with a relation between whole views. It adds a way to train representations while retaining the language model's generative task.

Compare VL-JEPA, which predicts an answer embedding from visual input, and LeJEPA for explicit distribution regularisation.

References

  1. Huang, H., LeCun, Y., & Balestriero, R. LLM-JEPA: Large Language Models Meet Joint Embedding Predictive Architectures. 2025.
  2. Authors’ LLM-JEPA implementation.

You May Also Enjoy

The Standard Library Modules Worth Knowing Before You pip install

8 minute read

Published:

Python ships with a counter, a queue, a memoiser, a CLI parser, a test framework’s smaller cousin, and a JSON-safe date library already installed, most of what people reach for a package to solve is one import away.

Files and Context Managers: Why with Is Not Optional

9 minute read

Published:

A file object that falls out of scope without with does get closed eventually, but ‘eventually’ means whenever the garbage collector gets around to it, which is not a promise any program handling more than a handful of files can live with.

Errors and Exceptions: EAFP, the Hierarchy, and Reading a Traceback

9 minute read

Published:

Exceptions in Python are not exceptional. They are a normal control-flow mechanism that the language leans on so heavily that the idiomatic style is to try the operation and handle the failure, rather than check first, which is faster, shorter, and free of race conditions.