Research map / Large language models / Language, safety and interpretability

Interpretable Next-token Prediction via the Generalized Induction Head

NeurIPS 2025 · Large language models › Language, safety and interpretability · approach: Pruning and circuits in LLMs · 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 the Generalized Induction-Head Model (GIM), an interpretable retrieval module that combines exact n-gram matching with fuzzy matching via a distilled neural similarity model, improving Infini-gram for language modeling and fMRI response prediction.

Problem

Interpretable next-token predictors such as Infini-gram lag far behind black-box LLMs, largely because exact n-gram matching cannot adapt to novel contexts or minor input variations like typos and rephrasings.

Approach

Generalize the induction head into an interpretable in-context retrieval module that scores sequence similarity by a small distilled Fuzzy Matching Model (trained to match LLaMA-2-7B next-token distributions) and predicts via a similarity-weighted count over tokens following matches, integrated with Infini-gram.

Key results

• Infini-gram with GIM improves BabyLM-test next-token accuracy by about 25 points over Infini-gram with OpenWebText (41.8 vs 16.7 for GPT-2 tokenizer, 43.2 vs 20.1 for LLaMA-2 tokenizer), narrowing the gap to LLMs.
• On fMRI, GIM matching achieves mean correlation 0.087 vs 0.072 for Eng1000 (20% improvement) across all voxels.
• GIM (fuzzy) with speculative decoding achieves up to 2.27x speed-up over LLaMA2-70B alone on H100s.

Measured effects reported in the paper

Surprising findings

• Infini-gram built on BabyLM-dev, a tiny 17M-token corpus, still outperforms much larger reference corpora on BabyLM-test.
• Fuzzy matching gives little gain in the fMRI setting, attributed to noise and temporal smoothing.

Limitations

• Does not close the performance gap with black-box LLMs, especially for reasoning or world-knowledge tasks.
• Relies on the input context, giving limited gains when context is short or uninformative.
• Fuzzy matching provides little benefit for fMRI, and fMRI analysis uses only one subject.

Experimental setup

Datasets: BabyLM OpenWebText Pile-train Pile-val FineWeb Pile-of-law fMRI podcast dataset (Huth lab, UTS03)

Benchmarks: BabyLM-test FineWeb Pile-val fMRI voxel correlation on UTS03

Metrics: next-token accuracy voxel-wise correlation speculative decoding speed-up acceptance rate

Baselines: Infini-gram GPT-2 LLaMA-2-7B Eng1000 LLaMA2-70B

Builds on: Infini-gram induction heads kNN-LM Eng1000

Closest papers by meaning

• RABBiT: Rapidly adaptive BOLD foundation model via brain-tuning for accurate zero-shot and few-shot prediction of speech-elicited responses in the brain (NeurIPS 2026 · similarity 0.85)
• Crafting Interpretable Embeddings for Language Neuroscience by Asking LLMs Questions (NeurIPS 2024 · similarity 0.84)
• Language Models as Zero-shot Lossless Gradient Compressors: Towards General Neural Parameter Prior Models (NeurIPS 2024 · similarity 0.84)
• On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.84)
• Interpreting the Repeated Token Phenomenon in Large Language Models (ICML 2025 · similarity 0.84)
• LensLLM: Unveiling Fine-Tuning Dynamics for LLM Selection (ICML 2025 · similarity 0.84)