Source-linked AI summary
Don't Forget! Decomposing the Training Dynamics of Memorization in Language Models
Florian Eichin, Philipp Mondorf, Andrei Mircea, Yupei Du, Barbara Plank, Michael A. Hedderich
TL;DR
Memorization can help language models fit long-tail training distributions, but its development during training remains poorly understood. The paper decomposes memorized-sequence loss trajectories across training and parameters in Pythia, studying both frequent recited and rare recollected sequences. It finds distinct alignment patterns associated with memorization and forgetting, and shows the decomposition can support earlier prediction and post-hoc intervention.
Problem
How and where memorization develops during training remains poorly understood, limiting understanding of its training dynamics.
Source: §1 · p. 1Method
The authors extend ExPLAIND to decompose memorized-sequence loss trajectories across training, parameters, and data, analyzing recited and recollected sequences in Pythia.
Source: §1 · p. 1Results
Memorized sequences align strongly with updates from their own occurrences, especially in lower layers; recited sequences show stronger negative alignment with other training influences, while decomposition features improve prediction over a cross-entropy baseline.
Source: §1 · pp. 1–2Takeaways & Limitations
The decomposition supports early prediction, and post-hoc intervention on a small set of influential parameters ablates memorization in the final model.
Source: §1 · p. 2The methodology omits AdamW first-moment influences and estimates the second moment, choices that could bias results; prediction is a proof of concept rather than a full-dataset competitor.
Source: §5 · p. 9
Abstract
from arXiv · showhide
Memorization has been proposed as a mechanism to explain how language models fit the tail of their training distributions, but its training dynamics are not understood well. In this work, we take a fine-grained look at memorization by decomposing the loss trajectory of memorized sequences over training and model parameters. Across the Pythia family, we study memorization of duplicated training sequences (recitation) and rare ones (recollection). We find that memorization in both cases is characterized by sequence-level gradient alignment, though recitation suffers from misalignment with other training influences which causes forgetting, explaining the necessity for higher duplication of these examples. We further show that the lower model layers are the most involved in memorization and forgetting. Predicting memorization, our decomposition improves over a cross-entropy baseline, especially in larger models and early in training. Intervening on a small set of highly influential parameters we are able to ablate memorization in the final model. Together, these findings advance our understanding of how memorization develops during training and offer insights for predicting and intervening on it.
1 INTRODUCTION
The paper asks how memorization develops across training, parameters, and data, and whether understanding these contributions helps predict and intervene on memorization. Its decomposition links memorization to alignment with a sequence’s own training occurrences and to how other training influences affect it.
1 INTRODUCTION
- The study addresses limited understanding of how and where memorization develops during training, a gap that restricts prediction and intervention.The paper notes privacy and data security as applications where intervention is desirable.Source: §1 · p. 1
- The authors extend ExPLAIND to analyze loss trajectories across training, parameters, and data for recitation and recollection in Pythia.Recitation concerns highly frequent sequences; recollection concerns rare sequences.Source: §1 · p. 1
- Memorized sequences align strongly with updates from their own occurrences, particularly in lower model layers.The decomposition also considers how remaining training influences interact with memorization.Source: §1 · pp. 1–2
- Highly duplicated recited sequences show stronger negative alignment with other training data and weight decay, while recollected sequences show less negative alignment.The authors suggest repeated exposure helps recited sequences avoid forgetting; the lower negative alignment is consistent with recollection despite few occurrences.Source: §1 · p. 2
- Decomposition-derived features improve memorization prediction over a cross-entropy baseline, especially in larger models and early in training.Post-hoc intervention on a small set of influential parameters, mostly in lower attention and MLP layers, ablates memorization in the final model.Source: §1 · p. 2
2 RELATED WORK
Prior work studies memorization in relation to training-distribution tails, corpus statistics, model scale, and model components. This study extends that work with earlier prediction based on loss-decomposition features and a unified decomposition validated through prediction and intervention.
2 RELATED WORK
- Prior Pythia studies find memorization predictable during pretraining and scaling with model size, while this work predicts it in earlier phases using loss-decomposition features.Prashanth et al. (2025) study corpus statistics and a duplication-based taxonomy; Tirumala et al. (2022) report larger models memorize faster and forget less, with the latter also found and explained here.Source: §2 · p. 2
- Earlier localization work implicates MLP layers and attention heads in recall, which the authors say their results confirm.Other studies identify embeddings as important and describe memorization and generalization as entangled; this work also indicates their interdependence in early pretraining.Source: §2 · p. 2
- The study extends Eichin et al. (2026)’s decomposition with explicit alignment between training influences and predictions, then validates findings through prediction and causal intervention.The passage characterizes this as a unified and exact approach combining attribution across data, parameters, and training.Source: §2 · p. 2
3 METHODOLOGY
The methodology decomposes loss changes into interpretable training-example and parameter influences, then analyzes how those influences shape memorized-sequence predictions over training. Experiments use sampled Pythia sequences classified by extractability and training-set frequency, with simulated updates adapted to available checkpoints.
3 METHODOLOGY
- The authors extend ExPLAIND to decompose loss trajectories into per-sample and per-parameter contributions, further separating sensitivity, update magnitude, and alignment.The decomposition is applied to evaluation loss across training.Source: §3 · p. 2
- Directional alignment indicates whether an update decreases or increases prediction loss; near-zero alignment means the update does not affect that prediction.Influence also depends on test sensitivity and update magnitude, with update magnitude scaled relative to recent gradient history.Source: §3 · p. 3
- The analysis tracks how memorized-sequence occurrences and the rest of training work toward or against memorization solutions.Influences are further decomposed across parameter groups, including model layers, to localize learning contributions.Source: §3 · p. 3
- Experiments use Pythia models from 160m to 6.9B, focusing on Pythia-1B while running all experiments across the family unless stated otherwise.Pythia provides checkpoints and training data, as well as an exhaustive search of duplicated and memorized sequences.Source: §3 · p. 3
- Sequences are classified as recited if they occur more than five times and recollected if they occur up to five times, using 32-extractability as the memorization criterion.The study filters repeated-subsequence and code samples and includes non-extractable control groups matched by occurrence frequency; it samples 32 sequences per class.Source: §3 · p. 4
- For Pythia, the authors simulate updates at available checkpoints, omit the first-moment term, estimate the second moment, and inject sampled sequences into a training batch.They report that 10–20% of sampled sequences could not be verified as 32-extractable, raising a robustness question about requiring strict 32-extractability.Source: §3 · p. 4
4 RESULTS
Memorization emerges early, but recited sequences face stronger interference from other training influences than recollected ones. Decomposition features help predict memorization, localize it especially to lower layers, and identify parameters whose ablation suppresses it.
4 RESULTS
- Recited and recollected loss trajectories separate from controls by 256 and 1,000 steps, respectively; by 10k, the model greedily decodes over 80% of memorized sequences.Token-level extractability reveals early learning that strict 32-extractability lags, with a large fraction of memorization learning occurring before step 20k.Source: §4 · pp. 4–5
4.1 DECOMPOSITION OF MEMORIZED SEQUENCES
- Memorized sequences align more strongly with updates from their own occurrences, while recited examples experience increasingly negative alignment with other-data updates and recollected examples cluster near zero.Recited instances show self-alignment around 0.3–0.4, compared with about 0.25 for recollected instances and about 0.18 for controls; the authors interpret the contrast as repeated exposure helping resist overwriting.Source: §4.1 · pp. 5–6
4.2 MEMORIZATION PREDICTION
- After step 256, CE+decomposition achieves 90.6% accuracy at step 10k, versus 84.4% for cross-entropy only and 75.0% for decomposition only, then peaks at 96.1% by step 40k.Source: §4.2 · p. 6
- For Pythia-6.9B at step 10k, CE+decomposition reaches 88.3% accuracy versus 77.3% for CE only, and CE alone does not catch up until the final checkpoint.Across model scales, cross-entropy becomes a weaker early predictor as model size increases; adding decomposition features recovers much of that early predictability.Source: §4.2 · p. 7
4.3 LOCALIZING MEMORIZATION IN PYTHIA
- Memorized examples are most sensitive to model changes in lower layers, where self-alignment is also highest; memorization self-aligns most in MLP layers.Input embeddings receive the largest training updates, while attention QKV alignment is also higher for memorized examples than for controls.Source: Figure 4 · p. 7, Figure 5 · p. 7, §4.3 · p. 8
- In lower layers 0–3, other-data and weight-decay updates align more strongly with recited memories than recollected ones, consistent with greater ongoing forgetting of recited examples.The authors connect these degrading influences and high test sensitivity in lower layers to recited sequences' need for continued restoration.Source: Figure 6 · p. 7, §4.3 · p. 8
4.4 CAUSAL MEMORIZATION ABLATIONS
- At step 2,000, ablating the top-0.0002 parameters drives Pythia-1B 32-extractability toward zero, while smaller interventions keep changes in control-sample cross-entropy below 0.1.The ablation is restricted to internal parameters because embedding-ranked interventions disrupt control behavior; the authors also report that one attention head can ablate memorization in models up to 1B parameters.Source: Figure 7 · p. 8, §4.4 · pp. 8–9
- Embedding-ranked ablations disrupt control behavior and appear to target training artifacts, while intervening on a single attention head can ablate memorization in models up to 1B parameters.The authors therefore restrict their main parameter interventions to internal parameters.Source: §4.4 · pp. 8–9
5 DISCUSSION & CONCLUSION
The discussion frames memorization as an early-forming process whose later dynamics mostly refine existing memories, while some forgetting influences weaken with model scale. It also identifies methodological constraints and limits the findings’ scope to verbatim memorization in Pythia models up to 6.9B parameters.
5 DISCUSSION & CONCLUSION
- Most memorization mechanisms form early: losses flatten around 5k steps, influences stabilize around 20k, and later dynamics largely refine existing memories.The discussion also places memorization and generalization in parallel, with phase transitions around 5k and 20k steps.Source: §5 · p. 9
- As model scale increases, constant forgetting effects from training and weight decay in lower model regions weaken; weight decay is slightly aligned with memorization in the largest models.The authors suggest these patterns could explain previous observations that memorization scales with model size.Source: §5 · p. 9
- The methodology omits AdamW’s first-moment influence and estimates its second moment, choices that could bias the results.The authors also describe their memorization prediction as proof-of-concept rather than a full-dataset competitor; computing ExPLAIND features over the entire dataset is infeasible.Source: §5 · p. 9
- Post-hoc interventions validate the analysis but do not establish how to prevent memorization during training, and their effects on robustness and other predictions remain unclear.The discussion limits the findings to verbatim memorization in Pythia models up to 6.9B parameters; generalization beyond that scope remains open.Source: §5 · pp. 9–10
A.1 PROOF OF EQUATION 1
The proof specializes ExPLAIND’s AdamW decomposition to sequence loss by applying its output-function result and accounting for training batches and optimizer details.
A.1 PROOF OF EQUATION 1
- The AdamW theorem decomposes loss using fixed per-sample weights and differentiable model and loss functions along the parameter segments.The setup assumes AdamW steps with positive learning rates and specified moment parameters.Source: §A.1 · p. 14
- The output-function corollary lets the decomposition apply to sequence loss through the model’s stacked next-token outputs and a fixed normalizer N.For a training batch, the sequence losses are summed with weight 1 per sequence.Source: §A.1 · pp. 14–15
- The specialization accommodates batch averaging and norm clipping by scaling each sequence contribution to match the gradient supplied to AdamW.Under clipping, the realized clipping factor is treated as fixed when attributing the update.Source: §A.1 · p. 15
A.2 ESTIMATING THE SECOND-MOMENT STATE OF PYTHIA
Because Pythia checkpoints lack optimizer state, the authors estimate AdamW’s second moment from held-out gradients, correcting for the difference between sample and training batch sizes.
A.2 ESTIMATING THE SECOND-MOMENT STATE OF PYTHIA
- The reconstruction treats checkpoint second moments as expected squared per-batch gradients over an effective 20-step EMA horizon, assuming approximate stationarity.Pythia’s training batch size is B = 1024.Source: §A.2 · pp. 15–16
- A smaller sample batch overestimates the target squared batch gradient by (1/b − 1/B)σ2, potentially by a factor of B/b in the noise-dominated regime.The bias does not vanish with more samples.Source: §A.2 · p. 16
- The estimator uses N held-out per-sequence gradients to form a sample mean and variance for estimating the expected squared gradient at the training batch size.The held-out data are disjoint from the batches being decomposed.Source: §A.2 · pp. 15–16
- With N = 256, relative update error is roughly 4% under Gaussian noise and remains below 10% for gradients with κ ≤13.The true 20-step EMA has relative fluctuations of roughly 30%.Source: §A.2 · p. 16
B ADDITIONAL EXPERIMENT SETTINGS
The experiments use deduplicated Pythia GPT-NeoX checkpoints and reconstruct AdamW while adapting compute and precision to model size.
B ADDITIONAL EXPERIMENT SETTINGS
- The study uses deduplicated Pythia GPT-NeoX checkpoints and does not decompose the 12B model because of computational cost.The 6.9B model uses bfloat16 to reduce memory; other computation defaults to float32.Source: §B.1 · p. 16
- AdamW reconstruction uses Pythia’s learning-rate schedule, betas (0.9, 0.95), epsilon 10^-8, weight decay 0.01, gradient clipping, and zero initial first moments.Each decomposition is restricted to one H200 GPU.Source: §B.1 · p. 16
- The path integral in ExPLAIND’s test feature map uses 100 integration steps, following Eichin et al. (2026)’s recommendation.Source: §B.1 · p. 16
B.2 DATA
The appendix specifies the sampled memorization classes, decomposition checkpoints, and procedures used for prediction and evaluation.
B.2 DATA
- The data comprise 32 examples from each of four classes defined by extractability and continuation count: recited, recollected, duplicated controls, and controls.Each example uses the first 64 tokenizer tokens; 128 prediction examples pass a text filter.Source: §B.2 · pp. 16–17
B.3 EXPLAIND DECOMPOSITION AND FEATURES
- The decomposition evaluates checkpoints from steps 0 through 130,000, retains sequences as separate accumulated groups, and defines sequence loss as token loss sum divided by 2,048.The schedule samples early checkpoints densely and later checkpoints at 1,000- or 10,000-step intervals.Source: §B.3 · p. 17
B.4 MEMORIZATION PREDICTION
- Prediction uses five checkpoint-specific decomposition features, including alignment and feature norms, while BLiMP evaluation compares summed sentence log probabilities in minimal pairs.The features include alignment with the original training update, matching sequence update, and weight decay contribution.Source: §B.4 · p. 17
- Ridge linear models compare decomposition features, cross-entropy, and their combination using stratified six-fold cross-validation and held-out accuracy.Feature standardization is fitted only on each fold’s training portion.Source: §B.4 · p. 17
B.5 CAUSAL ABLATION
The ablation procedure zeros selected parameters based on accumulated influence scores, then measures memorization and control-sample effects.
B.5 CAUSAL ABLATION
- At each checkpoint, the procedure sums parameter influences from 64 memorized sequences across that and all earlier checkpoints, then zeros the most negative fractions.Source: §B.5 · p. 17
- Ablations exclude normalization parameters and are evaluated with embeddings included or excluded, using fractions relative to eligible parameters in each scope.Source: §B.5 · p. 17
- Random masks with the same parameter count serve as a baseline; outcomes include 32-token extractability, token loss, control cross-entropy, and changed control tokens.Source: §B.5 · p. 17
C.1 TRAINING DYNAMICS AND GRADIENT NORMS ACROSS SCALES
Across Pythia scales, the section examines memorization dynamics, prediction, parameter interventions, and a shared mechanism in lower-layer attention heads. The decomposition features improve memorization prediction, while targeted interventions and alignment analyses probe where memorization develops and how it interacts with training.
C.1 TRAINING DYNAMICS AND GRADIENT NORMS ACROSS SCALES
- Memorized sequences show strong self-alignment; recited sequences have more negative alignment with other training data and weight decay than recollected sequences, with lower layers prominent in self-alignment.Source: §C.1 · p. 17
C.2 PREDICTION
- Decomposition features improve held-out memorization prediction over the cross-entropy baseline, particularly in larger models and early in training.Source: §C.2 · p. 18
C.3 CAUSAL PARAMETER INTERVENTIONS
- Ablations examine 32-extractability and control loss across scales, and locate selected parameters by depth and layer type with embeddings included or excluded.Source: §C.3 · p. 18
C.4 A SHARED MEMORIZATION MECHANISM IN LOWER LAYER ATTENTION HEADS
- Recited instances remain highly aligned with each other throughout training, while other instances follow this pattern only early; until about step 20k, lower-layer attention heads largely drive alignment.Source: §C.4 · p. 18
- In smaller models, ablating attention head 0 in layer 0 reduces 32-extractable sequences to 5 recited and 1 recollected, while control cross-entropy rises about 0.25.Source: §C.4 · p. 18
- The attention-head result generalizes to models around 1B, but the naive search and intervention did not reproduce it at 1.4B and larger.Source: §C.4 · p. 18