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

Paper: arXiv:2509.14252; OpenReview ICLR 2026 Code: galilai-group/llm-jepa Code reference: main @ ea0017c6 (2026-04-15)

1. Motivation (研究动机)

这篇论文想解决的不是“再训练一个更大的 LLM”，而是一个训练目标层面的错位：当前 LLM 的 pretraining、finetuning 和 evaluation 大多仍围绕 input-space reconstruction / next-token prediction 展开，即让模型在 token 空间复原下一个 token；但 Yann LeCun 系列 JEPA 思路强调，许多有用抽象并不需要、也不应该强迫模型重构原始输入，而应在 embedding space 中学习不同 view 之间的可预测关系。视觉里的 I-JEPA、V-JEPA 已经证明了 latent-space prediction 可以减少对像素级细节的依赖并形成更抽象的表示；语言模型却长期缺少一个既保留生成能力、又把 JEPA-style embedding objective 接入 LLM 训练的简洁做法。

论文把这个问题落在一个很具体的语言场景：许多 NLP / coding 数据天然包含同一知识的两个 view，例如自然语言描述 Text 与正则表达式 / SQL / 答案程序 Code。标准 finetuning 会学习从 Text 生成 Code，但它只在 token 级别惩罚输出错误；它不要求 Text 的 hidden representation 与 Code 的 hidden representation 在 latent space 中形成简单、稳定、可预测的关系。因此，一个模型可能 next-token loss 下降，但它的表示空间仍然混乱，尤其是当评测需要精确匹配、执行正确性或 reasoning answer 时，token-space objective 可能把大量容量花在表面形式而不是抽象对应关系上。

研究这个问题有两个实际价值。第一，如果 JEPA loss 能作为 regularizer 改善表示结构，那么它可能在不牺牲生成能力的前提下提升 finetuning accuracy，这比单纯扩大模型或数据更便宜。第二，如果 Enc(Text) 到 Enc(Code) 的映射被压到更窄、更近似线性的子空间，模型在多种 “input view → output view” 任务上会更容易泛化：正则表达式、SQL、GSM8K final answer、paraphrase pretraining 后的情感分类，都可以被看作从一个语义 view 预测另一个语义 view。第三，LLM-JEPA 若能通过 attention mask 和 loss dropout 降低额外开销，就为把 JEPA 扩到更大模型或更长 pretraining 提供了工程入口。

Figure 1 解读：左侧汇总了跨模型、跨数据集的 finetuning accuracy；LLM-JEPA 在 Llama3 / Gemma2 / OpenELM 和 SYNTH、Spider、GSM8K、TURK 等任务上普遍高于标准 baseline。右侧展示随 epoch 增长的 overfitting 对比：标准 NTP finetuning 在后期更容易掉头，而 LLM-JEPA 仍保持或继续提升，说明 JEPA term 更像一种 representation regularizer，而不是简单增加训练步数。

2. Idea (核心思想)

核心 insight 是：对含有自然双视图的数据，不只让 LLM 从 Text 自回归生成 Code，还要要求模型在最后一层 hidden state 中能从 Text 的 embedding 预测 Code 的 embedding。也就是说，LLM-JEPA 把标准 next-token prediction 作为生成项保留下来，同时加入 JEPA latent prediction 项，让生成能力和抽象表示能力在同一个 causal LM 内协同训练。

这篇论文的关键创新并不是引入一个外部 encoder 或额外 projector，而是把 LLM 自身复用为 encoder 和 tied-weights predictor：Enc(·) 取最后一层最后 token hidden state；Pred(·) 通过在 Text 后追加若干 predictor special tokens，让同一个 Transformer 的自注意力继续处理该表示。当 $k=0$ 时 predictor 退化为恒等映射；当 $k>0$ 时最后一个 <|predictor_k|> token 的 hidden state 作为预测向量。

