PaperlayerView exact revision on arXiv

Source-linked AI summary

ER-JEPA: Experience Replay Improves Joint-Embedding Predictive Learning in Language Models

Jingnan Pu, Zi-En Fan, Feng Lian

arXiv:2609.36952v1cs.CLcs.AI

TL;DR

Token-level objectives may favor local coherence over representations for comprehensive perception and reasoning, and replay addresses the gap between scheduled example visits. ER-JEPA adds an episodic replay path to LLM-JEPA, and it outperforms LLM-JEPA across multiple datasets and replay policies at matched compute.

  • Problem

    Next-token training may favor local coherence over representations for comprehensive perception and reasoning, while examples are supervised only when they appear in a mini-batch.

    Source: §1 · pp. 1–2

  • Method

    ER-JEPA extends LLM-JEPA with an episodic path that revisits stored training examples to add supervision for token prediction and representation alignment.

    Source: §1 · p. 2

  • Results

    ER-JEPA consistently outperforms LLM-JEPA across multiple datasets and, on SYNTH at 240.31 PFLOPs, content replay achieves 83.65 ± 1.48% accuracy versus 68.75 ± 4.92%.

    Source: §1 · p. 2, §3.2 · p. 5, §5 · p. 9

  • Takeaways & Limitations

    The replay branch, rather than a specific retrieval rule, is identified as the central source of improvement.

    Source: §1 · p. 2

    Replay-weight tuning requires grid search, adding computational cost, and ER-JEPA retains LLM-JEPA’s need for datasets with paired views.

    Source: §5 · p. 9

Abstract

from arXiv · show

Large language models (LLMs) excel at token-level generation but may learn undesirable abstract semantics and lack comprehensive perception. LLM-JEPA mitigates this by aligning different views of the same underlying knowledge via a joint-embedding predictive architecture (JEPA). However, strong alignment does not necessarily lead to accurate, stable predictions. To address this, we propose ER-JEPA, which adds an episodic replay path to LLM-JEPA. ER-JEPA stores training pairs in a memory. At each step, it stores and retrieves relevant data to provide additional supervision. This enables learning from both the current batch and stored training pairs, providing additional supervision for token prediction and representation alignment. Experiments across multiple datasets (NL-RX, GSM8K, Spider, and NQ-Open) demonstrate that ER-JEPA consistently outperforms LLM-JEPA.

1 Introduction

Token-level training can favor local coherence over representations for comprehensive perception and reasoning, while strong alignment alone does not ensure reliable predictions. ER-JEPA adds replay of past examples to LLM-JEPA and consistently outperforms it across multiple datasets.

1 Introduction

  • Autoregressive next-token prediction supports fluent generation, but may favor local coherence over representations needed for comprehensive perception and reasoning.Source: §1 · p. 1
  • LLM-JEPA aligns different views of shared knowledge, yet errors persist after alignment converges and some correct predictions become incorrect during training.This motivates asking how to correct persistent errors and retain correct predictions.Source: §1 · p. 2
  • ER-JEPA replays stored training examples during later updates, providing added supervision for token prediction and representation alignment; the replay path is removed at inference.Source: §1 · p. 2
  • Content-based, uniform, and hard selection share the replay path, and the paper identifies that path—not a particular retrieval rule—as the central source of improvement.Source: §1 · p. 2
  • ER-JEPA consistently outperforms LLM-JEPA across multiple datasets, supporting benefits for abstraction and generalization.Source: §1 · p. 2

2 Methodology

ER-JEPA retains LLM-JEPA’s token-prediction and representation-alignment objectives and adds replay of stored examples. The replay branch selects past examples, recomputes their representations with the current model, and supplies additional training signals.

2.1 ER-JEPA Overview and Objective

  • Figure 2 contrasts next-token training, LLM-JEPA’s embedding-space prediction, and ER-JEPA’s added replay path for token and representation supervision.Source: Figure 2 · p. 3
  • The objective retains LLM-JEPA’s generative and JEPA abstraction losses, with λ weighting the JEPA term and β controlling replay strength.Source: §2.1 · pp. 2–3

