Replay memory improves prediction and alignment in language models
ER-JEPA: Experience Replay Improves Joint-Embedding Predictive Learning in Language Models
Computation and LanguageArtificial Intelligence
Summary
Large language models are good at predicting words but sometimes learn meanings that aren't quite right or miss important context. The authors worked on a method called LLM-JEPA that tries to teach these models by comparing different views of the same knowledge. They improved this by adding a memory system called ER-JEPA that stores and reuses past training examples to give the model more guidance. This extra memory helps the model make better predictions and understand language more accurately across multiple tests.
What this means in practice
- •For language model engineers: Improve language model training to produce more accurate token predictions and better semantic understanding by incorporating experience replay techniques.
- •For natural language processing teams: Enhance downstream language tasks by using models with replay-enhanced predictive learning for more stable and aligned representations.
Authors
Jingnan Pu, Zi-En Fan, Feng Lian
Abstract
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.