LLM-JEPA: Large Language Models Meet Joint Embedding Predictive Architectures

📄 arXiv: 2509.14252v2 📥 PDF

作者: Hai Huang, Yann LeCun, Randall Balestriero

分类: cs.CL, cs.AI

发布日期: 2025-09-11 (更新: 2025-10-07)

🔗 代码/项目: GITHUB

💡 一句话要点

提出LLM-JEPA，将联合嵌入预测架构应用于LLM的预训练和微调，显著提升性能。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 联合嵌入预测架构 预训练 微调 嵌入空间学习 对比学习 自然语言处理

📋 核心要点

  1. 现有LLM训练主要依赖输入空间重建，而视觉领域的联合嵌入预测架构（JEPA）表现更优，存在领域差异。
  2. LLM-JEPA将JEPA的思想引入LLM训练，通过嵌入空间预测提升模型性能，适用于预训练和微调。
  3. 实验结果表明，LLM-JEPA在多个数据集和模型上显著优于标准LLM训练方法，并具有更好的鲁棒性。

📝 摘要（中文）

大型语言模型（LLM）的预训练、微调和评估依赖于输入空间重建和生成能力。然而，在视觉领域，基于联合嵌入预测架构（JEPAs）的嵌入空间训练目标远优于其输入空间对应方法。语言和视觉训练方式的这种不匹配引出了一个自然的问题：语言训练方法能否从视觉方法中学习一些技巧？缺乏JEPA风格的LLM证明了为语言设计此类目标的挑战性。在这项工作中，我们提出了朝这个方向迈出的第一步，开发了LLM-JEPA，这是一种基于JEPA的LLM解决方案，适用于微调和预训练。到目前为止，LLM-JEPA能够在各种模型中显著优于标准LLM训练目标，同时对过拟合具有鲁棒性。这些发现在众多数据集（NL-RX、GSM8K、Spider、RottenTomatoes）和来自Llama3、OpenELM、Gemma2和Olmo系列的各种模型中观察到。

🔬 方法详解

问题定义：现有LLM的训练范式主要依赖于输入空间的重建和生成能力，例如预测下一个token。然而，这种方式可能不够高效，并且容易受到过拟合的影响。视觉领域的JEPAs方法在嵌入空间进行训练，已被证明更加有效。因此，如何将JEPAs的思想引入到LLM的训练中，克服传统方法的局限性，是一个亟待解决的问题。

核心思路：LLM-JEPA的核心思路是在LLM的训练过程中，不再直接预测输入空间的内容（例如下一个token），而是在嵌入空间中进行预测。具体来说，模型学习预测同一输入的两个不同视角（例如，经过不同噪声处理的版本）的嵌入表示。通过这种方式，模型可以学习到更加鲁棒和泛化的特征表示，从而提高性能。

技术框架：LLM-JEPA的整体框架包括两个主要部分：编码器和预测器。编码器负责将输入文本转换为嵌入表示。预测器接收编码器输出的嵌入表示，并预测同一输入的另一个视角的嵌入表示。训练过程中，通过最小化预测的嵌入表示与目标嵌入表示之间的差异来优化模型参数。该框架可以应用于LLM的预训练和微调阶段。

关键创新：LLM-JEPA最重要的创新点是将JEPAs的思想从视觉领域引入到LLM领域。与传统的输入空间预测方法相比，LLM-JEPA在嵌入空间进行预测，可以学习到更加鲁棒和泛化的特征表示。此外，LLM-JEPA的设计使得它可以很容易地应用于现有的LLM架构，而无需进行大规模的修改。

关键设计：LLM-JEPA的关键设计包括：1) 如何生成同一输入的两个不同视角。例如，可以通过随机masking、token shuffling等方式对输入文本进行处理。2) 如何定义嵌入表示之间的差异。可以使用余弦相似度、均方误差等指标来衡量预测的嵌入表示与目标嵌入表示之间的差异。3) 如何选择合适的编码器和预测器架构。可以使用Transformer、MLP等架构作为编码器和预测器。

📊 实验亮点

实验结果表明，LLM-JEPA在多个数据集（NL-RX、GSM8K、Spider、RottenTomatoes）和各种模型（Llama3、OpenELM、Gemma2和Olmo）上都取得了显著的性能提升。例如，在某些数据集上，LLM-JEPA的性能提升超过了5%。此外，实验还表明，LLM-JEPA对过拟合具有更好的鲁棒性，这意味着它可以更好地泛化到未见过的数据。

🎯 应用场景

LLM-JEPA具有广泛的应用前景，可以应用于各种自然语言处理任务，例如文本分类、机器翻译、文本生成等。该方法可以提高LLM的性能和鲁棒性，使其在实际应用中更加可靠。此外，LLM-JEPA还可以促进LLM的进一步研究，例如探索更加有效的嵌入空间训练方法。

📄 摘要（原文）

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.