它与 SimCSE / sentence embedding contrastive learning 的根本区别在于，LLM-JEPA 不把模型变成一个只输出 embedding 的 encoder，也不放弃生成评测；它仍优化 token-level ${\mathcal{L}}_{\mathrm{LLM}}$，只是额外约束 latent space。它与常规 instruction finetuning 的区别也很明确：常规 finetuning 只关心输出 token 是否对，LLM-JEPA 还要求输入 view 和输出 view 在 representation space 中形成低维、近似线性的可预测关系。

3. Method (方法)

3.1 总体框架：LLM loss + JEPA loss

论文先定义标准 LLM 目标。给定 token 序列 ${\mathrm{Text}}_{1:L}$，普通 next-token prediction 可写成： ${\mathcal{L}}_{\mathrm{LLM}}({\mathrm{Text}}_{1:L-1},{\mathrm{Text}}_L)=\mathrm{XEnt}\left(\mathrm{Classifier}\left(\mathrm{Enc}({\mathrm{Text}}_{1:L-1})\right),{\mathrm{Text}}_L\right).$ LLM-JEPA 在此基础上加入 embedding-space prediction： ${\mathcal{L}}_{\mathrm{LLM}-\mathrm{JEPA}}=\underset{\text{generative capabilities (LLM)}}{\underbrace{\sum _{\ell =2}^L{\mathcal{L}}_{\mathrm{LLM}}({\mathrm{Text}}_{1:\ell -1},{\mathrm{Text}}_{\ell })}}+\lambda \underset{\text{abstraction capabilities (JEPA)}}{\underbrace{d(\mathrm{Pred}(\mathrm{Enc}(\mathrm{Text})),\mathrm{Enc}(\mathrm{Code}))}}.$ 其中 $\lambda \ge 0$ 控制 JEPA term 的权重，$d$ 是 embedding metric。论文正文说 metric 可选、并在文字中倾向 cosine similarity；released code 的默认实现是 $1-\mathrm{mean}(\cos (e_{text},e_{code}))$，同时保留 --jepa_l2、--jepa_mse、--infonce 作为 ablation 选项。

Figure 2 解读：左侧是方法结构：Text 和 Code 是同一知识的两个 view，LLM-JEPA 让 Text 的 representation 经 tied predictor 后接近 Code 的 representation。右侧给出两个任务例子：NL-RX-SYNTH 中 Text 是正则表达式的自然语言描述、Code 是 regex；Spider 中 Text 是数据库和查询描述、Code 是 SQL。这个图的重点是：JEPA loss 不依赖“图像增强”那类 view，而是利用语言/代码数据本来就有的 paired views。

3.2 Encoder 与 predictor 的具体含义

Enc(x) 在论文和代码中都不是一个新网络，而是 causal LM 最后一层最后 token 的 hidden state。对聊天模板数据，Text view 通常对应 user message，Code view 对应 assistant message；released code 在 load_and_prepare_dataset() 中分别构造 input_ids_user、input_ids_assistant、labels_user、labels_assistant。标准 full prompt 的 labels 用于 next-token training；user / assistant 两个 view 的 labels 多数位置被置为 -100，只作为取 hidden state 的输入。

Pred 的实现也尽量不增加参数。代码会把 <|predictor_1|> 到 <|predictor_10|> 加入 tokenizer；当命令行传入 --predictors k 时，它在 user content 后追加 <|predictor_k|>，并取最后一个 predictor token 的 hidden state 作为 $Pred(Enc(Text))$。这与论文中的 $k\in \{0,\dots ,K\}$ 对齐：$k=0$ 时没有额外 predictor token，$Pred(x)=x$；$k>0$ 时让 LLM 自身多做若干步 latent processing。

直觉上，LLM-JEPA 迫使模型把“输入问题/描述”压成一种能线性或近似线性到达“输出程序/答案”的表示。如果只有 NTP loss，模型可以靠局部 token 统计拟合输出，而不必让 input representation 与 output representation 对齐；JEPA loss 则把这种对齐变成训练信号。它不是替代生成 loss，因为单独用 JEPA term 不知道如何产生可解析的文本；它更像把表示空间整理成更利于生成任务的坐标系。

