Research map / Label-efficient and robust learning / Knowledge distillation

Memory Grafting: Scaling Language Model Pre-training via Offline Conditional Memory

NeurIPS 2026 · Label-efficient and robust learning › Knowledge distillation · approach: Few-shot and metric distillation · 0 citations (OpenAlex)

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 Memory Grafting, which fills a conditional n-gram memory table with offline frozen hidden states from a large grafting LLM, retrieved by exact longest-match lookup with a hash-based Engram fallback, and reports gains over MoE and vanilla Engram baselines at 0.92B and 2.8B scales.

Problem

Scaling conditional memory (e.g., Engram) requires learning large n-gram memory tables from scratch during pre-training, which is expensive and sometimes ineffective for compact models.

Approach

Instead of learning memory values from scratch, precompute frozen final-token hidden states of frequent n-grams from a pretrained grafting model, retrieve them via exact longest-match suffix lookup, adapt them with lightweight projections and gates, and fall back to trainable hashed Engram memory on misses.

Key results

• At 2.8B total / 0.55B activated with 100B tokens, Memory Grafting averages 53.86 vs 52.43 for vanilla Engram and 51.95 for MoE.
• At 0.92B with 50B tokens, MG-Qwen3.5 averages 46.98 vs 45.62 for MoE and 45.03 for vanilla Engram.
• Memory Grafting reaches 0.92x training throughput and 0.95x inference throughput relative to the MoE baseline; a 3M-entry table takes about 3 A100 GPU-hours to build from Qwen3.5-35B-A3B.

Surprising findings

• Vanilla Engram underperforms the MoE baseline at the 0.92B scale (45.03 vs 45.62), attributed to hash collisions from a limited embedding table.
• GLM-4.7-Flash grafted memory gives only a marginal gain (45.83), and its memory geometry is strongly anisotropic with low effective rank.

Limitations

• No comparison with knowledge distillation baseline due to compute cost.
• Gains depend on grafting model quality and source layer choice; deeper layers can hurt.
• Results limited to two recipient scales and a single benchmark suite evaluated with lm-evaluation-harness; memory bank requires ~24GB of frozen storage.

Experimental setup

Datasets: Nemotron-CC

Benchmarks: ARC-Challenge ARC-Easy BoolQ Social IQA RACE LAMBADA WinoGrande PIQA HellaSwag

Metrics: accuracy average accuracy training loss CKA throughput memory usage

Baselines: MoE baseline Engram MoE Product-Key Memory UltraMem STEM MoLE

Builds on: Engram DeepSeekMoE Mixture-of-Experts

Closest papers by meaning

• Language Model Memory and Memory Models for Language (NeurIPS 2026 · similarity 0.90)
• Frozen Memory Is Not Enough: Rethinking External Memory as Extraction (NeurIPS 2026 · similarity 0.89)
• Steering Information Utility in Key-Value Memory for Language Model Post-Training (NeurIPS 2025 · similarity 0.89)
• TAID: Temporally Adaptive Interpolated Distillation for Efficient Knowledge Transfer in Language Models (ICLR 2025 · similarity 0.89)
• Mitigating Memorization in Language Models (ICLR 2025 · similarity 0.88)
• Memorization Sinks: Isolating Memorization during LLM Training (ICML 2025 · similarity 0.88)