Title: Interpretable Next-token Prediction via the Generalized Induction Head
Venue: NeurIPS 2025
Accepted: yes
Keywords: interpretability

ngram

language modeling

fmri

neuroscience
Abstract: While large transformer models excel in predictive performance, their lack of interpretability restricts their usefulness in high-stakes domains. To remedy this, we propose the Generalized Induction-Head Model (GIM), an interpretable model for next-token prediction inspired by the observation of “induction heads” in LLMs. GIM is a retrieval-based module that identifies similar sequences in the input context by combining exact n-gram matching and fuzzy matching based on a neural similarity metric. We evaluate GIM in two settings: language modeling and fMRI response prediction. In language modeling, GIM improves next-token prediction by up to 25%p over interpretable baselines, significantly narrowing the gap with black-box LLMs. In an fMRI setting, GIM improves neural response prediction by 20% and offers insights into the language selectivity of the brain. GIM represents a significant step toward uniting interpretability and performance across domains. The code is available at https://github.com/ejkim47/generalized-induction-head.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces the Generalized Induction-Head Model (GIM), an interpretable framework for next-token prediction.  By integrating exact n-gram matching with fuzzy neural similarity matching, GIM overcomes the limitations of traditional models.  It achieves up to 25% higher next-token prediction accuracy than interpretable baselines in language modeling and a 20% improvement in fMRI response prediction.  GIM bridges the gap between interpretability and performance, providing transparent predictions  into language processing mechanisms.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. The paper introduces the Generalized Induction-Head Model (GIM), a novel framework that integrates exact n-gram matching with fuzzy neural similarity metrics to enable transparent next-token prediction. \n2. This paper highlights GIM's interpretability through exact n-gram matching, which overcomes the black-box nature of traditional transformers.\n \nWeaknesses\n1. GIM’s exact n-gram matching relies on fixed-size sliding windows (e.g., 1024 tokens), which restricts its ability to capture long-range dependencies. When input contexts are insufficient to contain informative matches, the model’s effective n drops, leading to performance degradation, highlighting limitations in modeling extended contextual understanding.\n2. GIM is designed as a task-specific framework for next-token prediction rather than a generalizable model to decode LLM internal mechanisms. This framework does not directly address how pre-trained transformers (e.g., GPT-2) generate predictions. For example, GIM’s hybrid matching does not clarify whether LLM success stems from explicit pattern retrieval or implicit semantic compression.\n3. The fuzzy matching module, trained via knowledge distillation from LLMs, may introduce semantic ambiguity when input distributions deviate from training data. While the neural similarity metric improves robustness to minor variations, it might struggle with out-of-domain inputs.\n\nQUESTIONS:\n1. Is there any structural relationship between the GIM and the induction heads in traditional transformers? Can GIM provide interpretability for prediction mechanisms in real-world large language models?\n2. Which component of the matching strategy - the n-gram matching or the fuzzy neural similarity matching - plays a decisive role in the accurate prediction of language models? Specifically, how does the hybrid design of these two strategies help improve the prediction accuracy and robustness of the model for input changes?\n3. Given GIM's context length limitation, which improvement direction - such as dynamic window scaling, hierarchical context aggregation, or external memory integration - is most promising to enhance the model's generalization ability for long-range dependencies and sparse input contexts?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces the Generalized Induction-Head Model (GIM), a fully interpretable architecture for next-token prediction inspired by \"induction heads\" found in large language models (LLMs). Unlike black-box LLMs, GIM offers transparency by retrieving similar input sequences using both exact and fuzzy matching, enabling direct attribution of predictions to input context. It combines traditional n-gram approaches with a learned neural similarity metric to match sequences even under minor variations. GIM significantly improves prediction accuracy in both language modeling and neural response prediction tasks.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n - It builds on insights from the mechanistic interpretability literature to develop a better language model, rather than limiting those insights to controlled or toy settings.\n - It integrates Infini-gram with GIM to effectively leverage contextual information and enhance next-token prediction accuracy.\n - It evaluates the effectiveness of the proposed GIM model not only for language modeling but also for fMRI response prediction.\n\n**Weaknesses**\n - Although GIM achieves comparable next-token prediction accuracy to GPT-2 over long contexts, I'm not confident it would perform similarly to GPT-2 on broader benchmarks such as MMLU, code, or math. These tasks are crucial, as many real-world applications of language models rely heavily on strong performance in these areas.\n - While the paper compares GIM’s performance to baselines like Infini-gram, it does not compare it to pretrained language models analyzed using interpretability techniques to identify vital tokens. For instance, we can estimate the importance of an input token using the derivative of output w.r.t that token. It would be valuable to see whether the contextual information identified by GIM could also be uncovered through such methods.\n - Some parts of the manuscript could be improved; please refer to the Questions section for details.\n\nWhile the idea of augmenting n-gram language models with contextual information is compelling, it's unclear why one would choose GIM over pretrained models. The main advantage seems to be GIM's ability to attribute next-token predictions to the existing corpus and/or context. However, similar attribution can be achieved using interpretability techniques applied to pretrained models. Therefore, it’s important to compare both approaches in terms of interpretability and evaluate their effectiveness more comprehensively. Such a comparison would help determine whether the proposed GIM is practically viable.\n\nQUESTIONS:\n- Provide more details of the Fuzzy Matching Model.\n   1. Line 131 mentions that it consists of a “few transformer layers”. How many?\n   2. Line 131 mentions, it is “designed to output feature embeddings.” Please provide more details on what exactly is the output feature embeddings and how are they computed.\n - Could you describe the $w_{:i-1}$ notation in some detail?\n - Line 180 mentions the 10B-token OpenWebText dataset. However, Table 1 reports it as 9.04B. Why the discrepancy?\n - Line 181 mentions that GIM (exact) outperforms infini-gram by a margin of 5.5% to 20%. I’m unsure where is 5.5% coming from?\n - In Table 1, why is tokenizer information under Type column? Additionally, could you describe what does reference corpus mean in the context of GPT-2 and LLaMA-2-7B?\n - Line 238 mentions next-token distribution. I’m unsure what exactly it means in the fMRI setting. I understand text is converted to embedding and then a linear model is trained for each voxel. But how is this next-token distribution computed?\n - It would also be helpful if you could describe the response variable in the fMRI setting. What exactly is this response variable and how is it used to compute the principal components?\n - Line 302-305 mentions incorporating other mechanistic interpretability insights to improve GIM. How do you envision doing it?\n\nLIMITATIONS:\nThe paper already notes GIM’s limited improvement when the input context is short. Beyond that, its performance is likely to be poor on symbolic reasoning tasks such as variable binding. For example, given the following code snippet, GIM would likely predict 5 as the next token instead of the correct value 15: \n`x = 5; x += 10; x =`", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes the Generalized Induction-Head Model (GIM), an interpretable approach to next-token prediction designed to address the lack of transparency in large transformer models. GIM integrates exact n-gram matching with a fuzzy retrieval mechanism based on embedding similarity. The model is evaluated on language modeling and fMRI neural response prediction tasks. It shows improved performance over interpretable baselines and narrows the performance gap with black-box models. Unlike its closest baseline, Infini-gram, the proposed model performs matching solely within the input context, without relying on an external corpus. As the authors demonstrate, this approach can improve performance, since in-context matching better aligns with the distribution of the input query.\n\nSTRENGTHS AND WEAKNESSES:\nWhile I am not familiar with the n-gram literature for LLMs, I found GIM to be a thoughtful and well-motivated extension of ideas from Infini-gram. Conceptually, it's quite similar—both aim to predict the next token based on prior context—but GIM takes a distinct approach by restricting retrieval strictly to the input context, which makes it more scalable and predictions more aligned to the input query distribution as demonstrated in Figure 4.\n\nAnother contribution of this paper is the fuzzy matching mechanism, in contrast to Infini-gram’s reliance on exact n-gram matches which requires a large external corpus for better matching. This makes GIM more robust to small changes in phrasing or noise, such as typos, and gives it more flexibility in practice. Surprisingly, even without access to an external corpus, GIM outperforms Infini-gram, which speaks to the effectiveness of using in-context information alone.\n\nI also appreciated the paper's extension of the model to fMRI prediction tasks. This demonstrates GIM’s potential beyond standard language modeling, particularly in exploring how meaning is represented in the brain.\n\nThat said, I do see limitations. Because GIM depends solely on the input context for retrieval, its performance might not scale well in open-ended or highly diverse tasks where context is sparse or uninformative. This trade-off between interpretability and generalization is important and the authors could explore it more directly in future work.\n\nQUESTIONS:\n1. How sensitive are the results to the use of Jensen–Shannon divergence in the fuzzy matching procedure? Were other divergence measures, such as KL divergence, considered or evaluated?\n2. While Jensen–Shannon divergence is introduced formally in equation (1), the implementation appears to rely on cosine similarity between embeddings. Could the authors clarify how the formulation in (1) maps onto the practical implementation in (2)?\n3. In equation (4), what is the purpose of the indicator function? It seems the summation is not necessary if we only look for cases where w_j = w_i.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper tackles the trade-off between performance and interpretability in language models by introducing the Generalized Induction-Head Model (GIM). Drawing inspiration from “induction heads” observed in large language models, GIM is a retrieval-based, interpretable model that operates entirely within the input context. It combines two core components: an exact n-gram matcher that flexibly identifies the longest possible match (rather than relying on a fixed n), and a fuzzy matcher based on a neural network to handle variations. GIM complements models like Infini-gram by offering richer in-context retrieval, helping address the limitations of external corpus-based search.\n\nEmpirically, GIM closes much of the gap with black-box LLMs: it improves next-token prediction accuracy by up to 25% over interpretable baselines and boosts fMRI response prediction by 20%. The model also proves useful for speculative decoding, enabling lossless acceleration of LLMs. Overall, this work shows how insights from reverse-engineering LLM internals can be repurposed to build transparent models that remain competitive—challenging the notion that interpretability must come at the expense of performance.\n\nSTRENGTHS AND WEAKNESSES:\nThis paper presents a compelling approach to improving language modeling while maintaining a high degree of interpretability. Conceptually, the architecture draws a neat parallel with modern Transformers: the Infini-gram module resembles an interpretable feed-forward network, retrieving world knowledge from an external corpus, while the Generalized Induction-Head Model (GIM) plays a role analogous to attention, capturing local, in-context patterns in a transparent way. This line of work is important for probing the performance limits of interpretable models. Insights from this research could help guide the design of future architectures that are both effective and transparent.\n\nThat said, the current model still falls short of state-of-the-art black-box LLMs in terms of raw performance—despite notable gains over prior interpretable baselines. A promising direction for future work would be to explore hybrid approaches that integrate GIM-like mechanisms with black-box models, especially for tasks requiring strong in-context learning. Such combinations may strike a better balance between interpretability and capability. Overall, I see this as a valuable and timely contribution to the field.\n\nQUESTIONS:\nHave you considered using GIM as a layer or integrating it as an attention head, similar to how the N-gram head was implemented in [1] ?\n\n[1] https://arxiv.org/abs/2401.12973\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}