3.3 Custom additive attention mask

实现难点是：要同时得到完整 Text→Code 生成 loss、独立的 Enc(Text)、独立的 Enc(Code)。如果直接拼接 Text+Code，Code 的 hidden state 会 attend 到 Text，不能作为独立 view 的 target embedding；如果分别 forward，则成本高。论文给出 block-wise causal additive mask：每个 view 内部仍是 causal mask，不同 view 之间用 $-\infty$ 隔断注意力。

Released code 中 RepresentationTrainer._build_additive_mask(k) 构造 $k\times k$ 上三角 -inf causal mask；build_with_additive_mask() 为 batch 建立形状为 (batch_size * 2, 1, seq_length, seq_length) 的 4D mask，并把 full prompt 与 user+assistant block 合并到一次 model call 中。代码里有一个重要工程更新：README 明确说原始实现需要两次额外 forward 来分别 encode Text 和 Code，最新版可用 --additive_mask 通过 4D additive mask 合并；若不启用 --additive_mask，forward() 仍会分别对 full prompt、user view、assistant view 做 model call。

论文公式与 released code 实现差异：论文正文在 mask 小节写的是“two forward passes instead of three”，而当前 main@ea0017c6 的 README / finetune.py 已支持 --additive_mask，把 JEPA view encoding 合并为一个 batch-doubled forward；但 run.sh 中的主要示例命令没有默认传 --additive_mask，所以复现实验时需要明确确认是否启用该 flag。另一个差异是论文公式写通用 metric $d$，正文举例有 ${\ell }_2$，代码默认走 cosine distance，--jepa_l2 / --jepa_mse / --infonce 才切到 ablation metric。

3.4 Random JEPA-loss dropout 与 gamma/lambda

论文还讨论计算开销。普通 LLM-JEPA 因为要取 view embeddings，比标准 NTP 更贵；为了在相同 PFLOPs 下比较，作者提出 random JEPA-loss dropout：以概率 $\alpha$ 跳过 JEPA term，只做普通 LLM loss。代码用 --jepa_ratio 表示“保留 JEPA 的 batch 比例”，所以若论文中的 dropout rate 为 $LD=\alpha$，代码应设置 jepa_ratio=1-\alpha。README 给出的成本近似是：平均每步计算为 $(2-\alpha )X$ 的标准 finetuning 成本；--same_flop 会据此调整 checkpoint / 训练步数，用于同 FLOPs 对比。

附录还引入 $\gamma$ 控制生成项： ${\mathcal{L}}_{\mathrm{LLM}-\mathrm{JEPA}}=\gamma {\sum }_{\ell =2}^L{\mathcal{L}}_{\mathrm{LLM}}({\mathrm{Text}}_{1:\ell -1},{\mathrm{Text}}_{\ell })+\lambda d(\mathrm{Pred}(\mathrm{Enc}(\mathrm{Text})),\mathrm{Enc}(\mathrm{Code})).$ 这个 ablation 很关键：当 $\gamma =0$ 时模型只生成空输出，NL-RX-SYNTH accuracy 为 $0.00\pm 0.00$。因此 LLM-JEPA 不能被理解成“去掉 autoregressive training”；它必须保留 LLM loss，JEPA term 的作用是 regularization 和 representation shaping。

3.5 Source-grounded pseudocode

