[ICLR 2025 Submission] MSSR: Harmonizing Human Memory Logic with LLM Continual Fine-Tuning
MSSR: Memory-Aware Adaptive Replay for Continual LLM Fine-Tuning
MSSR (Memory-Aware Adaptive Replay) is a continual learning framework for LLMs that integrates cognitive memory theories, specifically the Ebbinghaus forgetting curve, to manage experience replay. It achieves state-of-the-art results across 11 sequential tasks and multiple backbone models (Qwen2.5, Llama-3.1, Gemma-2), significantly reducing catastrophic forgetting in reasoning-intensive benchmarks.
TL;DR
Catastrophic forgetting remains the "Achille's heel" of Large Language Models (LLMs) when fine-tuned on a sequence of tasks. MSSR (Memory-Aware Adaptive Replay) bridges cognitive science and deep learning by adopting the Ebbinghaus forgetting curve to schedule data replay. By modeling per-sample memory strength and utilizing an expanding-interval "spaced repetition" schedule, MSSR achieves SOTA retention on reasoning tasks with negligible computational cost.
The Core Conflict: Adaptation vs. Retention
When we fine-tune an LLM on Task A and then move to Task B, the gradient updates for the new task often overwrite the weight configurations supporting Task A. This is Representation Drift.
Current solutions like "Fixed Replay" are too "dumb"—they replay old data at constant intervals regardless of whether the model has already mastered it. "Accuracy-based" methods are "smarter" but expensive, requiring constant evaluation on validation sets to detect when a model "forgets." MSSR proposes a third way: Predictive Modeling of Forgetting.
Methodology: The Ebbinghaus Logic
The authors propose that LLM forgetting isn't uniform. Like human memory, a model's "stability" on a specific sample increases every time it is successfully re-encountered (replayed).
1. Sample-Level Dynamics
MSSR tracks a Memory Strength () and Stability () for every sample.
- Decay Phase: Between replays, memory strength decays at a rate determined by the sample's difficulty (loss) and its current stability.
- Consolidation Phase: When a sample is replayed, its memory strength is reset to 1, and its stability is boosted. Crucially, the boost is higher if the interval since the last replay was longer (the Spacing Effect).
2. The Adaptive Scheduler
Instead of replaying at every steps, MSSR uses:
- Expanding Intervals: Intervals between replays grow (e.g., 1, 2, 4, 7... steps) as the model stabilizes.
- Decaying Replay Ratio: The proportion of old data in a batch starts high and tapers off as training progresses.
Figure 1: The closed-loop workflow of MSSR, showing the interplay between memory tracking and adaptive scheduling.
Experimental Battleground
The researchers tested MSSR against 11 sequential tasks, including high-stakes reasoning benchmarks like GSM8K (math) and MMLU (general knowledge).
Key Results:
- Mistral & Qwen Performance: MSSR consistently outperformed "Fixed" and "Loss-based" replays. In the 3-task setup (MMLU -> GSM8K -> MATH), MSSR reached an average score of 54.5 on Qwen2.5-7B, compared to 52.1 for vanilla fine-tuning.
- The Long Horizon: In the 11-task setting, MSSR’s ability to protect the "first" tasks in the sequence was significantly better than baselines.
Table 2: Performance across 11 tasks. Note how MSSRfull dominates the early-task metrics (AGNews, SQuAD) even after 10 subsequent training stages.
Efficiency Analysis
One might fear that tracking status for thousands of samples is heavy. However, MSSR uses scalar-only updates and lazy state updates.
- Wall-clock time: +5% overhead.
- Peak Memory: +6% overhead. Compared to the 2x or 3x training time required for "Accuracy-based" triggers (which necessitate frequent full evaluations), MSSR is remarkably lean.
Critical Insight: Why it Works
The "magic" of MSSR lies in its Inductive Bias. By assuming the model follows a decay curve, it forces the optimizer to focus on "fragile" knowledge during the early stages of a task transition and "anchor" that knowledge through spaced repetition. This aligns with the Stability-Plasticity Dilemma: it grants the model plasticity for new tasks while enforcing stability on old ones only when the "predictive risk" of forgetting is high.
Conclusion & Future Look
MSSR proves that we don't need complex reinforcement learning or constant evaluations to solve forgetting. A simple, cognitively-grounded mathematical model of memory decay is enough to steer LLM fine-tuning toward long-term stability. Future work could potentially expand this to Online Learning, where the model learns from a continuous stream of unlabelled web data.
Senior Editor's Note: MSSR is a refreshing take on CL. While most "SOTA-chasing" papers add complexity (distillation, GAN-based generators), MSSR looks back at 19th-century psychology (Ebbinghaus) to simplify the 21st-century problem of LLM drift.
