Research map / Large language models / Language, safety and interpretability
Memorization Sinks: Isolating Memorization during LLM Training
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Memorization Sinks (sequence-ID-dependent dropout of a per-sequence memorization neuron pool in MLPs) to isolate memorization of repeated natural text so it can be removed post-hoc without hurting generalization.
Problem
Post-hoc removal of memorized repeated natural text from LLMs degrades general capabilities because memorization is mechanistically entangled with generalizing components.
Approach
Allocate a pool of MLP neurons as memorization sinks that are activated deterministically per sequence ID (sequence-tied dropout) during training, so memorization accumulates in sink neurons shielded from interference; the sinks are removed at inference/unlearning time.
Key results
- On the larger-scale SlimPajama+TinyStories setup, memorization sinks close at least 50% of the train-validation loss gap, while validation loss matches the repeated-data baseline and beats deduplication.
- Dropping memorization neurons raises loss on repeated TinyStories sequences to roughly 66% of a standard model that does not memorize, while validation loss is comparable to standard training.
- Sequence-ID noise up to 10% is tolerated; at 50% noise isolation of memorization fails.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use sequence-tied dropout and remove memorization sink neurons after training | Memorization Sinks (sequence-tied dropout) · TinyStories validation set | validation loss | No clear effect |
| Use sequence-tied dropout and remove memorization sink neurons after training | Memorization Sinks (sequence-tied dropout) · TinyStories repeated training sequences | loss on repeated sequences | Helped |
| Use sequence-tied dropout instead of deduplicating repeated TinyStories data | Memorization Sinks (sequence-tied dropout) · SlimPajama and TinyStories pretraining; TinyStories validation set | validation loss | Helped |
| Use sequence-tied dropout instead of standard training with repeated TinyStories data | Memorization Sinks (sequence-tied dropout) · SlimPajama and TinyStories pretraining | training-versus-validation loss gap: by at least 50% | Helped |
| Restrict gradients from repeated sequences to memorization neurons and from other sequences to generalization neurons | Gradient Masking · TinyStories validation set | validation loss | Hurt |
| Drop out memorization neurons after gradient-masked training | Gradient Masking · TinyStories validation set | validation loss | Hurt |
| Apply post-hoc Integrated Gradients localization rather than pruning | Post-hoc localization methods · Repeated natural TinyStories sequences | sequence forgetting | Hurt |
| Apply post-hoc localization to repeated natural sequences rather than atypical canaries | Post-hoc localization methods · TinyStories repeated sequences and TS-Canary | model degradation | Hurt |
| Increase the fraction of memorization neurons activated per sequence to p=0.7 instead of p=0.3 (ablation) | Memorization Sinks (sequence-tied dropout) · TinyStories repeated training sequences | loss on memorization examples: p = 0.3; p=0.7 | Hurt |
| Use inconsistent sequence IDs for 10% of repeated documents (ablation) | Memorization Sinks (sequence-tied dropout) · TinyStories repeated training sequences | isolation of memorization: up to 10% | No clear effect |
| Use inconsistent sequence IDs for 50% of repeated documents (ablation) | Memorization Sinks (sequence-tied dropout) · TinyStories repeated training sequences | isolation of memorization: 50% noise | Hurt |
| Increase repetition to 20x (ablation) | Memorization Sinks (sequence-tied dropout) · SlimPajama and TinyStories pretraining | memorization mitigation: 20x | Helped |
Surprising findings
- Forced gradient-masking localization significantly worsens generalization compared to standard training, and dropping memorization neurons still degrades validation loss.
- For natural repeated text, post-hoc pruning and integrated gradients cause greater model degradation than for canary sequences.
Limitations
- Experiments are mostly small-scale and controlled; larger-scale effects on capabilities are unverified.
- Requires sequence-ID metadata grouping duplicated examples; robustness to noisy IDs only up to 10%.
- Robustness to adversarial extraction techniques not verified.
Experimental setup
Datasets: TinyStories SlimPajama SmolLM training mixture
Benchmarks: TinyStories validation loss SlimPajama+TinyStories validation loss
Metrics: validation loss sequence forgetting (loss increase on repeated sequences aft model degradation (validation loss increase after removal)
Baselines: Standard training Standard: No Rep (deduplication) Gradient Masking Hard Concrete pruning Integrated Gradients
Builds on: gradient routing Hard Concrete pruning integrated gradients neuron localization nanoGPT SmolLM
Closest papers by meaning
- Mitigating Memorization in Language Models (ICLR 2025 · similarity 0.91)
- Mitigating Memorization Where It Happens (NeurIPS 2026 · similarity 0.89)
- Be like a Goldfish, Don't Memorize! Mitigating Memorization in Generative LLMs (NeurIPS 2024 · similarity 0.89)
- Rethinking LLM Memorization through the Lens of Adversarial Compression (NeurIPS 2024 · similarity 0.88)
- Learning the Signature of Memorization in Autoregressive Language Models (NeurIPS 2026 · similarity 0.88)
- To Each (Textual Sequence) Its Own: Improving Memorized-Data Unlearning in Large Language Models (ICML 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.