Research map / Large language models / Language, safety and interpretability
Chunk-Distilled Language Modeling
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 CD-LM, a training-free decoding method that interleaves LM token generation with retrieval of multi-token chunks from a trie datastore keyed by LM hidden states, improving perplexity, efficiency, and knowledge injection.
Problem
Token-by-token LLM generation is inefficient, and adapting pretrained LMs to new domains or knowledge normally requires costly retraining or RAG that adds overhead without efficiency gains.
Approach
Store high-probability multi-token chunks with their preceding contexts in entry-token tries, retrieve chunks by cosine similarity of the LM's own hidden states, and accept them stochastically via a piecewise-linear mapping so that multiple tokens are emitted in one decoding step; a dynamic program gives exact sequence probabilities for perplexity.
Key results
- KCD-LM with GPT-2 small reduces WikiText-103 test perplexity from 34.83 to 22.90, versus 32.19 for kNN-LM, and reaches 24.95 on Medical vs 39.66 for kNN-LM at its best setting.
- SCD-LM with GPT-2-XL-conversational reduces mean token time by 19.59% and saves 43.33% of forward passes on MT-Bench-80.
- On private information injection, ECD-LM with GPT-2-XL reaches 75.7% PII accuracy vs 0% for the base LM and 46.4% with in-context learning.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Inject retrieved chunks from a domain-specific teacher | GPT-2 137M with KCD-LM · WikiText-103, Medical, Federal Register law, and Dockerfile code test sets | perplexity: WikiText: 34.83 -> 22.90; Medical: 51.68 -> 24.95; Law: 11.41 -> 8.24; Code: 106.44 -> 50.77 | Helped |
| Inject retrieved chunks from a domain-specific teacher | GPT-2 137M with KCD-LM · WikiText, Medical, Law, and Code test generation | MAUVE: WikiText: 0.016 -> 0.032; Medical: 0.006 -> 0.011; Law: 0.015 -> 0.040; Code: 0.024 -> 0.053 | Helped |
| Use KCD-LM instead of kNN-LM | GPT-2 137M · WikiText-103, Medical, Federal Register law, and Dockerfile code test sets | perplexity: WikiText: 32.19 -> 22.90; Medical: 39.66 -> 24.95; Law: 11.10 -> 8.24; Code: 89.88 -> 50.77 | Helped |
| Retrieve self-memory chunks during decoding | GPT-2-XL, LLaMA-2, and Mistral with SCD-LM · MT-Bench-80 | token time saved (TTS, %): GPT-2-XL: 13.74 -> 19.59; LLaMA-2: 2.44 -> 14.89; Mistral: -1.23 -> 11.75 | Helped |
| Retrieve self-memory chunks during decoding | GPT-2-XL, LLaMA-2, and Mistral with SCD-LM · MT-Bench-80 | forward passes saved (FPS, %): GPT-2-XL: 23.77 -> 43.33; LLaMA-2: 6.75 -> 32.32; Mistral: 5.86 -> 24.52 | Helped |
| Use a question-specific unique datastore instead of a shared datastore | SCD-LM · MT-Bench-10 | token time saved (TTS, %): GPT-2-XL: 9.28 -> 13.31; LLaMA-2: 8.42 -> 15.94; Mistral: 8.22 -> 16.39 | Helped |
| Inject expert-curated Wikipedia hyperlink chunks | GPT-2-XL, LLaMA-2, and Mistral with ECD-LM · Alan Turing QA | unique entities: GPT-2-XL: 102 -> 145; LLaMA-2: 130 -> 153; Mistral: 143 -> 160 | Helped |
| Retrieve private information from a PII datastore | GPT-2-XL / 1.5B with ECD-LM · PII questions | accuracy (%): 0.0 -> 75.7 | Helped |
| Retrieve private information from a PII datastore rather than append it in the prompt | LLaMA-2 / 6.7B · PII questions | accuracy (%): 75.5 -> 77.5 | Helped |
| Raise the retrieval similarity threshold η (ablation) | SCD-LM · MT-Bench-80 generation | similarity to base LM generations | Helped |
| Lower the retrieval similarity threshold η (ablation) | SCD-LM · Mistral response to the Mars prompt | chunk use | Mixed |
| Use a unique datastore rather than a shared datastore (ablation) | SCD-LM on MT-Bench-10 · MT-Bench-10 | forward passes saved (FPS, %): GPT-2-XL: 31.13 -> 40.72; LLaMA-2: 24.67 -> 26.01; Mistral: 17.43 -> 50.03 | Mixed |
Surprising findings
- Self-distilled chunks sometimes lower perplexity below the base LM for LLaMA-2 and Mistral, suggesting explicit self-memories improve quality.
- Retrieval is restricted to chunks sharing the entry token, so only about 0.0003-0.01% of the datastore is searched per step.
Limitations
- Paper does not optimize retrieval overhead; engineering such as quantization or pruning is left for future work.
- Decoding is greedy and uses argmax retrieval, and perplexity computation requires retrieval at every position.
- Expert distillation evaluation relies on GPT-4o judges and a small human fluency study with English experts.
Experimental setup
Datasets: WikiText-103 GitHub Code (Dockerfile subset) Medical Instruction Dataset Pile-of-Law (Federal Register) MT-Bench
Benchmarks: WikiText-103 GitHub Code (Dockerfile) Medical Instructions Pile-of-Law MT-Bench-80 MT-Bench-10 Alan Turing QA
Metrics: perplexity MAUVE token time saved forward passes saved ROUGE-L BLEURT pairwise win rate entity count PII accuracy fluency Likert
Baselines: kNN-LM RETOMATON REST GPT-2 fine-tuned on domain data
Builds on: kNN-LM REST speculative decoding RETOMATON
Closest papers by meaning
- Dataset Decomposition: Faster LLM Training with Variable Sequence Length Curriculum (NeurIPS 2024 · similarity 0.88)
- RAT: Bridging RNN Efficiency and Attention Accuracy via Chunk-based Sequence Modeling (NeurIPS 2025 · similarity 0.88)
- One Token Can Help! Learning Scalable and Pluggable Virtual Tokens for Retrieval-Augmented Large Language Models (AAAI 2025 · similarity 0.88)
- LongWriter: Unleashing 10,000+ Word Generation from Long Context LLMs (ICLR 2025 · similarity 0.87)
- Compact Language Models via Pruning and Knowledge Distillation (NeurIPS 2024 · similarity 0.87)
- SepLLM: Accelerate Large Language Models by Compressing One Segment into One Separator (ICML 2025 · similarity 0.87)
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.