def prepare_jepa_views(example, tokenizer, predictors: int, reverse_pred: bool = False):
    messages = example["messages"]
    text_msgs = get_assistant_messages(messages) if reverse_pred else get_user_messages(messages)
    code_msgs = get_user_messages(messages) if reverse_pred else get_assistant_messages(messages)
 
    if predictors > 0:
        text_msgs[0]["content"] += f"<|predictor_{predictors}|>"
 
    full = tokenizer.apply_chat_template(messages, tokenize=True)
    text = tokenizer.apply_chat_template(text_msgs, tokenize=True)
    code = tokenizer.apply_chat_template(code_msgs, tokenize=True)
 
    return {
        "input_ids": full.input_ids,
        "labels": mask_user_tokens(full.input_ids, messages),
        "attention_mask": full.attention_mask,
        "input_ids_user": text.input_ids,
        "labels_user": [-100] * len(text.input_ids),
        "attention_mask_user": text.attention_mask,
        "input_ids_assistant": code.input_ids,
        "labels_assistant": [-100] * len(code.input_ids),
        "attention_mask_assistant": code.attention_mask,
    }

def build_additive_mask_batch(batch, last_token_fn, jepa_ratio: float):
    if jepa_ratio > 0.0 and torch.rand(1).item() > jepa_ratio:
        return only_full_prompt(batch), True
 
    bsz, seq_len = batch["input_ids"].shape
    mask = torch.full((bsz * 2, 1, seq_len, seq_len), -torch.inf, device=batch["input_ids"].device)
    full_last = last_token_fn(batch["input_ids"], batch["labels"], batch["attention_mask"])
    user_last = last_token_fn(batch["input_ids_user"], batch["labels_user"], batch["attention_mask_user"])
    code_last = last_token_fn(batch["input_ids_assistant"], batch["labels_assistant"], batch["attention_mask_assistant"])
 
    for i in range(bsz):
        full_len = full_last[i] + 1
        user_len = user_last[i] + 1
        code_len = code_last[i] + 1
        batch["input_ids_user"][i, user_len:user_len + code_len] = batch["input_ids_assistant"][i, :code_len]
        mask[i, :, :full_len, :full_len] = causal_additive_mask(full_len)
        mask[i + bsz, :, :user_len, :user_len] = causal_additive_mask(user_len)
        mask[i + bsz, :, user_len:user_len + code_len, user_len:user_len + code_len] = causal_additive_mask(code_len)
 
    return {
        "input_ids": torch.cat([batch["input_ids"], batch["input_ids_user"]], dim=0),
        "labels": torch.cat([batch["labels"], batch["labels_user"]], dim=0),
        "attention_mask": mask,
    }, False

def llm_jepa_loss(model, batch, gamma: float, lbd: float, metric: str = "cosine"):
    outputs = model(**batch["llm_inputs"], output_hidden_states=True)
    lm_loss = outputs.loss
 
    text_hidden = batch["text_hidden"]
    code_hidden = batch["code_hidden"]
    text_emb = text_hidden[torch.arange(text_hidden.shape[0]), batch["text_last"], :]
    code_emb = code_hidden[torch.arange(code_hidden.shape[0]), batch["code_last"], :]
 
    if metric == "l2":
        jepa = torch.linalg.norm(text_emb - code_emb, ord=2, dim=-1).mean()
    elif metric == "mse":
        jepa = torch.mean((text_emb - code_emb) ** 2)
    elif metric == "infonce":
        logits = F.cosine_similarity(text_emb[:, None, :], code_emb[None, :, :], dim=-1) / 0.07
        labels = torch.arange(logits.size(0), device=logits.device)
        jepa = F.cross_entropy(logits, labels)
    else:
        jepa = 1.0 - F.cosine_similarity(text_emb, code_emb, dim=-1).mean()
 
    return gamma * lm_loss + lbd * jepa

def train_with_jepa_dropout(trainer, batch, alpha: float):
    jepa_ratio = 1.0 - alpha
    if torch.rand(1).item() > jepa_ratio:
        outputs = trainer.model(input_ids=batch["input_ids"], labels=batch["labels"], attention_mask=batch["attention_mask"])
        return outputs.loss
    return trainer.compute_loss(trainer.model, batch)

Code reference: main @ ea0017c6 (2026-04-15) — pseudocode and mapping based on this commit