2.2 Replay-Based Training

  • Content-based, uniform, and hard replay share an objective; across checkpoints, ER-JEPA has more correct predictions despite slightly lower cosine similarities.Source: §2.2 · p. 3, Figure 3 · p. 4
  • Each step selects up to R memory entries, concatenates their stored source and target tokens, and recomputes their representations with the current model.Source: §2.2 · pp. 3–4
  • The replay loss is combined with the current-data loss to update the shared model, and the episodic pathway is removed at inference.Source: §2.2 · p. 4

3 Experiments: Episodic Replay Improves LLM-JEPA

ER-JEPA extends LLM-JEPA with episodic replay and outperforms it across five datasets. The experiments examine compute-matched performance, the value of historical samples, prediction quality, and inference cost.

3.1 Experimental Setup

  • Experiments use Llama-3.2-1B-Instruct on five datasets, measuring exact match for regex generation, math answers, and open-domain QA, and execution accuracy for Spider.The datasets are NL-RX-SYNTH, NL-RX-TURK, GSM8K, Spider, and NQ-Open.Source: §3.1 · p. 5

3.2 ER-JEPA Outperforms LLM-JEPA

  • ER-JEPA achieves higher performance than LLM-JEPA on all five datasets, spanning structured generation, mathematical reasoning, and open-domain question answering.The results suggest episodic replay improves the use of training data in LLM-JEPA.Source: §3.2 · p. 5
  • At 240.31 PFLOPs on SYNTH, content replay reaches 83.65 ± 1.48% accuracy versus 68.75 ± 4.92% for LLM-JEPA, a 14.90-point gain.Normalized AUC rises from 56.61 ± 3.41% to 73.73 ± 1.15%; all three replay policies outperform LLM-JEPA at every evaluated compute budget.Source: §3.2 · p. 5, Figure 4 · p. 6

3.3 Understanding the Role of Historical Replay

  • With matched token exposure, content replay reaches 83.65±1.48% accuracy versus 62.25±20.80% for token-matched LLM-JEPA.The authors conclude that extra token exposure alone does not reproduce the gain.Source: §3.3 · p. 6
  • Against the current-batch control, uniform replay reaches 85.04±0.54%, 3.80 percentage points higher, and is higher in all five seeds.Content and hard replay are also higher than the control by 2.41 and 2.17 percentage points, respectively.Source: §3.3 · p. 6
  • Replay maintains gains more consistently during training; ER-JEPA also improves correction of baseline errors, with over-generation correction rising from 53.07% to 84.51%.The rate of correct predictions becoming incorrect falls from 8.85% with LLM-JEPA to 6.65% with ER-JEPA.Source: §3.2 · p. 5, §3.3 · pp. 6–7, Figure 4 · p. 6

3.4 Fine-grained Analysis of Prediction Quality

  • ER-JEPA has higher mean exact-match accuracy than LLM-JEPA in all six regex-structure subsets and all four length groups.For expressions longer than 20 tokens, ER-JEPA reaches 78.53% versus 56.76% for LLM-JEPA; its shortest-to-longest accuracy drop is 11.27 percentage points versus 28.39 for LLM-JEPA.Source: §3.4 · pp. 7–8
  • After training, ER-JEPA removes the episodic memory and replay path and uses the same generation procedure as LLM-JEPA, adding no inference overhead.The inference procedure requires no additional computation or memory.Source: §3.4 · p. 8

4 Related Work

JEPA methods predict representations rather than reconstructing every input detail, with applications across vision and language. ER-JEPA builds on language-model JEPA by adding replay of past examples for additional training supervision.

JEPA foundations

  • JEPA predicts abstract states from context, while data2vec applies masked-input representation prediction across speech, vision, and language.LeCun proposed JEPA as a framework for learning world models without reconstructing every input detail.Source: §4 · p. 9

JEPA in vision

  • I-JEPA predicts target image-region representations from context, and V-JEPA extends latent prediction to video for downstream tasks with a frozen encoder.Other work broadens image prediction to photometric transformations and analyzes latent-prediction dynamics.Source: §4 · p. 9

JEPA in language and vision-language learning

  • LLM-JEPA combines next-token prediction with a JEPA loss over paired textual views; ER-JEPA adds replay of past examples for token-prediction and representation-alignment supervision.VL-JEPA predicts continuous text embeddings for vision-language tasks.Source: §4 · p. 9

5 Conclusion

ER-JEPA extends LLM-JEPA with replay of stored training pairs, and at matched training compute it outperforms LLM-JEPA across the three selection rules. The authors describe the current implementation as an initial step and identify loss-weight tuning and paired-view data requirements as open constraints.