Paper Concept	Source File	Key Class/Function
LLM-JEPA finetuning loop	finetune.py	RepresentationTrainer.forward, compute_loss
Text / Code view construction	finetune.py	load_and_prepare_dataset, get_user_messages, get_assistant_messages
Predictor tokens	finetune.py	setup_model_and_tokenizer, --predictors
Additive attention mask	finetune.py	_build_additive_mask, build_with_additive_mask, --additive_mask
JEPA metric ablations	finetune.py	--jepa_l2, --jepa_mse, --infonce
Random JEPA-loss dropout / same FLOPs	finetune.py, README.md	--jepa_ratio, --same_flop, FLOPs notes
Standard / JEPA experiment drivers	run.sh, run8bh200.sh	run_regular, run_jepa
Evaluation	evaluate.py	eval, spider_eval, exact / execution matching
Later STP code in same repo	stp.py, run_stp.sh	LinearPredictor, RepresentationTrainer; not the core arXiv 2509.14252 LLM-JEPA path

4. Experimental Setup (实验设置)

Datasets and scale. 论文主实验覆盖 NL-RX-SYNTH、NL-RX-TURK、GSM8K、Spider，并扩展到 NQ-Open、HellaSwag、paraphrase pretraining 后的 Rotten Tomatoes / Yelp。Released repo 的 JSONL 规模为：SYNTH 8000 train / 2000 test，TURK 8000 / 2000，GSM8K 7473 / 1319，Spider 6587 / 1447，NQ-Open 8000 / 1769，HellaSwag 8042 / 2000，paraphrase 330276 train-only，Rotten Tomatoes 8530 / 1066，Yelp 8000 / 2000。Spider 还需要 repo 中 spider_data.zip 的数据库用于执行 SQL 评测。

Baselines. 主要 baseline 是标准 next-token prediction / regular finetuning，即只优化 ${\mathcal{L}}_{\mathrm{LLM}}$。模型族包括 Llama-3.2-1B-Instruct、Llama-3.2-3B-Instruct、Llama-3.1-8B-Instruct、Gemma-2-2B-it、OpenELM-1_1B-Instruct、OLMo-2-0425-1B-Instruct、OLMo-2-1124-7B-Instruct，以及 reasoning model Qwen3-1.7B 与 R1-Distill-Qwen-1.5B。LoRA 实验把 rank 32/64/128/256/512 与 full fine-tuning 对比。

Metrics. NL-RX-SYNTH / TURK 使用生成 regex 的 exact match；GSM8K 使用 final result exact match；Spider 使用生成 SQL 的 execution-result exact match；NQ-Open / HellaSwag / Rotten Tomatoes / Yelp 使用对应任务的 accuracy。论文报告五次运行的 mean ± std，并给 paired single-tailed $t$-test 的 $p$-value。

Training config from actual scripts, not defaults. run.sh 的主 driver 使用 torchrun --nproc_per_node=8 finetune.py，固定 --num_epochs=4，seeds 为 82, 23, 37, 84, 4；Llama GSM8K / Spider 示例使用 learning_rate=1e-5、last_token=-2，扫描 lbd in {0.5, 2.0} 与 predictors in {0,1,2}。

run8bh200.sh 用 torchrun --nproc_per_node=8 finetune8bh200.py，README 说明该路径面向 NVIDIA H200 上的 up-to-8B 模型；示例为 OLMo-2-1124-7B-Instruct、learning_rate=8e-5、dataset=synth、lbd=2.0、predictors=0..4、last_token=-1。

finetune.py 的 CLI 默认值是 batch_size=4、grad_accum=4、max_length=512、bf16=True、gradient_checkpointing=True、dataloader_drop_last=True、report_to=none；这些是代码默认，不等同于论文表格中每个最佳配置，论文表格的 LR / $\lambda$ / $k$ 来自具体 grid search。