5 Conclusion

  • ER-JEPA adds an episodic replay path that revisits stored training pairs and provides additional learning signals.The replay path is removed after training, leaving the same architecture and inference cost as LLM-JEPA.Source: §5 · p. 9
  • At matched training compute, ER-JEPA outperforms LLM-JEPA with all three selection rules.Replay also corrects more errors during training and avoids most large accuracy drops.Source: §5 · p. 9
  • The replay-loss weight β requires grid-search selection, adding computational cost, and ER-JEPA inherits the need for datasets with paired views.The authors suggest adaptive weighting and more efficient, general data augmentation as future directions.Source: §5 · p. 9

A Detailed Experimental Setup

The experiments follow the LLM-JEPA protocol to isolate the effect of episodic replay. The objectives share a LoRA fine-tuning setup, training schedule, and evaluation across five fixed seeds.

A Detailed Experimental Setup

  • The experimental design follows the LLM-JEPA protocol to isolate episodic replay without changing data, model scale, or training configuration.It compares standard supervised fine-tuning, LLM-JEPA, and ER-JEPA objectives.Source: A Detailed Experimental Setup · p. 11
  • All three objectives use LoRA rank 512, learning rate 1e-5, and six training epochs for each model–dataset combination.The study says this LoRA setup was chosen because the LLM-JEPA study found it outperformed full-parameter fine-tuning under the same setting.Source: A Detailed Experimental Setup · p. 11
  • Experiments use five fixed random seeds, report mean performance across seeds, and assess significance with paired one-tailed t-tests over seed-matched runs.Training runs use a server equipped with four NVIDIA A800 GPUs.Source: A Detailed Experimental Setup · p. 11

B How Historical Replay Can Provide Additional Correction Directions

The analysis gives a conditional local account of how replay can improve the selected historical loss while preserving current-batch token probabilities and alignment residuals to first order. It also describes how regex structural subsets are assigned for evaluation.

B How Historical Replay Can Provide Additional Correction Directions

  • Replay directions preserve current token probabilities and alignment residuals to first order, though they may still change losses on historical examples.The projector K may contain only the zero vector.Source: B How Historical Replay Can Provide… · p. 12
  • For fixed β > 0, replay adds the same joint objective on a nonempty replay multiset; overlapping examples are counted separately and representations are recomputed.The comparison holds the selected replay multiset fixed.Source: B How Historical Replay Can Provide… · p. 12
  • Current-batch gradients, including repetitions or fixed reweightings of current examples, have no component in the defined subspace K, while replay can supply a descent direction for a selected historical loss.The analysis holds at a fixed parameter state and for a fixed replay multiset.Source: B How Historical Replay Can Provide… · pp. 11–12
  • If the replay gradient is nonzero, the replay loss is lower with replay than without it for sufficiently small positive step sizes, even when its projection onto K is zero.The comparison is stated in the limit as η approaches zero.Source: B How Historical Replay Can Provide… · p. 12
  • The analysis establishes a conditional local effect on replay loss; predictive-performance implications are evaluated empirically.The projector is an analytical device, while ER-JEPA uses the full replay gradient.Source: B How Historical Replay Can Provide… · p. 12

C.1 Regex Structural Subsets

  • Regex examples are assigned to every applicable structural subset; escaped literals and symbols inside character classes are not counted as separate operators.The subsets may overlap.Source: §C.1 · p. 13

C.2 Error Correction Across Seeds

On SYNTH, ER-JEPA has the highest reported mean accuracy among the three methods, and replay is associated with fewer local accuracy declines across individual seeds. The section also describes matched-input error analyses and example predictions.

C.2 Error Correction Across Seeds

  • For each seed, the error analysis evaluates both methods on the same 2,000 SYNTH inputs where Baseline erred, excluding syntax errors from length-based categories.A correction requires an exact target match, and correction rates use the Baseline error count in each row as the denominator.Source: §C.2 · p. 13

C.3 Examples Selected by a Fixed Rule

  • 83.65% ER-JEPA, 68.75% LLM-JEPA, and 53.63% Baseline mean accuracy on 2,000 SYNTH test inputs per seed across five seeds.Figure 11 illustrates selected Baseline errors corrected by ER-JEPA.Source: §C.3 · p. 13, Figure 11 · p. 14