Hyperparameter search. 论文先为每个 (model, dataset) 找最佳 baseline learning rate，搜索 $lr\in \{1e-5,2e-5,4e-5,8e-5\}$；然后固定该 LR，搜索 LLM-JEPA 专属 $(k,\lambda )\in \{0,1,2,3,4\}\times \{0.5,1,2,4\}$。附录指出最佳点没有简单模式，局部相邻点通常接近，但 tuning 仍是主要成本。

5. Experimental Results (实验结果)

5.1 主结果：fine-tuning 与 pretraining 均提升

在 NL-RX-SYNTH pretraining 上，Llama-3.2-1B-Instruct 的 ${\mathcal{L}}_{\mathrm{LLM}}$ accuracy 为 $54.38\pm 1.70$，LLM-JEPA 为 $60.59\pm 1.01$，$p=2.94e-4$，配置为 $lr=8e-5,\lambda =2,k=3$。在 NL-RX-SYNTH finetuning 的模型族对比中，Llama-3.2-1B-Instruct 从 $57.29\pm 5.32$ 到 $71.46\pm 1.34$；Gemma-2-2B-it 从 $33.65\pm 3.24$ 到 $43.12\pm 2.61$；OpenELM-1_1B-Instruct 从 $12.07\pm 1.81$ 到 $25.40\pm 2.40$；OLMo-2-0425-1B-Instruct 从 $87.09\pm 0.36$ 到 $87.52\pm 0.29$。这些数值说明收益在弱模型/困难任务上更大，在已经接近饱和的 OLMo 上更小但仍为正。

对 Llama-3.2-1B-Instruct 的跨数据集结果也一致：NL-RX-TURK 从 $22.49\pm 1.91$ 到 $30.94\pm 1.13$；GSM8K 从 $32.36\pm 0.58$ 到 $36.36\pm 0.20$；Spider 从 $47.52\pm 2.44$ 到 $50.55\pm 2.08$。在更偏 QA / reasoning 的扩展上，NQ-Open 从 $20.12\pm 0.41$ 到 $21.59\pm 0.40$，HellaSwag 从 $69.40\pm 0.99$ 到 $70.51\pm 1.20$；Qwen3-1.7B 的 GSM8K 从 $44.32\pm 0.39$ 到 $45.00\pm 0.40$，R1-Distill-Qwen-1.5B 从 $13.87\pm 1.01$ 到 $15.04\pm 0.15$。在 paraphrase pretraining 之后再 finetune，Rotten Tomatoes 从 $56.57\pm 1.66$ 到 $57.76\pm 1.33$，Yelp 从 $26.46\pm 0.92$ 到 $27.15\pm 0.93$。

5.2 表示结构证据：低奇异值与 t-SNE 分离

Figure 3 解读：左侧显示 $Enc(Text)-Enc(Code)$ 的 top-100 singular values。LLM-JEPA 曲线显著低于 base 和普通 finetuning，支持作者的解释：Text → Code 的映射被约束在更窄子空间内。右侧同时监控 next-token loss 与 JEPA prediction loss；普通 ${\mathcal{L}}_{\mathrm{LLM}}$ 能降低 NTP loss，但不会自动降低 prediction loss，而 LLM-JEPA 在保持类似 NTP loss 的同时显著降低 JEPA loss，所以 accuracy gap 更可能来自 representation shaping，而非 token loss 本身。

Figure 4 解读：这组 t-SNE 对比说明普通 NTP finetuning 会扰动原本的 Text / Code 结构，而 LLM-JEPA（图中 $k=0$）会让 paired views 的表示形成更清晰、更可解释的几何结构。它不是直接证明 causal mechanism，但与 SVD 和 least-squares 结果相互印证。

> Figure omitted from public site because `fig6_tsne_full.svg` is 33.1 MiB and exceeds the Cloudflare Pages single-file limit; the original asset remains in the local Obsidian vault.

Figure 6 解读：附录的完整 t-SNE 把 base、普通 NTP、LLM-JEPA $k=0$、LLM-JEPA $k=1$ 放在一起。base model 尚有一定结构；普通 finetuning 后结构被打散；LLM-JEPA 两种 $k$ 都能恢复或强化 paired representation 的对应关系。附录 least-squares 表进一步给出数值：base 的 ${\min }_X\Vert Enc(Text)X-Enc(Code){\Vert }_2$ 为 $3953.11$，普通 ${\mathcal{L}}_{\mathrm{LLM}}$ 为 $3035.01$，LLM-JEPA $k=1$ 为 $4.47$，$k=0$ 为 $4.04$；Avg. Top 100 Singular 也从 $310.73/341.80$ 降到 $94.84/16.82$。

5.3 Ablations：哪些设计真的重要

设计 ablation 在 NL-RX-SYNTH 上使用 $lr=2e-5,\lambda =1.0,k=1$。Baseline 为 $57.29\pm 5.32$，标准 LLM-JEPA 为 $71.46\pm 1.34$。改用 ${\ell }_2$-norm 会崩到 $2.22\pm 0.07$，MSE 仍有 $70.64\pm 2.05$，把 predictor token prepend 到 Text 为 $68.07\pm 2.57$，反向用 Code→Text 为 $65.70\pm 2.63$，InfoNCE 为 $34.40\pm 6.10$。这表明：方向、metric 和 predictor placement 都会影响效果；默认 cosine-like positive-pair distance 比 InfoNCE 更适合当前设置。

$\gamma /\lambda$ ablation 说明生成项不可删除。在 $lr=2e-5,\lambda =1,k=0$ 并固定 $\max (\gamma ,\lambda )=1$ 时，$\gamma /\lambda =0$ 得到 $0.00\pm 0.00$，$0.01$ 得到 $1.38\pm 0.06$，$0.1$ 得到 $45.80\pm 5.04$，$1.0$ 得到 $70.42\pm 2.36$，$10.0$ 得到 $67.52\pm 1.45$，$100.0$ 得到 $66.83\pm 3.89$，$\infty$（即 $\lambda =0$）回到 $57.29\pm 5.32$。最清楚的结论是：JEPA term 不能独立承担 generation；它需要足够强的 LLM loss 来保持输出格式与可读性。

5.4 计算效率：same PFLOPs 与 loss dropout

Figure 5 解读：这张图把 regular finetuning、标准 LLM-JEPA、以及不同 $LD$ 的 JEPA-loss dropout 放在相同 PFLOPs 轴上比较。即使考虑额外 forward 成本，LLM-JEPA 仍更快达到高 accuracy；更激进的 $LD=0.75$ 在许多点上接近或超过无 dropout 的 JEPA，说明并非每个 mini-batch 都需要 JEPA term 才能形成有效 regularization。

Random dropout 的表格也支持这一点：在每行相隔 4.83 PFLOPs 的设置中，最后一行 $LD=0,\lambda =1$ 为 $63.96\pm 1.07$，$LD=0.5,\lambda =2$ 为 $72.11\pm 2.18$，$LD=0.75,\lambda =4$ 为 $73.08\pm 1.28$。作者观察到 co-tuning 时保持 $\lambda (1-\alpha )$ 近似稳定是一个有用经验，因为 dropout 减少 JEPA 更新频率后，可以适当提高 $\lambda$ 补偿单次 JEPA signal。

5.5 LoRA、模型尺度与额外例子

Figure 7 解读：六个训练曲线覆盖 Llama on GSM8K、Llama on Spider、Gemma / OpenELM / OLMo on SYNTH、以及 Llama pretraining。它们共同展示了 LLM-JEPA 的两个现象：一是 early epoch 上升更快；二是在 baseline 已经开始震荡或过拟合时，LLM-JEPA 往往继续保持更稳的提升。不同模型和 LR 的曲线差异也提醒读者：$\lambda$、$k$、LR 的最优点不是通用常数。