C.4 Individual Seed Accuracy at Matched PFLOPs

  • Baseline accuracy declines locally in four of five seeds and LLM-JEPA declines in all five; such declines occur less often with replay.The authors suggest replay mitigates overfitting across individual runs; Figure 13 shows raw accuracy at six shared target PFLOPs for four seeds, without means or error bars.Source: §C.4 · p. 13, Figure 13 · p. 15

D ER-JEPA Implementation Details

ER-JEPA uses a fixed-capacity episodic memory to store and retrieve past training examples for replay. Replay policies select examples by content similarity, uniform sampling, or significance, while sharing the memory rules and replay loss.

D.1 Episodic Memory and Retrieval

  • Each active memory slot stores a source–target token pair, a sparse address, and management metadata; current mini-batch examples are excluded from replay.Content-based retrieval uses the current source representation as a cue, retrieves up to κ slots by cosine similarity, then deduplicates and truncates the batch’s selections to replay budget R.Source: §D.1 · p. 16

D.2 Shared Episodic Memory

  • When memory is full, ER-JEPA evicts the slot with the lowest replay-adjusted significance and updates selected slots’ significance after replay.Unselected slots retain their scores, and these memory-management rules are shared across the three selection policies.Source: §D.2 · p. 17

D.3 Replay Selection Policies

  • Content, uniform, and hard replay differ only in slot selection; all share the memory rules and replay loss, and return no selection when memory is empty.The three policies are content-based retrieval, equal-probability sampling, and significance-based selection.Source: §D.3 · p. 17

D.3.1 Content Replay

  • Content replay maps source representations to sparse addresses using a frozen random projection and Top-K selection, then retrieves slots by cosine similarity.The projection remains fixed throughout training; retrieved slots are ordered by similarity, deduplicated, and truncated to the replay budget.Source: §D.3.1 · pp. 17–18

D.3.2 Uniform Replay

  • Uniform replay samples stored examples with equal probability, without using cue similarity or significance, and estimates mean JEPA error over active memory.When the replay budget is no greater than the number of active slots, sampling is without replacement; when the memory is smaller than the budget, sampling is with replacement.Source: §D.3.2 · p. 18

D.3.3 Hard Replay

  • Hard replay selects up to R slots with the largest significance scores, but a large JEPA error indicates poor representation alignment, not necessarily an incorrect generated answer.Selected slots are returned in descending score order, with shorter selections filled to the replay budget.Source: §D.3.3 · p. 18

E Ablation Studies

The sparse-address and dense-representation variants achieve similar mean accuracy, with sparse retrieval higher by 0.16 percentage points.

E.1 Sparse and Dense Memory Addressing

  • Sparse addresses are formed by a frozen random projection and Top-K selection.The comparison tests sparse addresses against dense source representations for content replay.Source: §E.1 · p. 18
  • 86.37% for sparse retrieval versus 86.21% for dense retrieval, a 0.16 percentage-point difference in mean accuracy.Both variants use cosine similarity, with all other training settings fixed.Source: §E.1 · p. 18
  • The two addressing strategies achieve similar mean accuracy across the reported five-seed results.The reported statistics include mean accuracy, standard deviation, minimum, and maximum.Source: §E.1 · p. 18

E.2 Hyperparameter Ablations

Increasing episodic-memory capacity from 10 to 104 raises accuracy after four epochs. The ablations also vary JEPA loss weight λ and predictor depth k, while the capacity runs are not compute matched.

E.2 Hyperparameter Ablations

  • 84.98% versus 84.22% accuracy after four epochs as memory capacity increases from 10 to 104.The capacity ablation reports mean ± standard deviation over five seeds for Meta-Llama-3.2-1B-Instruct on NL-RX-SYNTH.Source: §E.2 · p. 18, Table 6 · p. 19
  • The objective ablation varies JEPA loss weight λ and predictor depth k, and Figure 15 reports accuracy for each setting.The figure marks the best result in bold.Source: §E.2 · p. 18, Figure 15 · p. 20
  • The memory-capacity learning curves show accuracy across four epochs and four capacities, but the runs are not compute matched.Figure 14 reports means over five seeds with one sample standard deviation error bars; the observed training costs are also not compute matched.Source: Table 6 · p. 19, Figure 14 · p. 20
Loading 2609.36952v1…