LoRA 表格显示 LLM-JEPA 对低秩适配同样有帮助。在 NL-RX-SYNTH、$lr=2e-5,\lambda =1,k=1$ 下，rank 32 从 $6.09\pm 0.55$ 到 $7.45\pm 1.87$；rank 64 从 $21.09\pm 1.90$ 到 $32.46\pm 1.26$；rank 128 从 $34.21\pm 2.82$ 到 $48.45\pm 3.66$；rank 256 从 $45.57\pm 4.52$ 到 $60.80\pm 2.31$；rank 512 从 $50.18\pm 5.15$ 到 $72.41\pm 2.94$。Full fine-tuning 则为 $57.29\pm 5.32$ vs $70.42\pm 2.36$。作者特别指出 rank 512 的 LLM-JEPA 已接近 full fine-tuning 的 LLM-JEPA，而普通 LoRA 仍明显落后。

Figure 8 解读：LoRA 情况下 LLM-JEPA 也比普通 ${\mathcal{L}}_{\mathrm{LLM}}$ 更抗过拟合，但作者承认趋势不如 full fine-tuning 稳定，可能因为 LoRA 的可训练容量有限。图中的标准差波动说明低秩约束会放大随机性，不能只看单次 run。

模型尺度方面，Llama-3.2-3B-Instruct 在 NL-RX-SYNTH 上从 $74.55\pm 3.58$ 到 $77.16\pm 3.66$；Llama-3.1-8B-Instruct 从 $35.77\pm 6.60$ 到 $63.57\pm 16.81$；OLMo-2-1124-7B-Instruct 从 $87.26\pm 0.27$ 到 $87.75\pm 0.33$。HellaSwag 的 $\lambda$ scaling 图显示 Llama-3.2-1B 可以把 $\lambda$ 扩到 1024 仍继续提升；但附录 $\gamma /\lambda$ 表也说明过度压低生成项会失败，所以不能把“更大 $\lambda$”泛化成无条件规则。

Figure 9 解读：这张 HellaSwag 曲线展示某些设置下 $\lambda$ 可以非常大，仍带来 accuracy 增益。它的正确解读不是“永远把 $\lambda$ 调大”，而是说明 JEPA regularization 的可用强度范围比常见小权重正则项更宽；最佳 $\lambda$ 仍依赖 dataset、model、$k$ 和是否保留足够 $\gamma$。

5.6 Limitations and conclusion

论文最重要的 limitation 是调参成本和计算成本。LLM-JEPA 新增 $k$、$\lambda$，附录明确说没有发现最佳点的简单规律，只能观察到相邻 grid point 多数时候相近；这意味着每个新 (model, dataset) 都要额外 sweep。计算上，若不使用 additive mask 或 dropout，JEPA view encoding 会显著慢于普通 finetuning；当前代码虽有 --additive_mask 和 --jepa_ratio，但主脚本是否启用需要复现者自己检查。

第二个 limitation 是适用场景依赖 paired views。LLM-JEPA 在 Text/Code、question/answer、paraphrase 这类天然双视图数据上定义清楚；对于没有明显 target view 的开放式 instruction data，需要设计 view construction，否则 JEPA target 可能退化成噪声。第三，论文的“结构更好”证据主要来自 SVD、t-SNE、least-squares 与 accuracy 的相关性，还不是严格因果证明；不过 loss 曲线中 NTP loss 基本重合、JEPA prediction loss 显著下降，是支持 causal story 的较强证据。

总体结论：LLM-JEPA 是把 LeCun 式 JEPA 从 vision latent prediction 推进到 causal LLM 的一个实用版本。它最有价值的部分是把生成 objective 与 representation objective 放进同一个 LLM 训练循环：生成项保证模型还能输出可评测文本，JEPA 项把 paired views 的 hidden states 拉到可预测的几何结构中。实验显示这种结构约束在多模型、多数据集、LoRA、pretraining+finetuning、reasoning model 上大多带来稳定正收益；工程上仍需解决 hyperparameter sweep、mask/default 脚本一致性和大规模训练成本。