Interpretable Next-token Prediction via the Generalized Induction Head
Authors: Eunji Kim
Sriya Mantena
Weiwei Yang
Chandan Singh
Sungroh Yoon
Jianfeng Gao
Venue: NeurIPS 2025 poster ()
Keywords: interpretability
ngram
language modeling
fmri
neuroscience
Primary area: neuroscience_and_cognitive_science
TL;DR: We introduce the Generalized Induction Model (GIM), a retrieval-based in-context module that enhances interpretable next-token prediction in language modeling and fMRI response prediction.

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.

Interpretable Next-token Prediction via the Generalized Induction Head
Eunji Kim
Sriya Mantena 1.3+
Weiwei Yang!
Chandan Singh'
Sungroh Yoon
Jianfeng Gao 11
Microsoft Research
2 Department of Electrical and Computer Engineering, Scoul National University Stanford University
Interdisciplinary Program in Artificial Intelligence, Scoul National University
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-HHes Head Model (GIM), interpretable model for next-token prediction inspired by the observation of "induction heads" in LLMs. GIM retrieval-based module that identifies similar sequences the input context by combining exact n-gram matching and fuzzy matching based on a neural similarity metric. We evaluate GIM in two settings: lan- guage modeling and fMRI response prediction. In language modeling, GIM improves next-token prediction by up to 25%p over interpretable baselines, sig- nificantly narrowing the gap with black-box LLMs. In fMRI setting, GIM improves neural response prediction by 20% and offers insight into the lan- guage selectivity of the brain. GIM represents significant step toward unit- ing interpretability and performance across domains. The code available at https: htapp:://gtthue wwt.cooo//eeeeeeeee
Introduction
While modern transformer models have achieved impressive performance across a wide array of next-token prediction tasks [1-3], these models remain black-boxes, limiting their use in real-world applications. Their opacity is detrimental fields such as neuroscience [4] and social science [5], where trustworthy interpretation, specifically, token-level attribution that traces outputs back to input data, is often the end goal. Their lack of transparency also hinders adoption in high-stakes applications such medicine [6], raising concerns around regulatory compliance, safety, and alignment [7-10].
As alternative these black-box models, interpretable models have been proposed for various tasks [11-13], but they continue to struggle on the task of next-token prediction. For example, in next-token prediction for natural language, the state-of-the-art interpretable model is Infini-gram [14], which trails GPT-2 by 30%p on the BabyLM dataset (see Table 1). Our analysis suggests that this performance gap stems from Infini-gram' 's inability to adapt novel contexts or handle minor input variations such as typos and rephrasings.
"Equal contribution. Work conducted during internship Microsoft Research. 'Corresponding authors.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

We address this gap by proposing the Generalized Induction- Head Model (GIM). GIM is inspired by the observation of "induction heads" in LLMs [15, 16] that support in-context learning by detecting and extending patterns in prior input. In pre-trained LLMss, this behavior arises implicitly and is inferred through post-hoc approximations over dense internal states. In contrast, GIM is not a post-hoc tool for interpreting opaque systems, but an inherently interpretable system that explicitly models this behavior in a transparent and auditable manner.
GIM is an interpretable retrieval-based framework that operates entirely within the model's 's input context to retrieve suggestions for next-token completion. It extends traditional exact matching by incorporating a lightweight fuzzy similarity function to match sequences that yield similar next-token distributions. Importantly, this neural component is used solely for scoring similarity between input phrases rather than generating outputs. The final next-token prediction is computed as a similarity- weighted distribution over tokens that follow the matched phrases, enabling each prediction to be directly attributed to specific input sequences, supporting full interpretability and human auditability. GIM is designed as a standalone, model-agnostic module that supports both exact and fuzzy matching and can be integrated across modalities.
We first evaluate the performance and interpretability of GIM in next-token prediction for language modeling. We integrate GIM into Infini-gram, and GIM improves next-token prediction accuracy by 25%p over Infini- gram using Web [17] as a reference corpus, significantly narrowing the performance gap with GPT-2 (see Table 1).
Second, focus on a single, real-world neuroscience problem, deviating from typical machine- learning conference paper. Grounding in a neuroscience context allows us to avoid common pitfalls in evaluating interpretation methods [18, 19] that seek to test "interpretability"" in the abstract. We find that when used to predict IMRI responses to language stimuli, GIM yields a 20% improvement over the state-of-the-an interpretable model (see Table 2), and its transparency enables the attribution of predicted neural responses in each region across the cortex to specific linguistic features.
Taken together, these results challenge the assumption that interpretability and predictive performance are fundamentally at odds, showing that reverse-eng eeceese-nngeerred neural components can be leveraged to enhance transparency. Importantly, our aim is not to claim parity with black-box LLMs, but to show that meaningful gains in predictive performance can be achieved without compromising transparency.
2 Related Work
N-gram language models Early language modeling techniques revolved around n-gram models [20, 21J, which generally stored next-token probabilities in large tables learned from data [22]. While largely surpassed by neural LLMs, recent works have continued to improve n-gram LMs, e.g., by scaling up the n n-gram reference data [23] and improving the n-gram probability representations using suffix arrays and suffix trees [24-26]. This line of work culminated in Infini-gran (aaii [14], which efficiently scales n-gram models to massive datasets and is the starting point for our work.
Bridging interpretable models and LLMs Some works have studied bridging n- n-gram models and LLMs. For example, Khandelwal et al. [27] interpolated neural LMs with n-gram model and Li et al. [28] trained a neural model to complement an n-gram model. Other approaches augment black-box LMs with nonparametric components, such as k-nearest neighbors [27, 29]. While these methods improve performance, they lack transparency in prediction behavior and token-level attribution.
Interpretable models have been proposed for simplified settings such as text classification. Some offer fully interpretable decision processes [30-32], while others offer partial interpretability by approximating model behavior with natural language concepts [33-37]. However, these approaches not designed for open-ended generation.
In parallel, there has been a recent surge of interest mechanistic interpretability, which seeks to understand what mechanisms are learned by transformer-bast based LLMSs [38-41]. This line of work identified induction heads toy LIM models [15] as well as large-scale pre-trained L.LMS [42, 16). Despite these efforts, frameworks to make these findings useful in real settings remain underexplored.

---

(a)
(b)
AVG: 5.23
Nex-tolen Acciracy
0.8
AVG: 6.86
0.6
Count
6000
Reference Data of lokens)
BabyLM 118
0.4
AVG: 0.39 22.2.
4000
OpenWebllex (394M)
train (3838)
0.2
2000
0.0
10
15
20
25
30
10
15
25
30
Effective
Effective
Figure I: Performance on the BabyLM dataset with Infini-gram built from various reference datasets. (a) Next-token prediction accuracy by effective with the dashed line indicating the average. (b) The histogram of the count for each effective n.
Natural language representations in fMRI recent years, predicting brain responses to natural language using LLM representations has become common in the field of language neuroscience [43- 47]. Predictive "encoding models" have been used to explore the relative contributions of syntax, semantics, and discourse to neural activity [48-55] and to study the cortical organization of language timescales [56, 57], sometimes making use LLMs to help annotate and generate stimuli [58- 61]. However, these models largely operate as black boxes. While they reveal which language representations best predict neural activity, they offer limited insight into where in the cortex these features exert their influence when in the linguistic stimulus they become relevant.
Separately, behavioral studies have examined how humans recall and process repeated text [62-65] and how similar recall patterns emerge in LLMs [66, 67]. Yet the cortical mechanisms involved in contextual recall remain unclear, motivating our investigation through interpretable modeling.
3 Methods
We begin by discussing Infini-gram, the scalable n-gram method for interpretable next-token predic- tion (Sec. 3.1), and introduce the Generalized Induction-Head Model (GIM) (Sec. 3.2). In describing the method, note that the method straightforwardly generalizes to generic next-token prediction tasks (e.g., fMRI responses, time series, video frames). We later show how GIM improves interpretable next-token prediction by integrating it with Infini-gram for language modeling (Sec. 4) and combining it with linear regression for fMRI response prediction (Sec. 5).
focus the familiar scenario of next-token prediction for language modeling. but
3.1 Preliminaries and Motivation: Infini-gram
Given an input text sequence, Infini-gram [14] searches a reference corpus for the longest exact suffix match to the input, then calculates the next-token distribution based on the token following each of the matches. This search made efficient by building large-scale suffix arrays that can scale to trillions of reference tokens. The length of the longest match is referred to the effective n, with the accuracy of the estimated probabilities increasing as the effective n becomes larger.
Infini-gram is limited by its reliance on exact matches, which becomes problematic under distribution shifts between the input and reference corpus. For instance, when evaluating on the BabyL [68] test set, Infini-gram built on larger corpora, such as OpenWeble on average, has fewer instances of higher effective n compared to the model built on the BabyLM (Fig. 1). With far larger corpora like Pile-train [69], Infini- gram is able to increase the number of instances with a high effective n, resulting in improved performance. However, Infini-gram built on BabyLM, which contains only 0.005% of the tokens found in Pile-train, still achieves the highest performance. This highlights the difficulty Infini-gram faces when there is substantial gap between the reference corpus and the input prompt, making hard to find matching cases with large effective n. We address this limitation with the concept of the induction head.
[17], shows lower performance and,
"cttpe: h/bbbyla....////// wit.u..
3

---

(a)
Next-token probability
Distillation
s x:)
Feature Embedding
LLM
KLD
Fuzzy Matching Model
22
2g
24
22
2g
24
27'
22
24
Neg
Pos Neg
Neg Pos Neg
22
23
(b)
Input Prompt
Mr. and Mrs. Smith were chatting in the Living poom when Hrs, Smith toned a atrange noise, promptinng Mr.
8 $FM()
104 A h-1-- 1.4-1
Figure 2: (a) Overview of training Fuzzy Matching Model via distillation from pretrained LLM. (b) Calculation of sequence similarity within the input prompt for next-token prediction.
3.2 Building a Generalized Induction Head
LLMs excel in-context learning by capturing the statistical distribution of tokens in a given context. One key mechanism enabling this capability is the induction head, a critical component in LLMs responsible for recognizing and extending repeated sequences [15, 16, 42]. Induction heads operate by detecting prior occurrences of a token sequence and leveraging this recurrence for next-token prediction (e.g., given [A][B] ... [A], the model predicts [B]. However, these heads emerge implicitly within LLMs, making their operation difficult to interpret and control.
To this end, we introduce the Generalized Induction-Hedde Model (GIM) that explicitly models this behavior in a structured, interpretable manner. functions like Infini-gram but is restricted to the input context. GIM treats the end of the context query, searches for the best match within the context and takes the token following the match as the next-token prediction.
What constitutes "good match"? When identifying n n-gram-level matches in context, exact matching can perform well if a high effective n is guaranteed (Sec. 3.1), but it can be overly restrictive minor changes such as rephrasings or typos. To remedy this, allow for fuzzy n-gram matching, which makes the model more robust to minor changes. Since the fuzzy matching is performed at the level of n-grams, predictions remain interpretable and auditable by a human.
Fuzzy matching requires appropriately computing the similarity between sequences. While similarity can be defined in many ways, in building an induction head, desire two sequences to be similar if they yield similar next-token distributions. To quantify this, we define the similarity between two sequences, 21 and 22, for fuzzy matching using -hannon divergence (JSD):
s(x1, x2) = exp (-JSD (P (21) (nex)))))))
where (-)x is the estimated next-token probability distribution for given sequence.
Computing s efficiently One approach for computing would be to use a pre-trained LLM to obtain Pnexa. but this can be computationally expensive. Instead, we develop small Fuzzy Matching Model, which consists existing LLMss This model is designed to output feature embeddings that facilitate the calculation of next token probabilities for similarity assessments. With the Fuzzy Matching Model, the similarity between 1j and 222 whose feature embeddings from the model are 61 and e2, is obtained as follows:
few transformer layers and is trained via knowledge distillation from
cxp (- (1 CosSim (e1, e2)) /T), (11,
where T is a temperature, which is set to 0.1. The Fuzzy Matching Model is trained with combination of Cross Entropy (CE) loss and reverse Kullback-Leible divergence (KLD) loss (Fig. 2(a)). In each training batch, we generate similarity pairs from randomly sampled sequences. The CE loss aids in identifying the most similar pairs. The reverse KLDD loss guides the model to follow the overall similarity distribution, ensuring that close pairs receive high scores while distant pairs receive low scores. Further details can be found Appendix A.2.
Predicting the next token Given the similarity scoring function SyM- we construct an induction head that yields the predicted next-token probability distribution pffurry) given input sequence

---

Input Prompt
submiLLed paper
InformaLion
Mhill thera
another International contasence deep leazning,
GIM (exact)
GIM (luzzy)
Sequence from Inpul Context
Similarity
Next-lokn
4 (Rraay) (M)
Metched Sequence
from Input Context
Next-toknn
inlucina (sla)
ry piper to the piger to the confersnce the conference
Matching
0.25
0:11
conference
0.25/SUM
0.1 0.11//S0
Naural
0.98/BUX
conference on conference
Neural
1/2
1/2
... another Internationa. conference
0.76
0.76/SUX
choee conference oc
Figure 3: Overview the GIM pipeline. GIM predicts the next token by efficiently searching for potential next-token completions in the input context with either exact or fuzzy matching.
J. To achieve this, we identify matches for the end of x, w using a sliding window size (Fig. 2(b)). We then count the occurrence of each token wy in the vocabulary set V following these matches and normalize to obtain the next-token probability:
Ciuzzy (w, (wi-k- k- 1:1- (x)
induction (w;-1
(3)
Ew, EV Cluzzy --i-
where Cruzzy (W; h- 11--
Σ
Imy=w; 9PM
This similarity score serves floating count for the next token. In cases where the sequences
Ii and 22 are exactly matched, as in the case of Infini-gram, we have $TM (2)) 1, which is equivalent to increasing the count by one. The window size k specifies the number of tokens to be considered in fuzzy matching.
3.3 Prediction of Generalized Induction head Model
By employing both the Infini-gram algorithm and Furzy Matching Model, GIM searches for the most relevant aatchh-iihhrr exact or fuzzy-within the preceding tokens given query at the end of the context (Fig. 3). Once match is identified, retrieves the token that followed the prior occurrence as the next-token prediction. By explicitly modeling this process, our method provides transparent and controllable alternative to implicit in-context learning mechanisms LLMs.
Results: Next-token Prediction for Language Modeling
4.1 Experimental Setup
Datasets evaluation We use 4 text datasets for evaluation: BabyLM [68], OpenWebText [17], Pile [69]. and FineWeb ([70): sample-10BT subset), using some the reference corpus and some as test datasets (Table 1). When testing, we report performance on 100k sequences randomly sampled with context length of 1024 and stride of 512 [14, 27]. We evaluate next-token prediction via accuracy, i.c whether the top-predicted token was the correct token.
Baselines We compare against Infini- gramm sole baseline, it is the state-of-the-arr n-gram model and the only fully interpretable model with token-level attribution for generation. We found that it consistently outperformed prior interpretable language models, e.g., standard 5-gram model based on Open Web Text achieves 26.4% accuracy in next-token prediction on the Pile-val, lower than Infini-gram's 's 27.1 %. For detailed discussion of interpretable frameworks, please refer to Sec. 2.
'The BabyLM test set contains fewer than 100k sequences, yielding approximately 32k and 34k cases for the GPT-2 and LLaMA-2 tokenizers, respectively.
'We do not use perplexity, as the sparse next-token predictions from n-gram models often assign zero probability to the top-ranked token, resulting in undetined extremely high perplexity scores [14].
5

---

Table 1: Next-token prediction accuracy (%) for language modeling. Gray shading represents
alignment between the reference corpus and the test dataset.
Reference Corpus
Test Dataset FineWeb
Model
Type
# of Tokens
BabyL.M-test
Pile-val
Tokenizer: GPT-2
GIM (exact)
36.7
17.2
37.0
GIM (fuzzy)
41.1
25.2
38.7
Infini-gram
37.6
14.7
16.0
M-
17.4M
+GIM
42.2 (14.6)
25.3 (+10.5)
40.0 (+24.0)
Infini-gram
16.6
20.1
Pile-val
383M
+GIM
41.5 (+24.)
25.5 (+5.4)
Infini-gram
16.7
25.5
22.7
OpenWebText
9.04B
+GIM
41.8 (x25.1)
27.2 (+1.7)
42.7 (420.0
Unknown
~0B
LLM (GPT-2)
46.9
39.0
52.3
Tokenizer: LLaMA-2
GIM (exact)
37.0
19.6
GIM (fuzzy)
42.7
28.3
38.5
Infini-gram
39.0
17.1
13.2
BabyLM-dev
18.9M
+GIM
43.1 (+4.1)
28.6 (+11.5)
39.6 (426.4)
Infini-gram gram
19.0
24.1
Pile-val
394M
+GIM
42,9 (x23.9)
28.4 (+4.3)
Infini-gram
20.1
29.5
27.1
OpenWebText
10.3B
+GIM
43.2 (+23.1)
30.3 (+0.8)
42.1 (+15.0)
Infini-gram
33.5
39.3
49.2
Pile-train
383B
+GIM
49.4 (+15.9)
38.0
50.3 (+1.1)
Unknown
2T
LLM (LLaMA2-7B)
62.2
57.1
64.4
Integrating GIM with Infini-gram For language modeling, we integrate GIM with Infini-gram, enabling the use of both reference corpus statisties and in-context distributions:
P texact) (y/ær)
Thoo and 10
Ply[=) =
P
induction (y|z)
(exact)
floo and n2 >
(5)
(furzy)
(y|z)
Otherwise,
where no and fi the effective when matching from reference corpus the input context, respectively. When these values are low, fuzzy matching is employed to compensate for the limited effective n. When the effective n values from both the input context and reference corpus are equal, the estimate from the input context is prioritized. The hyperparameter 7 determines how frequently exact matching is used over fuzzy matching; we tokenizers, respectively. based on cross volosavali results (see Appendix A.3 for details).
8 and 9 for the GPT-2 and LLaMA-2
4.2 Improving Next-token Prediction Accuracy with Contextualization
Prediction performance of in-context matching GIM relies solely on the input context to predict the next token (limited to 1024 tokens in our evaluation). Table shows that, despite this, GIM (exact) can outperform Infini-gram, which uses the Open WebText dataset as a reference corpus, comprising approximately 10B tokens when tokenized with LLaMA-2 and 9.04B with GPT-2, by margin of 5.5%p to 20%p on the BabyLM and Pile datasets. Infini Infini-gram gram using BabyLM-dev as the reference corpus slightly outperforms GIM (exact) on the BabyLM with performance gaps of 0.9%p and 2.0%p for the GPT-2 and LLaMA-2 tokenizers, respectively. under the aligned setting of reference corpus and input context. As shown in Fig. 4(a), Infini-gram fafini-gaam (green) performs better in cases with high effective n. even surpassing LLM (blue). However, significantly more cases have a low effective (histogram), where GIM (exact) (orange) outperforms Infini-gram. This finding underscores that in-context matching reflects the input query's distribution, leading to more accurate next-token predictions than reference matching, even when the reference corpus contains abundant tokens, especially under distribution shifts between the reference corpus and the test input.

---

(a)
(b)
10000
8000
Noxt token Accuracy
Accuracy
0.8
8000
6000
0.66
Count
6000
0.6
4000
0.4
4000
0.4
0.2
2000
0.0
10
20
25
15
20
25
Effective from nffini gramm IPlle train)
Effective from GIM (exact)
LIM (LLaMa2-78)
Infinigramm (Pi (Piestrai))
(ecact))
Huzzyy
Figure 4: Comparison next token prediction accuracy on BabyLM-tes dataset, depending on effective from (b) Infini-gras gram and (b) GIM (exact). LLaMA-2 tokenizer is used.
(a) Input Prompt:
Breaschuickersheiin
commune,
the*
Exach-Matched Sequence in
Token
Effective
the* Elsenheim commune. It 1a in Grand Est In the* Bourgheim is commune. It is in Grand Est in the"
"Baa"
Ohnenheim 1s
* "Bas"
12
"Bas"
(b) Input Prompt:
Simpson still delays taking kick,
"in"
Fuzzy-Matched Sequence Context
Next Token "in"
Similarity
breakaway down the left,
coming"
0.160
Three minutes later Simpson ran" "aaa but he grabbed it again at the second attempt before it went" just little wide.
"on"
0.083
"over" 0.075
"... was
ran* "into"
0.075
Figure 5: GIM's token-level attribution by tracing predictions (a) exact (b) fuzzy matches in prior context. Yellow highlight shows the match, and red box marks the source of the final prediction.
Prediction improvements from GIM GIM (fuzzy), using Fuzzy Matching Model, consistently outperforms GIM (exact) with margin of 1. 1.7%p to 8. 8.7%p (Table 1). This improvement is particularly evident in cases with low effective n. As illustrated in Fig. 4(b). the majority of cases within the input context have low effective (histogram), indicating that finding exactly matched long sequences within the limited amount of tokens challenging. Fuzzy matching helps to provide better estimations for next-token predictions in these scenarios. Specifically, when the effective is less than 3, GIM (fuzzy) (yellow) demonstrates better performance than GIM (exact) (orange). Since many cases fall into this range, the overall accuracy of GIM (fuzzy) is higher.
The improvements achieved through the use of induction and fuzzy matching enable Infini-grane with GIM outperform infini-gramt built on 383B tokens, improving performance by up 15.9%p. While enlarging the reference corpus boosts performance, GIM offers a more efficient alternative to scaling from 10.3B to 383B tokens- 38-fold increase. Moreover, GIM is complementary approach that can be applied orthogonally to Infini-gram, regardless of the size of the reference corpus.
4.3 Qualitative Example of GIM Prediction
Fig. 5 shows examples of explanations provided by GIM. In the first case, the prompt exactly matches a 12-gram in the context, so GIM follows to predict the next token. the second, no exact match exists for the prompt ending in "comes", but GIM finds the most similar sequence ending in "coming" and follows it for prediction. These cases illustrate how GIM predicts from retrieved sequences, with transparency into which tokens contribute and how they are combined.
5 Results: Next-token prediction for fMRI Responses to Natural Language
Understanding how and where semantic information is represented across the human brain is central objective in neuroscience. In this work, we extend prior modeling frameworks that learn mappings between natural language stimuli and corresponding neural responses across voxels, which are small three-dimensional regions of the brain. [71, 43]
7

---

5.1 Experimental Setup
We analyzed publicly available data from [72] and [73], in which three human participants
Text
TR Features Concalenation 1000 and PCu
[passus
says you got
listened to 20+ hours of English-laguage pod-
Lancoos Interpolation
Interpolaned Eng embecdings
cast narratives while their fMRI responses were recorded across 95,556 cortical voxels. Our goal was to predict the brain response of each voxel from the language input heard by the participant'. We extracted text embeddings from the input story, then fit linear models to map these embeddings to fMRI responses the training split (24 stories), and evaluated performance the test split (2 stories) using bootstrapped ridge regression. Embeddings are extracted in various ways (described below) for each word in the input, and then interpolated to make predictions for the fMRI data that is recorded at 2-second time of repetition (TR) intervals. To model temporal delays in the IMRI signal, we add 4 time-lagged duplicates of the input features. See more fMRI details in Appendix A.6.
Eng1000 ambidding
Top 100 PCs R, Ry
Text near matched th:
TR
Text
En 22 jchis qeeat and Hace)
paik the (win]
lett oathens
MRR Respoonee
R. Eshaatio
R, Pya
Ry Rya
Figure 6: fMRI feature construction. (a) At each TR 1, we retrieve a prior TR 1 with various matching methods. (b) We extract the top 100 principal com- ponents (PCs) of neural response changes before t*, (c) These are concatenated to interpolated Eng 1000 embeddings for fMRI signal prediction at TR t.
IMRI prediction baselines We use Eng1000 as primary baseline, the tfte-of-theear inter- pretable model for predicting fMRI responses to narrative stories from a seminal study of language selectivity [71]. Each element in an Eng embedding corresponds
co-occurrence statistic with a different word. We also compare against LLaMA2-70B [74] embeddings, which achieve the highest performance on this task [75] but are not interpretable. LLaMA embeddings extracted with a 16-word sliding window, using the final-layer embedding of the last token in each window.
GIM for fMRI prediction We construct our GIM for fMRI by searching the preceding story text in an IMRI session for semantic matches and retrieving the changes in the recorded brain response that follows each match. Specifically, to predict the fMRI response, Re, for the TR L, we first find the TR/ tr for which the text input yields the highest cosine similarity to the next-token distribution of the text input at TR L 1. Next, we isolate the change fMRI responses following TR t": we take the difference in the top 100 principal components of the response R.- Ra- and use them as features. To deal with potential time delays in the fMRI signal, we additionally concatenate these features with the top 100 principal components of Rç- Ri--2 and Rç- R These features, along with the interpolated Eng Eng100 1000 embeddings, form the full input to the linear model predicting fMRI response at TR t (see Fig. 6). When constructing the GIM for fMRI, we search over the most recent 1024 words and their corresponding fMRI responses. To measure similarity between two texts, we use the predicted next-word distributions yielded by GIM (exact) in the input context (P Prenetto in Eq. (5)), which call GIM matching.
Baseline matching methods We compare GIM matching against three baseline matching strategies. First, we use the predicted next-word distributions yielded by exact n-gram matching in the 10B- token Open OpennWeTex, reference corpus (exatt) in Eq. (5)), which we call Infini- gram matching. Second, Random matching selects a random preceding TR as a match. Third, Naive n-gram matching searches for exact match to the most recent 4-word n-gram in the input context, without relying on predicted next-word distributions that our GIM matching method relies on. Table A6 shows additional experiments with fuzzy matching methods that show little performance gain, likely due to noise and temporal smoothing in fMRI signals that diminishes the advantage of fuzzy matching.
5.2 Prediction Improvements from GIM Matching
Table 2 shows the average correlation values across all voxels for each similarity model. Eng 1000, the primary interpretable baseline, achieved a mean test correlation of 0.072. In contrast, GIM matching
6 bttpp://github haw//pppelllorrrstttteeeaaaaaooe
We report results for subject UTS03 due to high fMRI data quality, including superior repeatability, minimal motion, and strong encoding model performance [72].
8

---

(4) Diffeencce Prediction ferfrmance Between GIM Masching and Eng 1000 Baseine
Table 2: IMRI test prediction performance for different models. Black-box encodings use LLaMA-2. Error bars show 95% CI.
Mean Correlation
Feature Model
Top 10% Voxels
Voxels
Matching Precistion Performance
Fing 1000
0.072:100044
0.220=0.0012
Voxel on Performance
Random matching
0.069:10.003
0.197:000012
(Correlation)
Naive ngram matching
0.068:+0.0003
0. 0.0000
+ Infini-gram matching
0.069:40.0003
0.200=0.0012
GIM matching
0.087:10.005
0.265
Figure 7: (a) Difference in the correlation per- formance between the GIM matching and the Eng1000 baseline, visualized across the cortex. (b) Correlation performance of GIM matching.
Black-box encodings
0.096::0.0005
0.208
achieves a mean correlation of 0.087, 20% improvement over Eng1000. When predicting the top-10% of voxels, GIM Matching achieves mean correlation of 0.265, again 20% improvement over Eng1000, and only 1% lower than the black-bos box LLaMA-2 model (mean correlation 0.268). In contrast, other matching-based baselines unable improve over Eng 1000: The Naive n-gram matching baseline achieves a correlation of 0.068, and random matching achieves a correlation of 0.069.
Fig. 7 visualizes voxel-wise differences in test correlation performance between GIM matching and the Eng 1000 baseline across the cortex. In line with prior studies linking model performance to functional localization [47, 45, 71, 43], GIM significantly improves prediction in regions highted in red such the Occipital Face Area (OFA) and Intraparietal Suleus (IPS). These gains reflect GIM's use of contextual input, in contrast to the static embeddings used by Eng1000, and suggest that these highlighted regions may contribute contextual language processing.
Describing improvements from GIM To understand the improvements provided by matching, we summarize the text for inputs where each matching procedure (GIM and Infini- Infini-gram) performs well. We use an LLM to do the summarization, following recent works in LLM interpretability [76, 77]. We first identify phrases in the input story where model's performance (average absolute error across voxels) exceeds the baseline performance by more than one standard deviation (see an example in Fig. 8). Then, prompt GPT-4 ([2); gpt-4-0613) to generate descriptions for these phrases.
Fig. 8 presents the unedited L.LM descriptions GIM matching is described as capturing Emotionally or Narratively Critical Phrases, aligning with the idea that induction improves performance by tracking local context in story, e.g., phrases that "are critical to the plot and character development". In contrast, infini-gram matching is described as capturing Brief. Stand-Alone Phrases, matching the intuition that Infini-gram excels in capturing context that is not story specific, but "can stand alone with minimal context".. To test these descriptions, we prompt GPT-4 to classify the identified phrases in two test stories using only the descriptions. This yields 61% % accuracy, a moderate but significant improvement over chance (binomial test p 0.032). See all phrases and prompts in Appendix A.6.
6 Discussion
GIM constitutes a significant step toward building mechanistically interpretable language models inspired by pre-trained LLMs. Unlike black-box models or partially interpretable approaches. GIM provides full transparency in next-token prediction while substantially narrowing the performance gap between interpretable and black-box architectures across two diverse domains. Importantly, GIM is not a general-purpost LLM or a tool to decode its internals: it isolates and reimplements a single observed capability, induction via repetition, as fully interpretable module. This shows that high- performance behaviors implicitly learned by LLMs can be transparently reconstructed to advance
Irrelevant preceding text such as "Sure here is the answer" is removed from the response.

---

TAM
TON
20T2
THE
1000
200
RO0S
FMC
Tan
4V07
54.
LLM Summerization
Vill Graannaatt
lock reatting
No qpalias Sieennc
(b) Moddl
Decription
GIVI Matching
Emotioaaly Narrativlyy
Thme planes
iregoned wowterttteeeecoo and conzed the stay's resolaiios
(c)
Chcaalal froes
Maiching
Drief Stand Alone Phrass These then planes, shmn
staeece, dar can sind alaae with muinal counest. They tyskally
eream connire, eveylay woscs ard independox.
C LLM fiaaiin
Figure 8: Qualitative illustration of how GIM and Infini-gram matching improve performance. (a) Highlighted phrases in the input story where the model outperforms the baseline. (b) Summary of the highlighted phrases by an LLM to characterize each matching method. (c) Classification of highlighted phrases in test stories based on the LL -LM-eerrratedd summaries.
performance in the interpretable modeling space. The transparency of GIM makes it well-suited for language modeling scenarios that require complete auditing, such as analyzing scientific texts or medical notes [78]. GIM's transparency also supports neuroscience research, the fMRI analyses conducted here are a suggestive starting point for understanding how context is stored and recalled in the human cortex. GIM can further serve testbed for analyzing how context modulates the recall of specific semantic categories, like people and places, across the cortex, extending prior work with static embeddings [71]. Additionally, improvements from GIM Matching may help build encoding models that can more rapidly adapt to local context, which can be used in downstream applications such brain decoding [73] or brain-computer interfaces [79].
GIM shows limited gains when the input context is short or uninformative. Its modular design enables the available context to be expanded through
generation [80] or external memory. Like KNN-LMs [81], GIM's n- n-gram-basse based reasoning also struggles with tasks requiring deeper reasoning. Future work may explore hybrid approaches that pair GIM with black-box models for better trade-offs. Our speculative decoding setup, where GIM serves as a transparent draft generator verified by larger L.M (Appendix A.5), illustrates one example in this direction. Another promising direction is expanding GIM beyond induction heads, integrating additional mechanistic components
such indirect object identifiers [42], numerical representations [82], retrieval heads [80], iteration
heads [83], concept-leved induction heads [84]. instruction-followinn heads [85]. or interpretable LLM submodules [86-88]. Finally, GIM's ability to model context-dependent patterns makes well-suited for other sequential domains that require interpretability, e.g., it could be extended to study long-range dependencies in electronic health records [89], audio/speech models [90, 91], genomics [92] or financial time-series analysis [93].
Acknowledgements
We would like to thank Lucas Liu, Ziyang Wu, and Paul Smolensky for insightful discussions and feedback throughout this work, which significantly contributed to the development of our ideas. This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) RSo (No.RS-2021-1D 2021-1121134,, Artificial Intelligence Graduate School Program (Seoul National University), RS-2022-1122099 No.RS-2025 02263754, Clmann-Centris Embodied AI Agents with Autonomous Decision- Makingl, the BK21 FOUR program of the Education and Research Program for Future ICT Pioneers, Seoul National University in 2025, the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. 2022R 107777772, 2022R 75833839081 1), and InnoCORE program of the Ministry of Science and ICT (1.250021.01).
10

---

References
[1] Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neclakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
[2] OpenAL. GPT-4 technical report, 2023.
[3] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llaa 3 herd of models. arXiv preprint arXiv:2407221783 2024.
[4] Shailee Jain, Vy A Vo, Leila Wehbe, and Alexander G Huth. Computational language modeling and the promise of in silico experimentation. Neurobiology of Language, 5(1):80-106, 2024.
[5] Caleb Ziems, William Held, Omar Shaikh, Jiaao Chen, Zhehao Zhang, and Diyi Yang. Can large language models transform computational social science? Computational Linguistics, 50(1):237-22 2024.
[6] Arun James Thirunavukarasu, Darren Shu Jeng Ting, Kabilan Elangovan, Laura Gutierrez, Ting Fang Tan, and Daniel Shu Wei Ting. Large language models medicine. Nature medicine, 29(8):1930- 194199999999,,,---1 2023.
[7] Bryce Goodman and Seth Flaxman. European union regulations on algorithmic decision-making and "right to explanation". Al magazine, 38(3):50-57, 2017.
[8] Dario Amodei, Chris Olah, Jacob Steinhardt, Paul Christiano, John Schulman, and Dan Mané. Concrete problems AI safety. arXiv preprint arXiv:1606.06650 2016
[9] lason Gabriel. Artificial intelligence, values, and alignment. Minds and machines, 30(3):411-433 2020.
[10] Chandan Singh, Jeevana Priya Inala, Michel Galley, Rich Caruana, and Jianfeng Gao. Rethinking interpretability in the era of large language models. arXiv preprint arXiv:2402.07661 2,..4400..176,,22244 2024.
[11] Cynthia Rudin, Chaofan Chen, Zhi Chen, Haiyang Huang, Lesia Semenova, and Chudi Zhong. Interpretable machine learning: Fundamental principles and 10 grand challenges. Statistic Surveys, 16:1-85, 2022.
[12] Amaud Mignan and Marco Broccardo. One neuron versus deep learning in aftershock prediction. Nature, 574(7776;B1-E33 2019.
[13] Wooseok Ha, Chandan Singh, Francois Lanusse, Srigokul Upadhyayula, and Bin Yu. Adaptive wavelet distillation from neural networks through interpretations. Advances in Neural Information Processing Systems, 44:20669-2063 2021.
[14] Jiacheng Liu, Sewon Min, Luke Zettlemoyer, Yejin Choi, and Hannaneh Hajishirzi. Infini-gram: Scaling unbounded n-gram language models to trillion tokens. In First Conference on Language Modeling, 2024.
[15] Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann. Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads, arXiv preprint arXiv:2209. 200.2.09....9..... 2022.
[16] Ekin Akyürek. Bailin Wang. Yoon Kim, and Jacob Andreas. In-context language learning Arhitectures and algorithms, 2024.
[17] Aaron Gokaslan and Vanya Cohen. Openwebtext corpus. OpenkebTextCorpps. 2019.
http://Skyl1onn007 titu..
[18] Julius Adehayo, Justin Gilmer, Michael Muelly, Ian Goodfellow, Moritz Hardt, and Been Kim. Sanity checks for saliency maps. In Advances in Neural Information Processing Systems, 2018.
[19] Finale Doshi- Velez-Velez.nno and Been Kim. A roadmap for a rigorous science of interpretability. arXiv preprint arXiv: 772.008608 2017.
[20] James H Martin and Daniel Jurafsky. Speech and language prucessing: An introduction to natural language processing, computational linguistics, and speech recognition, volume 23. Pearson/Prentice Hall Upper Saddle River, 2009.
[21] Slava Katz. Estimation of probabilities from sparse data for the language model component of a speech recognizer. IEEE transactions on acoustics, speech, and signal processing, 35(3):400- -01, 1987.
[22] T. Brants, Ashok Popat, Peng Xu, Franz Josef Och, and Jeffrey Dean. Large language models in machine translation. In Conference Empirical Methods in Natural Language Processing, 2007. 2007.
11

---

[23] Miltiadis Allamanis and Charles Sutton. Mining source code repositories at massive scale using language modeling, 2013 10th Working Conference Mining Software Repositories (MSR), pages 207-216, 2013.
[24] Herman Stehouwer and Menno van Zaanen. Using suffix arrays as language models: Scaling the n-gram. 2010.
[25] Casey Redd Kennington, Martin Kay, and Annemarie Friedrich. Suffix trees as language models. In International Conference on Language Resources and Evaluation, 2012.
[26] Ehsan Shareghi, Matthias Petri, Gholamreza Haffari, and Trevor Cohn. Compact, efficient and unlimited capacity: Language modeling with compressed suffix trees. In Conference on Empirical Methods in Natural Language Processing, 2015.
[27] Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization; Nearest neighbor language models. iclr, 2020.
[28] Huayang Li, Deng Cai, Jin Xu, and Taro Watanabe. Residual learning of neural text generation with n-gram language model. In Findings of the Association for Computational Linguistics: EMNLP 2022, 2022.
[29] Sehastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George Bm Van Den Driessche, Jean- Bann-Bapisste Lespiau, Bogdan Damoc, Aidan Clark, et al. Improving language models by retrieving from trillions of tokens. In icml, 2022.
[30] Armand Joulin, Edouard Grave, Piotr Bojanowski, and Tomas Mikolov. Bag of tricks for efficient text classification, arXiv preprint (:Xivv:607.077555 2016.
[31] Bofang Li, Tao Liu, Zhe Zhao, Puwei Wang, and Xiaoyong Du Neural bag-of-ngrams. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 31, 2017.
[32] Chandan Singh, Armin Askari, Rich Caruana, and Jianfeng Gao. Augmenting interpretable models with large language models during training. Nature Communications, 14(1))7913, 2023.
[33] Yue Yang, Artemis Panagopoulou, Shenghao Zhou, Daniel Jin, Chris Callison Burch, and Mark Yatskar. Language in a bottle: Language model guided concept bottlenecks for interpretable image classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 19187- 19197, 2023.
[34] Chung-En En Sun, Tuomas Oikarinen, Berk Ustun, and Tsui-Wei Weng. Concept bottleneck large language models. arXiv preprint arXiv:2412.07992, 2024.
[35] Chandan Singh, John Morris, Alexander M Rush, Jianfeng Gao, and Yuntian Deng. Tree prompting: Effficient task adaptation without fine-tuning. In Proceedings of the 2023 Conference Empirical Methods in Natural Language Processing, pages 6253-6267, 2023.
[36] Jean Feng, Avni Kothari, Luke Zier, Chandan Singh, and Yan Shuo Tan. Bayesian concept bottleneck models with IIm priors. arXiv preprint arXiv:2410. 15555, 2024.
[37] Denis Jered MeInerny, Geoffrey Young, Jan- Anl-illem van de Meent, and Byron C Wallace. Chill: Zero-shot custom interpretable feature extraction from clinical notes with large language models. arXiv preprint arXiv:2302. 12343, 2023.
[38] Daking Rai, Yilun Zhou, Shi Feng, Abulhair Saparov, and Ziyu Yao. A practical review of mechanistic interpretahility for transformer-based language models. arXiv preprint arXiv:2407.02644. 2024.
[39] Kevin Meng, David Bau, Alex Andonian, and Yonatan Belinkov. Locating and editing factual associations in gpt. Advances in neural information processinn systems, 35, 2022.
[40] Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, et al. A mathematical framework for transformer circuits. Transformer Circuits Thread, 1, 2021.
[41] Naomi Saphra and Sarah Wiegreffe. Mechanistic? arXiv preprint arXiv:2410.09087 2024.
[42] Kevin Wang, Alexandre Variengien, Arthur Conmy, Buck Shlegeris, and Jacob Steinhardt. Interpretability in the wild: a circuit for indirect object identification in GPT-2 small. arXiv preprint arXiv:2211.000 1.00993, 2022.
[43] Shailee Jain and Alexander Huth. Incorporating context into language encoding models for fmrr. Advances in neural information processing systems, 31, 2018.
12

---

[44] Leila Wehbe, Ashish Vaswani, Kevin Knight, and Tom Mitchell. Aligning context-based statistical models of language with brain activity during reading. In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Procexsing (EMNLP), pages 233-243, Doha, Qatar, October 2014. Association for Computational Linguisties.
[45] Martin Schrimpf. Idan Asher Blank, Greta Tuckute, Carina Kauf, Eghbal Hosseini, Nancy Kan- wisher, Joshua B Tenenbaum, and Evelina Fedorenko. The neural architecture of language: Integra- tive modeling converges on predietive processing. Proceedings of the National Academy of Sciences, 118445)c21066411 18, 2021.
[46] Mariya Toneva and Leila Wehbe. Interpreting and improving natural natural-languaa language processing (in machines) with natural language- processing (in the brain). In H. Wallach, H. Larochelle, A. Beygelzimer, F. d'Alché- Bue, Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[47] Ariel Goldstein, Zaid Zada, Eliav Buchnik, Mariano Schain, Amy Price, Bobbi Aubrey, Samuel A. Nastase, Amir Feder, Dotan Emanuel, Alon Cohen, Aren Jansen, Harshvardhan Gazula, Gina Choe, Aditi Rao, Catherine Kim, Colton Casto, Lora Fanda, Wernrr Doyle, Daniel Friedman, Patricia Dugan, Lucia Melloni, Roi Reichart, Sasha Devore, Adeen Flinker, Liat Hasenfratz, Omer Levy, Avinatan Hassidim, Michael Brenner, Yossi Matias, Kenneth A. Norman, Orrin Devinsky, and Uri Hasson. Shared computational principles for language processing in humans and deep language models. Nature Neuroscience, 35:3)336- 380, March 2022. Number: Publisher: Nature Publishing Group.
[48] Michael C.-K. Wu, Stephen V. David, and Jack L. Gallant. Complete functional characterization of sensory neurons by system identification. Annual Review of Neuroscience, 29:477-505, 2006.
[49] Charlotte Caucheteux. Alexandre Gramfort, and Ban- King. Disentangling syntax and semantics in the brain with deep networks. In Proceedings of the 38th International Conference on Machine Learning, pages 1336-1348 1348. PMLR. July 2021. ISSN: 2640-3498.
[50] Carina Kauf, Greta Tuckute, Roger Levy. Jacob Andreas, and Evelina Fedorenko. Lexical-semantic content, not syntactic structure, is the main contributor to ann-brain similarity of fmri responses in the language network. Neurobiology of Language, 5(1):7 42, 2024.
[SI] Aniketh Janardhan Reddy and Leila Wehbe. Can fmri reveal the representation of syntactic structure in the brain? In Advances in neural information processing systems, volume 34, pages 9843-9856, 2021.
[52] Alexandre Pasquiou, Yair Lakretz, Bertrand Thirion, and Christophe Pallier. Information-restricte neural language models reveal different brain regions' sensitivity semantics, syntax, and context. Neurobiology of Language, 4(4):61 1 1-636, 2023.
[53] Khai Loong Aw and Mariya Toneva. Training language models for deeper understanding improves brain alignment. CORR, 2022.
[54] Sreejan Kumar, Theodore R. Sumers, Takateru Yamakoshi, Ariel Goldstein. Uri Hasson, Kenneth A. Norman, Thomas L. Griffiths, Robert D. Hawkins, and Samuel A. Nastase. Reconstructing the cascade of language processing in the brain using the internal computations of transformer-based language model. bioRxiv, June 2022.
[55] SubbaReddy Oota, Manish Gupta, and Mariya Toneva. Joint processing of linguistic properties in brains and language models. Advances in Neural Information Processing Systems. 18001-44148 2023.
[56] Shailee Jain, Vy Vo, Shivangi Mahto, Amanda LeBel, Javier S Turek, and Alexander Huth. Interpretable multi-timescale models for predicting fmri responses to continuous natural speech. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin. editors, Advances in Neural Information Processing Systems, volume 33, pages 13738 137499. Curran Associates, Inc., 2020.
[57] Catherine Chen, Tom Dupré la Tour, Jack L Gallant, Daniel Klein, and Fatma Deniz. The cortical representation of language timescales is shared between reading and listening. Communications Biology, 7(1):284, 2024.
[58] Greta Tuckute, Aalok Sathe, Shashank Srikant, Maya Taliaferro, Mingye Wang, Martin Schrimpf, Kendrick Kay, and Evelina Fedorenko. Driving and suppressing the human language network using large language models. Nature Human Behaviour, 8(3):544-56) 2024.
[59] Vinamra Benara, Chandan Singh, John X Morris, Richard Antonello, lon Stoica, Alexander G Huth, and Jianfeng Gao. Crafting interpretable embeddings for language neuroscience by asking Ilms questions. Advances in neural information processing systems, 37:124137, 2024.
13

---

[60] Chandan Singh, Richard J Antonello, Sihang Guo, Gavin Mischler, Jianfeng Gao, Nima Mesgarani, and Alexander G Huth. Evaluating scientific theories as predictive models in language neuroscience. bioRxiv, pages 2025-08, 2025.
[61] Richard Antonello, Chandan Singh, Shailee Jain, Aliyah Hsu, Jianfeng Gao, Bin Yu, and Alexander Huth. A generative framework hridge data-driven models and scientific theories in language neuroscience, 2024.
[62] Alan Baddeley. Working memory. Science, 255(5044):556-559, 1992.
[63] Ovid JL Tzeng. Positive recency effect in delayed free recall, Journal of Verbal Learning and Verbal Behavior, 12(4))436-439, 1973.
[64] Jeanne T Amlund, Carol Anne M Kardash, and Raymond W Kulhavy. Repetitive reading and recall of expository text. Reading Research Quarterly, pages 49-58, 1986.
[65] TR Miles, Guillaume Thierry, Judith Roberts, and Josie Schiffeldrin. Verbatim and gist recall of sentences by dyslexic and non-dyslexic adults. Dyslexia, 12(3):177 11944177-19, 2006.
[66] Aditya R Vaidya, Javier Turek, and Alexander G Huth. Humans and language models diverge when predicting repeating text. arXiv preprint 23X.xx3310006000 2023.
[67] Mathis Pink, Vy A Vo, Qinyuan Wu, Jianing Mu, Javier S Turek, Uri Hasson, Kenneth A Norman, Sebastian Michelmann, Alexander Huth, and Mariya Toneva. Assessing episodic memory in ilms with sequence order recall tasks. arXiv preprint 2oXiv:2400.08153 2024.
[68] Alex Warstadt, Aaron Mueller, Leshem Choshen, Ethan Wilcox, Chengxu Zhuang, Juan Ciro, Rafael Mosquera, Bhargavi Paranjabe, Adina Williams, Tal Linzen, and Ryan Cotterell. Findings of the BabyLM challenge: Sample-efficien pretraining on developmentally plausible corpora. Proceedings of the BabyLM Challenge at the 27th Conference on Computational Natural Language Learning, pages 1 14, Singapore, December 2023. Association for Computational Linguisties.
[69] Leo Gao, Stella Biderman, Sid Black, Laurence Golding, Travis Hoppe, Charles Foster, Jason Phang, Horace He, Anish Thite, Noa Nabeshima, et al. The pile: An 800gb dataset of diverse text for language modeling, arXiv preprint arXiv:2101.0000 000077 2020.
[70] Guilherme Penedo, Hynek Kydlíčck, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. The fineweb datasets: Decanting the web for the finest text data at scale. Advances in Neural Information Processing Systems, 37:30811-30849 2024.
[71] Alexander G Huth, Wendy A De Heer, Thomas L Griffiths, Frédéric E Theunissen, and Jack L Gallant. Natural speech reveals the semantic maps that tile human cerebral cortex. Nature, 532(7600):443- 4588, 2016.
[72] Amanda LeBel, Lauren Wagner, Shailee Jain, Aneesh Adhikari-Desai, Bhavin Gupta, Allyson Morgenthal, Jerry Tang, Lixiang Xu, and Alexander G Huth. natural language fmri dataset for voxelwise encoding models. Scientific Data, 10(1):555, 2023.
[73] Jerry Tang, Amanda LeBel, Shailee Jain, and Alexander G Huth. Semantic reconstruction of continuous language from non-invasive brain recordings. Nature Neuroscience, pages 858-886, 2023.
[74] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shrut Bhosale, et al. Llama Open foundation and fine-tuned chat models. arXiv preprint Skdi:2307.092880 2023.
[75] Richard Autonello, Aditya Vaidya, and Alexander Huth. Scaling laws for language encoding models in fmri. Advances in Neural Information Processing Systems, 36, 2023.
[76] Ruiqi Zhong, Charlic Snell, Dan Klein, and Jacob Steinhardt. Describing differences between text distributions with natural language. In International Conference on Machine Learning, pages 27099- 27116. PMLR, 2022.
[77] Lisa Dunlap, Yuhui Zhang, Xiaohan Wang, Ruiqi Zhong, Trevor Darrell, Jacoh Steinhardt, Joseph Gonzalez, and Serena Yeung-Levy. Describing differences in image sets with natural language. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 24199- 24208, 2024.
[78] Zonglin Yang, Xinya Du, Junxian Li, Jie Zheng, Soujanya Poria, and Erik Cambria. Large language models for automated open-domain scientific hypotheses discovery. arXiv preprint 02272 2023.
14

---

[79] Luis Fernando Nicolas-Alonso and Jaime Gomez-Gil. Brain computer interfaces, review. sensors. 12(2):1211-1279. 2012.
[80] Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu, Retrieval head mechanistically explains long-context factuality, 2024.
[81] Shangyi Geng, Wenting Zhao, and Alexander M Rush. Great memory, shallow reasoning: Limits of ls nn-Ims, arXiv preprint arXiv:2408. 1815, 2024.
[82] Joshua Engels, Isaac Liao, Eric J Michaud, Wes Gurnee, and Max Tegmark. Not all language model features are linear. arXiv preprint arXiv:2405. 14860, 2024.
[83] Vivien Cabannes, Charles Amal, Wassim Bouaziz, Xingyu Yang, Francois Charton, and Julia Kempe. Iteration head: mechanistic study of chain-of-thought. Advances in Neural Information Processing Systems, 1091111199122 2024.
[84] Sheridan Feucht, Eric Todd, Byron Wallace, and David Bau, The dual-route model of induction. arXiv preprint arXivv2504.00022, 2025.
[85] Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao, and Tuo Zhao, Tell your model where to attend: Post-hoc attention steering for L.LMs. arXiv preprint arXiv:231 1.02662, 2023.
[86] Steven Bills, Nick Cammarata, Dan Mossing, Henk Tillman, Leo Gao, Gabriel Goh, Ilya Sutskever, Jan Leike, Jeff Wu, and William Saunders. Language models can explain neurons in language models, 2023.
[87] Chandan Singh, Aliyah R Hsu, Richard Antonello, Shailee Jain, Alexander G Huth, Bin Yu, and Jianfeng Gao. Explaining black box text modules in natural language with language models. arXiv preprint arXiv:2305.0988 098633 2023.
[88] Trenton Bricken, Adly Templeton, Joshua Batson, Brian Chen, Adam Jermyn, Tom Conerly, Nick Turner, Cem Anil, Carson Denison, Amanda Askell, Robert Lasenby, Yifan Wu, Shauna Kravec, Nicholas Schiefer, Tim Maxwell, Nicholas Joseph, Zac Hatfield-Dodds, Alex Tamkin, Karina Nguyen, Brayden McLean, Josiah Burke, Tristan Hume, Shan Carter, Tom Henighan, and Christopher Olah. Towards monosemanticity; Decomposing language models with dictionary learning. Transformer Circuits Thread, 2023. https:/ttansfoomeerrrrciii http::/rrrrrriiiiuuuuu////////e pwtp://22//ooooommmtttto http://ooooommnnteeeetteeeeeeeee features/index..tnl)
[89] Pawel Renc. Yugang Jia, Anthony E Samir, Jaroslaw Was, Quanzheng Li, David W Bates, and Arkadiusz Sitek. Zero shot health trajectory prediction using transformer. NPJ digital medicine, 7(1):256, 2024.
[90] Charlotte Cancheteux, Alexandre Gramfort, and Jean- Rémi King. Evidence of predictive coding hierarchy in the human brain listening to speech. Nature human behaviour, 733)4400-41,, 2023.
[91] Riki Shimizu, Richard Antonello, Chandan Singh, and Nima Mesgarani. Interpretable embeddings of speech enhance and explain brain encoding performance audio models. arXiv preprint arXiv:2507.16000 2025.
[92] Žiga Avsec. Vikram Agarwal, Daniel Visentin, Joseph R Ledsam, Agnieszka Barnkkaaaar Kyle Taylor, Yannis Assacl, John Jumper, Pushmeet Kohli, and David R Kelley. Effective gene expression prediction from sequence by integrating long-range interactions. Nature methods, 2021.
[93] Bryan Lim, Sercan Ô Ank, Nicolas Loeff, and Tomas Pfister. Temporal fusion transformers for interpretable horizohorizon time series forecasting. International Journal of Forecasting, 37(4))1748 (744-174481774, 2021.
[94] Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani. Self-attenioon with relative position representations. arXiv preprint 1rXiv:1803.02/55 2018
[95] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations, 2019.
[96] Peter Henderson, Mark Krass, Lucia Zheng, Neel Guha, Christopher D Manning. Dan Jurafsky, and Daniel Ho. Pile of law: Learning responsible data filtering from the law and a 256gb open-soures legal dataset. Advances in Neural Information Processing Systems, 35:29217-29233 2022.
197] Shinji Nishimoto, Alexander G Huth. Natalia Bilenko, and Jack L Gallant. Eye movement-invariant representations in the human visual system. Journal of vision, 17(1):11-1,, 2017.
198] Bruce Fischl. Freesurfer. Neuroimage, 6202):77-788,,011
15

---

A Appendix
A.1 Impact Statement
We introduce the Generalized Induction-head Model (GIM), which improves the performance of fully interpretable models while maintaining transparency, making interpretable models more viable for high-stakes applications. For example, in medical note generation, interpretable models can enhance transparency, enabling clinicians to audit Al-generated text and reduce the risk of hallucinations or biased outputs. Additionally, GIM's token-level grounding improve faireess in language models and mitigate bias in automated seriiionn-making GIM also achieved significant speedups in speculative decoding compared to inference with LLaMA2-70B alone, making it suitable for deployment in compute-limited settings (see Appendix A.5). Despite these advantages, GIM does not fully close the performance gap with black-box models, particularly for tasks requiring extensive reasoning or broad world knowledge. Its reliance on input context may also limit effectiveness in some scenarios where high-quality data is not available.
A.2 Training of Fuzzy Matching Model
Architecture of Fuzzy Matching Model We train two Fuzzy Matching Models, one using the GPT-2 tokenizer and the other using the LLaMA-2 tokenizer. With GPT-2 tokenizer, Fuzzy Matching Model consists of four transformer layers, whereas it comprises three transformer layers when using LLaMA-2 tokenizer. Since relative position is crucial for calculating similarity, we incorporate Relative Positional Encoding [94], with a maximum relative position of 32 for the GPT-2 tokenizer and 64 for the LLaMA-2 tokenizer. The vocabulary embeddings are initialized with those from GPT-2 and LLaMA2-7B. ensuring that the number of heads and embedding dimensions align with the specifications of GPT-2 and LLaMA2-7B.
Creating similarity pair with LLMs For both Fuzzy Matching Model, we use LLaMA2-7B as teacher model. OpenWebTex. and Pile-train datasets for training each Fuzzy Matching Model that use GPT-2 or LLaMA-2 tokenizer. During training, we randomly sample sequences of 32 or 64 tokens with batch size of 128 or 256, resulting in 4,096 or 16,384 next-token prediction probabilities per batch. From these, we sample distant 3,584 or 4,096 queries and 512 keys and create similarity pairs (3, 584 x 512 or 4, 096 x 512) by calculating similarity based on Equation (5). The models trained using a combination of CE loss and reverse KL.D loss, with equal weights (1-0). We adopt most of the training settings from the 1 10 for training. Gradients are accumulated over 16 iterations, and we use the Adam W optimizer [95] with a learning rate of 0.0001 and a weight decay of 1. The learning rate follows cosine schedule with warmup over the first 1,000 iterations, and training continues for 15,000 or 20,000 iterations. Training is conducted on four NVIDIA A100 GPUs.
Ablation study on Fuzzy Matching Model training We conduct an ablation study on the positional encoding strategy and training process of Fuzzy Matching Model using the OpenWebTex WebText dataset to distill it from L-AMA-2-7B. The study evaluates the contributions of Relative Positional Encoding, reverse KLD loss, and CE loss to the model's effectiveness. As shown in Table A 1,, next-token prediction accuracy improves significantly when both reverse KLD and CE losses are included, demonstrating their complementary roles in optimizing the Fuzzy Matching Model. With CE loss, Forward KLD loss is less effective than reverse KLD loss. Furthermore, using Relative Positional Encoding instead of Sinusoidal Positional Encoding leads to better performance, highlighting the advantages of incorporating relative positional information for enhanced fuzzy matching capabilities.
We also perform an ablation study on the training data. Table A2 shows that the performance difference between Open WebText and Pile-train is minimal. When trained on Pile-of-Law [96], more domain-specific corpus, Fuzzy Matching Model exhibits slightly lower performance. This suggests that domain specificity may slightly limit the generalization ability of the fuzzy matching module. Nevertheless, the approach remains robust even with more Nomain-speciin training data.
"nttpa://nugginngaacs http://ggg..gg.c.co///ccaa/ccnnn htt.:/oow.oooooooooeeooooorraaaa uncopyrighted
10 hutppe:/ htstpe///iithbb con/karpstay/nindf?
16

---

Table A1: Ablation study on training of Fuzzy Matching Model. Next-token accuracy (%) of GIM (fuzzy) on the BabyLM-tee is reported. LLaMA-2 tokenizer is used.
Positional Encoding
Reverse KLD loss
Forward KLD loss
CE loss
Accuracy
Relative
43,2
Relative
42.8
Relative
Relative
41.9
Sinusoidal
37.0
Table A2: Ablation study training Fuzzy Matching Model with different datasets. Next-token accuracy (%) of GIM (fuzzy) on the BabyLM-rest is reported. LLaMA-2 tokenizer is used.
Dataset
Accuracy
OpenWehText
43.2
Pile-train
42.7
Pile-of-law
41.8
A.3 Determination o
To build GIM by integrating the three types of estimations, we first need to determine the threshold for effective n., denoted т. To identify the optimal value of we conducted cross-validation using the BabyLM training set (100M tokens). BabyLM consists of six datasets: open_subtite subtitles, mnc_spoken gutenberg, childes, simple__ wiki, and switchboard. Since switchboard con- tains only 2M tokens, we exclude it from the experiment. For the remaining datasets, we use each dataset validation set, while the other four used as the reference corpus to build Infini-gra gram. We then compare the performance changes of -afmn GIM (exact), and GIM (fuzzy) depending on effective 10k samples are used for evaluating each dataset.
As shown in Figure AI, anfini-gram outperforms GIM (exact) when the effective n exceeds for the GPT-2 tokenizer and for the LLaMA-2 tokenizer. Therefore, we set r to 8 and 9 for the respective tokenizers.
Tokenizer: GPT-2
Tokenizer: Llama-2
18800
12100
12003
1000
Count
5200
6003
0
2000
10
12
1.4
1/
13
12
14
16
Effectlve
Effectiin
Figure A1: Comparison of next-token accuracy.
A.4 Language Modeling Results Extended
Experimental details We use diverse datasets as reference corpus for Infini-gram. We use Infini- gram that is released by authors! "I for Pile-traint and 1 For Maby and Open WebText,
hhstpa///nffin http://aatitngaa.ntitnn wrw......app..o......t
14 _04_pieerriin__ 1lama
44 "v__pileval, 11ama and v4 14.piletrain_gt22
17

---

Table A3: Ablation study on components of Infini-gram with GIM. Next-token accuracy (%) on BabyLM is reported.
Reference Corpus
BabyLM-dev
Pile-val
OpenWebText
Pilc-train
Infini gram with GIM
43.1
42.9
43.2
49.4
w/o GIM (fuzzy)
42,2
36.9
38.3
46.6
w/o GIM (exact)
43.0
42.8
43.1
49,3
w/o Infini-gram
42.9
Infini-gram
39.0
19.0
20.1
33.5
Table A4: Speed of speculative decoding (SP). Accept. denotes the acceptance rate (%). The mean and standard deviation of runs are reported.
BabyL.M-test
Pile-val
Draft Model
Large Model
SP
Accept. rate (%)
Speed
Accept. rate (%)
Speed
ms/token (L)
Up (t)
ms/token (L)
Up (1)
LLaMA2-7B
30.2+0.0
30.2=0.1
TmnLLLMA-IIIIB
LLaMA2-7B
78.7
21.3+0.0
1.42
78.340.1 1
21.3=0.6
1.42
A40X1
GIM (fuzzy)
LLaMA2-7B
74.9+1.1
17.7+0.7
1.71
71.2+0.5
20.1=0.4
1.50
L1AMA2-13B
52.4+0.0
52.0=0.2
TInyLLaMA-I.IB
LLaMA2-13B
78.2+0.0
26.7+0.5
1.96
77.620.1
26.3=0.5
1.98
GIM (fuzzyy
LLaMA2-13B
73.5+0.1
24.8+0.1
2.11
69.8+0.2
27.8=0.1
1.87
LLAMA2-13B
26.4+0.1
26.3+0.4
LLaMA2-7B
LLaMA2-13B
78.9=0.0
24.7+0.0
1.07
78.6=0.0
25.1=0.3
ToyLLaMA-LIIB
LLaMA2-13B
78.340.1
20.740.1
1.28
77.640.1
21.5+0.1
H100X0
GIM (fuzzyy
LLaMA2-13B
73.2+0.3
13.3+0.2
1.98
69.9=0.1
14.9=0.1
1.77
LLaMA2-70B
71.2+0.1
71.020.2
LLaMA2-7B
LLaMA2-70B
77.240.2
38.3+0.5
1.86
77.8+0.2
37.4+0.3
1.90
TinyLLaMA-L.IB
LLaMA2-70B
75.5+0.1
35.340.2
2.02
76.3+0.4
33.9=0.6
2.10
GIM (fuzzy)
LLaMA2-70B
68.520.6
31.4+0.7
2.27
66.620.6
33,340.6
2.13
we build our own Infini-gram. We use public code to build and inference Infini-gram and GIM (exact)15, During inference, the maximum length for exact matching with Infini-gram is 500, and use window size k for fuzzy matching as 32 and 64 for GPT-2 and LLaMA-2 tokenizers, respectively.
Ablation study on Infini-gram with GIM We conduct an ablation study to assess the impact of each component in Infini-rram with GIM. Table A3 reports next-token accuracy when individual components are omitted. Excluding GIM (fuzzy) results in more significant performance drop than removing GIM (exact). This underscores the importance of fuzzy matching in handling diverse contexts and improving adaptability, as reflected in Table 1, where GIM (fuzzy) outperforms GIM (exact). Since both components act as induction heads, they exhibit complementary roles-when one is removed, the other partially compensates for its absence. Only when using Pile-train reference corpus, omitting Infini- Infini-gram gram leads to the most substantial performance decline. It is worth noting that when the reference corpus lacks similarity to the test dataset's 's distribution (e.g.. Pile-val, Open WebText, and Pile-train), the performance of thiii-gramm falls significantly below the scenario where it is not utilized at all. This highlights the sensitivity of Infini-gran to the quality and relevance of the reference corpus.
A.5 Speculative Decoding
GIM offers both interpretability and efficiency. When combined with LLMs in speculative decoding, enhances prediction accuracy while significantly boosting inference speed.
14 taattpa://intiii gran. gran.1o/pkk....... ht/pk//.pp...oooo
attpp /ttp//gtttbtt..c/eee wat.cco///eeeeee Con/AleeVanO/ietee geaaaagrennnae tren/tree/niin
18

---

Table A5: Speed of speculative decoding (SP). reported.
The mean and standard deviation of runs are
MMbyeM
Pile-val
Fine Web
Draft Model
Large Model
ms/token (1)
Speed (t)
ms/token (J)
Speed Up (1)
ms/token (1)
Speed Up (t)
I.a.MA2-13
26.4s0.1
26.3:0.4
GIM (fuzzyy
LLaMA2-13H
13.3+0.2
1.98
14.9+0.1
1.77
14.9=0.3
1.76
GIM
LLaMA2-13B
23.1
1.14
22.8+0.3
1.15
23.0z0.7
1.14
LLaMA2-70B
71.2#0.1
7100+022
71.1=0.2
GIM (fuzzy)
LLaMA2-70B
31.4+0.7
2.27
33.340.6
2.13
33.2=1.0
2.15
GIM
LLaMA2-70B
42.0+0.7
1.70
41.6+10
1.71
40.421.2
1.76
Experimental details To evaluate the efficiency of GIM (fuzzy). we compare the inference time for speculative decoding with TinyLLaMA 16 and LLaMA2-7B [74]. We evaluate speculative decoding by generating up 1024 tokens, using prefix of 1024 tokens. The speed of decoding may vary depending on the computational environment. To ensure robust evaluation across different setups, we conduct experiments in two environments: one with a single NVIDIA A40 GPU and 128 CPU cores, and another with two NVIDIA H100 GPUs and 64 CPU cores. Greedy sampling is used for token generation, and each experiment is repeated three times with different random seeds.
Induction improves speculative decoding performance Table A4 demonstrates the speed-up effect of speculative decoding with GIM (fuzzy). GIM (fuzzy) relies solely on the induction power derived from the input context to predict the next token, leading to lower acceptance rates compared LLMs. Despite this, its inference speed is remarkably fast, and it often matches the predictions of large models. As a result, the speed improvement can exceed 2x compared to using LLaMMA-70BB alone. most cases, GIM (fuzzy) achieves even greater speed gains than when using an LLM as draft model for speculative decoding.
Additionally, would like note that speculative decoding with GIM (fuzzy) and pretrained LLM not only accelerates the inference speed of the pretrained model but also enables explainable predictions based the given input context. When accurate predictions more challenging cases, we rely
be made through
interpretable methods, we utilize this process for interpretability. on larger model that, while less interpretable, delivers better performance for accurate predictions. Thus, this approach provides a balanced method that addresses both interpretability and accuracy, in addition to enhancing efficiency.
Table AS reports the inference times for GIM (fuzzy) and GIM using speculative decoding, with the Text dataset serving the reference corpus for Infini-gram. We find matches with a maximum of 64 tokens for both exact and fuzzy matching. The experiments are conducted on two NVIDIA H100 GPUs and 64 CPU cores. Although GIM requires time for generation average than GIM (fuzzy), still significantly reduces inference time compared to relying solely on a large model for inference.
Explanation Figure A2 presents several examples of explanations provided by GIM. Even an exact match fails to yield good match, when the probability of subsequent tokens is similar, the fuzzy matching model can predict with high similarity, enabling successful fuzzy matching, enabling successful fuzzy matching, and improving next-token prediction.
A.6 IMRI results extended
Data details We analyze publicly available data collected from prior work [72, 73]. Methods from the previous study are summarized here for completeness. Functional magnetic resonance imaging (IMRI) data was recorded from three healthy participants as they listened to Anglish-language podcast stories over Sensimetrics S14 headphones. Participants were only instructed to listen to the stories. No explicit behavioral responses were required.
For the collection of training data, each participant completed approximately 20 hours of listening sessions across 20 separate sessions in which unique stories were presented. This produced 33,000
16 "aatpe://hugggnngfse http:iiinnnceeeeeetnreeeeeeeeeee hwwpa//iwwwiiiiiaaaaataaaa htt.aaaaaaaaa---aaaaeeeeeaeeeeo. Laterrediate-stee--655.-
19

---

GIM (exact)
(8) Input Prompt: Frontispicce, ceeeeeeee______ Paae 61_]] ANBUNHY BROWN AND HIS
"BR"
Sequence Context
Effective
PG70358 ADBUNNY BROWN AND H15 SIST"
BR
(b) Input Prompt: Then the chorus: "will you, won you, will you, von
*
Sequence from Context
Effective
Next Taken t
out in friendly voice:\n'Wil Wn'Will you, won "nn't you, will you, won"
13
GIM (fuzzy)
"it" (c) Input Prompt: .... Because he says it's Lincolnshire \nÑo, he didn't 'tl Antte said"
Sequence from Context
Similarity
Next Token
What ' Lincoinshire gotta do with it? \nBecauae he aaya"
0.680
1t
God that wind's s gone cold: \nI say"
0.210
that
Hell he don't know anything about gardening, you see! AnBut"
0.203
I
What 8 Lincolnshire gotta with it? \nBecause*
0.186
he
know why"
(d) Input Prompt: .... So taught
first week, second"
"week" Next Token
running Sequence from Context
Similarity
sixty odd the first"
week day
in every"
And running it and the first week
there, and one"
0.053
gu
so I taught him that the first"
0.042
week
we had to cancel because nabady turned up. \när one"
0.555
of
Figure A2: Examples of explanation of GIM from BabyLM-test. (a) and (b) show examples of exact matching while (c) and (d) show examples of fuzzy matching. The red box marks the source of the final prediction.
timepoints per voxel across the entire human cortex. For testing data collection, participants heard two-held-out stories five times each and third story ten times (one story per session). These repeated measurements were averaged to improve reliability. Signal-to-noise ratios for each voxel estimated using the mean-explainable variance approach from [97]. Analysis was restricted to voxels that were located within 8mm of the cortical mid-surface, yielding about 90,000 voxels per participant.
participants were healthy adults with normal hearing and gave written informed consent. The study protocol was approved by the Institutional Review Board of the University of Texas at Austin. The scans were acquired on a T Siemens Skyra MRI system at the University of Texas at Austin using 64-channel Siemens head coil. Functional images were obtained with gradient-echo EPI sequence (TR = 2.0 s, TE = 30.8 ms, flip angle 71% multi-band factor = 2, voxel size = 2.6 mm x 2.6 mm x 2.6 mm, matrix size 84 84, field of view 220 mm). Anatomical scans were collected using a T1-weihhted multi-echo MP-RAGE sequence with mm isotropic voxels, following the standard Freesurfer morphometry protocol [98].
Functional data was preprocessed with FSL 5.0 using the FMRIB Linear Image Registration Tool (FLIRT) for motion correction and alignment. Fach participant's runs were registered to a subject- specific template built from the first functional run of the first session, with all automated registrations manually verified for accuracy. Low-frequency signal drifts were removed using second-order Savitzky-Gon Golay filter with a 120 s window. To mitigate onset artifacts and detrending issues near scan boundaries, 20 (10 volumes) were discarded from both the start and end of each run. This eliminated the initial silent period and the first and last 10 s of each story. Each voxel's I's mean response then subtracted, and the remaining signal was normalized to unit variance.
To ensure temporal alignment between linguistic and neural data, word onset times were interpolated to the fMRI sampling rate using Lanczos interpolation with window size of 3. The hemodynamic response was modeled as finite impulse response (FIR) with four time lags (-8, -6, -4, and s), following the approach of [71]. For every subject æ, voxel fit an encoding model g(x.v) to
20

---

Table A6: IMRI Prediction Performance when using fuzzy matching. Error bars show 95% CI.
Mean Correlation
Feature Model
Tokenizer
Matching Model
All Voxels
Top 10% Voxels
Eng1000
0.072 = 0.0004
0.220 0.0012
Infini-gram Eng1000
GPT-2
0.069 + 0.0003
0.200 0.0012
GIM Matching F0g0000
GPT-2
0.087 + 0.0005
0.265 0.0011
Fuzzy Induction Matching Eng1000
GPT-2
GPT-2
0.076 = 0.0004
0.222 + 0.0011
Fuzzy Induction Matching Eng 1000
LLaMA-2
LLaMA2-70B
0.076 + 0.0004
0.225 + 0.0012
Fuzzy Induction Matching Eng 1000
GPT-2
Fuzzy Matching Model
0.076 + 0.0004
0.216 0.0011
Fuzzy Induction Matching Eng1000
LLaMA-2
Fuzzy Matching Model
0.077 = 0.0004
0.223 + 0.0012
predict the BOLD response B from the embedded stimulus, i.e. Biz,e) = 9(x,v) (H, Model evaluation was performed on the held-out test stories, using the trained encoding models to predict and assess voxel responses.
IMRI fuzzy induction head settings Similar to the GIM Matching technique described in Sec. 5.1, we construct an induction head for fuzzy matching. the fuzzy setting, we leverage the predicted next-word distributions obtained through fuzzy n-gram matching in the input context innduction calculate the cosine similarity between the next-word distributions of the current word and all prior candidate words.
in Equation (3)), which refer to Fuzzy Induction Matching. Specifically,
To account for the temporal resolution of fMRI, similarity values, aligning these values with the fMRI time scale. This allows us to identify the time point (TR) t" that maximally corresponds to the current time point based on the highest similarity.
apply Lanczos smoothing to the word-level
We evaluate several configurations for deriving the next-word distributions, including GPT-2, LLaMa- 2, the Fuzzy Matching model with the GPT-2 tokenizer, and the Fuzzy Matching Model with the LLaMA-2 tokenizer. See more details on Fuzzy Matching models Sec. 3.2.
Extended prediction performance results The prediction performance of Fuzzy Induction Match- ing Models is compared the performance of the GIM Matching Models and the Eng1000 baseline in Table A6. The Fuzzy Induction Model, in its highest- sethhestperforming configuration (using the Fuzzy Matching Model with the LLaMa2-700 tokenizer), achieves only 6.94% improvement prediction performance compared to the Eng (000 baseline. The lower relative performance of Fuzzy Induction Matching compared GIM Matching may due to the inherent noise and lower spatial and temporal resolution of fMRT data, which makes it challenging to detect subtle differences in neural activations associated with similar but non-identical stimuli.
21

---

Table A7: GPT-4 Prompts for Generating and Classifying Categories of Text. Ellipses (...) indicate omitted portions of the full prompts.
Title
Prompt
GPT-4 Prompt for Generating Calegory scriptions
have provided two test stories below. Specific phrases from each story have been picked out based on the performance of different encoding models. Can you describe the characteristies of the words and phrases that each category contains? Be specific about type of words, their context the story, any other relevant commonalitics. Write succinet descriptions for each category that would allow one to categorize phrases in other such stories accurately.
Category A: ['sh first she digs into her cutoff's in the", 'both need this right now .... ]
Category B: ['to everything or you make yourself scarce', 'my cigarettes and uh', ...]
Full Story: [['i reached over and secretly'l, ], l'undid my seatbelt' ) ...]
GPT-4 Prompt for Classifying Stages Based Descriptions
have attached category descriptions below. Based on the de- scriptions, in order, go through each short list of words (short phrase) in the story at the end and classify the segments into one of the categories. Rather than listing all the phrases in a category at time, list each phrase in order and label it as helonging to category Category A: Emotionally, or Narratively Critical
Category B: Brief, Stand-Alone Phrases
Full Story: [['i reached over and secretly') J.. ['undid seatbelt') J.
22

---

Where's The Smoke Story' Segmented by Highest Performing Model
my
CORN
the
no
DECRAOR
POU
WITh
ver
ENG
PESSART
SOU
the
and
Sreat
place
there
VOLRS
after
and
STUTT
bock
192
SFOUE
NOTE
DETS
COSOF
the
PROOOTOR
theugh
TRUST
TURE
TOOO
RAVT
POTT FO
RED
1he
bACk
VAST
ON
TWE
STEP
DOCKETS
INN
DAcE
KnD
STD
TURKEY
DOVINO
Dom
DHOO
the
ISSE
SOUTE
cream
partor
that
thn
GRTS
CNTT
INC
Tew
pretly
trat
What
DUE
SouTre
85
yeArS
Take
kerse
WCB
why
DACK
valking
talentiann
light
and
hey
then
the
and
Test
Taling
the
Seepin
JACE
thn
undemeed the
COCK
the
Day
SOTTE
SOUTE
tome
pratyy
200
10
goOd
Eye
TOOKE
SYU
Tucky
of
g010
DC
ONST
500
know
SOKE
ITASEN
SOUTIA
LOCE
NYS
SCRTY
DUT
BEEK
105
TANO
END
Shend's
aryttingg
Legend Huidirg
Infrni Gram matching
No cigsfioett difiennce
Figure A3: Test story 1 (Where's 's There's 's Smoke), highlighted in regions where the Infini-Gram matching and GIM matching models exceed baseline performance, measured by the average absolute error across voxels, by more than one standard deviation.
23

---

Hovv Matl Story Segmented Highwst Performing Hocel
Thouood
djtt
poople
with
gradusted
Dut
and
10
after
compa
drow
30
hape
Shange
IVD
Fine
thy
2004
mat
good
Besgge
Tandin
the
somabedy
Senta
famiy
the
Pm
Jtice
bot
rebecca
aed
department
Terreland
security
doeett
15A
Choumatance
distrtinng
casobon
han
monkeer
the
Fesse
wth
only
get
2000
may
warstedd
Ue
the
Scanda
dart
but
STAROO
sta
people
reth
coder
eigenbation
clssceekted
ret
saying
statfer
Imagieed
Fm
MTM
word
the
ARR
hac
Imsparuund
yoU
to
NT
benn
RM
tharess
Sand
itat
huve
nork
better
compooouded
than
that
zeather
<bct
down
WOLLS
FAXT
BOXK
Tare
JaY
and
called
batty
Sumting
Tealy
yeur
and
Itc
32c28l
sfac
cffcrett
poople
b20ry
betisdy
these
secodd
5005
Aayy
die
Legend
Infin-Gremm matching
GIM Matching
No significant difference
Figure A4: The first section of test story 2 (Have You Met Him Yer), highlighted in regions where the Infini-Gram and GIM matching models exceed baseline performance.
24

---

317
Loyorl GM Paidden
Figure A5: The second section of test story (Have You Met Him Yer), highlighted in regions where the fffii--rrmann and GIM matching models exceed baseline performance.
25

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We reviewed the abstract and introduction to ensure they accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer:
[Yes]
Justification:
We discussed the limitations of work in the Discussion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
26

---

Justification:
Our paper does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We included all necessary information to reproduce the experimental results in the main text and the appendix.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
27

---

Answer: [Yes]
Justification: We included the URL to the publicly available code in the main text. Addition- ally, experimental details are demonstrated in the main text and the appendix.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
ceeioonooooioooo olycy) for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: / //nips. cc/ mublic/guidess Caleee//ooSSubbieee bebbmmisssiins ssisnnolicyyyyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We specified all details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer:
[Yes]
Justification:
We added error bars in plots or reported standard deviations tables.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
28

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We included the information on the computer resources in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https atttp://neurrr http///eerripp.peeee httpp///ccc./iiiicccttieeeeeeeee
Answer: [Yes]
Justification: We have read through the Code of Ethics and confirmed that our research conforms to them.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification; We discussed broader impacts in the Appendix and limitations in the Discussion section.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
29

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
We did not release any data or models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We included citations for all cited papers.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
30

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We did introduce any new assets in the paper.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not include such experiments or researches.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
This paper does not encompass such potential risks.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
31

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: LLMs were not employed for purposes integral to the central aspects of this research.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: : http://eeur...eeeee. cc ccc//oofeeeeccee/225///i for what should or should be described.
32


Reviews:
[{"review_id": "3JYrFDsU2N", "reviewer": "Reviewer_vuy2", "summary": "This 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.", "questions": "Have 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", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "This 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.", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful and encouraging comments. We are especially grateful for the recognition of (1) our model’s interpretability and (2) its utility in exploring the performance boundaries of fully interpretable approaches. Below, we address each of the reviewer’s points in detail.\n\n### W1. Raw performance gap vs. black-box LLMs & Hybrid directions\nWe agree that GIM, while substantially improving over prior interpretable baselines, does not yet match the performance of state-of-the-art black-box LLMs, and that bridging this gap via hybrid approaches is a promising future direction. In fact, we have taken a first step in this direction using speculative decoding, as described in Appendix A.5. Specifically, we employ GIM as the draft model and a black-box LLM (LLaMA-2-7B/13B/70B) as the large model. This setup leverages the interpretability of GIM for generating candidate tokens while relying on the large model for final validation. Because GIM’s predictions are fully decomposable into exact or fuzzy n-gram matches, every accepted token from the draft stage carries transparent attribution, thus maintaining interpretability even when coupled with a powerful black-box verifier. We view such GIM–LLM hybrids as a promising direction for not only faster inference but also improved balance between transparency and performance. We will revise the Discussion section to further highlight these trade-offs and future directions (please see our response to Reviewer WRxr, **W2: GIM’s scope and relation to LLM internals**).\n\n### Q1. GIM as a layer or attention head\nThank you for the suggestion. While our current implementation keeps GIM external to maintain full decomposability and interpretability, we agree that incorporating GIM as an internal component (e.g., as an attention head, similar to [1]) could open up interesting trade-offs. Our current design emphasizes transparency by ensuring that all probability mass originates from explicit matches. Embedding GIM within model internals might dilute this property, but it may also unlock new capabilities. We mention this as a promising direction for future work (Lines 301-302), particularly if such integration can retain partial attribution while improving flexibility.\n\nWe again thank the reviewer for their insightful and constructive feedback, and we hope our responses adequately address your concerns."}, {"review_id": "mvciLVrUnV", "reviewer": "Reviewer_WuPU", "summary": "This 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.", "questions": "1. 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.", "limitations": "Yes.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "While 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.", "was_revised": true, "final_justification": "This paper presents an innovative technique and has a very thorough evaluation. I believe it's a solid contribution to the field.", "rebuttal": "We thank the reviewer for their thoughtful and detailed feedback. We particularly appreciate the recognition of our paper’s contributions in three areas: (1) extending Infini-gram via in-context-only retrieval, (2) introducing fuzzy matching for improved robustness, and (3) validating the model across language and neuroscience domains.\n\n### W1. Limitations of in-context-only retrieval\nWe thank the reviewer for highlighting this consideration. While we agree that sparse or uninformative contexts can limit any model’s predictive ability, we do not view GIM’s reliance on input context as fundamentally at odds with generalization. In fact, many generalization behaviors in LLMs rely on identifying patterns in recent input [a, b]. GIM captures this same ability in an explicit and transparent way.\n\nGIM models this inductive behavior through both exact and fuzzy matching, enabling it to generalize across local variations even without access to large external corpora. As our results show, it often outperforms Infini-gram trained on massive datasets, demonstrating stronger generalization, especially in domains with repetitive or structured language. When richer or longer contexts are available, GIM benefits accordingly. It mirrors how black-box LLMs perform better when provided with more examples.\n\nImportantly, “input context” in GIM is not restricted to the literal text string passed at inference time. Its modular retrieval-based architecture allows it to seamlessly incorporate any additional information source, including external memory stores, retrieval-augmented generation (RAG) modules, or hybrid pipelines. For example, as we show in Appendix A.5, GIM can serve as a fast, interpretable draft generator in speculative decoding, with a larger model verifying predictions, demonstrating that accuracy and interpretability can be jointly achieved through composition.\n\nWe’ll clarify this perspective more directly in the revised discussion (please see our response to Reviewer WRxr, **W2: GIM’s scope and relation to LLM internals**), and hope this clarifies GIM’s design intent and extensibility across contexts.\n\n[a] Olsson, Catherine, et al. \"In-context learning and induction heads.\" arXiv preprint arXiv:2209.11895 (2022).\n\n[b] Brown, Tom, et al. \"Language models are few-shot learners.\" NeurIPS (2022).\n\n**Q1.** We chose JSD over other measures (such as KL divergence) primarily due to its symmetry. Unlike KLD, JSD is symmetric, which aligns with our use case: the fuzzy-matching model must learn the similarity between two n-gram next-token distributions that stand on equal footing, rather than comparing an “estimated” distribution to a “true” one. Because the two distributions are treated symmetrically, a symmetric metric like JSD is the appropriate choice.\n\n**Q2.** Equation (1) defines the teacher-side similarity using JSD over next-token distributions. However, our fuzzy matching model—a student model trained via knowledge distillation—does not need to generate full token distributions during inference. Instead, it produces embedding vectors for n-grams, and the similarity between them is computed via cosine similarity (Equation (2)). During training, the student is supervised to match the similarity scores given by the teacher (based on JSD), effectively learning to embed sequences such that their cosine similarity reflects the similarity of their predictive distributions. This formulation allows us to bypass full next-token prediction at inference time while still preserving the semantics of the JSD-based similarity.\n\n**Q3.** In Equation (4), we sum over all candidate spans $w_{j−k−1:j}$ in the input context and $w_j$ can be different from $w_i$. The indicator function $\\mathbb{1}_{w_j = w_i}$ ensures that we only include cases where the token following the matched n-gram ($w_j$) is equal to the specific token $w_i$ being scored. This allows us to accumulate a “floating count” of support for each token $w_i$ based on the similarity between preceding contexts.\n\nWe appreciate the reviewer’s attention to both the conceptual framing and technical details, and we hope our responses adequately address your concerns."}, {"review_id": "LLn4UL8bBX", "reviewer": "Reviewer_3WB2", "summary": "The 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.", "questions": "- 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?", "limitations": "The 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 =`", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "As mentioned in my rebuttal response, my primary concern with the current draft is its framing. Although the authors stated in the rebuttal that GIM represents a fundamentally different paradigm, focusing primarily on interpretability rather than performance, the current version of the manuscript suggests otherwise.", "rebuttal": "We thank the reviewer for their thoughtful and constructive review. We are especially grateful for the recognition of (1) extending mechanistic interpretability insights, (2) integrating Infini-gram with GIM to leverage context, and (3) validating the model across language and neuroscience domains.\n\n### W1. Role of GIM\nWe agree that tasks such as MMLU, code generation, and math require strong reasoning skills, which benefit from the broad parametric capacity of large black-box LLMs. However, GIM represents a fundamentally different paradigm: it is designed to maximize interpretability, not to replicate the full capabilities of LLMs or serve as a stand-alone replacement.\n\nOur contribution is to show that a large fraction of next-token prediction—particularly repetition-based induction—can be captured by a fully transparent mechanism, substantially narrowing the performance gap between interpretable and black-box models (see Tables 1 and 2). This marks a meaningful step for interpretable language modeling.\n\nAt the same time, GIM is explicitly composable to bridge transparency and general-purpose capability. As noted in Appendix A.5, we already demonstrate a hybrid approach where GIM acts as a transparent draft generator and a larger LLM serves as a verifier via speculative decoding. This design preserves interpretable predictions for accepted tokens while retaining the strengths of large models. We also refer the reviewer to our response to Reviewer WRxr, **W2: GIM’s scope and relation to LLM internals**, where we further clarify this distinction and outline our revision plan.\n\n### W2. Comparison to attribution-based interpretability\n\nWe appreciate the reviewer’s suggestion to compare GIM with attribution-based methods on pretrained models. While both GIM and attribution techniques may highlight context tokens that influence a prediction, they differ fundamentally in motivation, mechanism, and reliability.\n\nPost-hoc attribution methods, such as gradients or relevance propagation, estimate token importance after a parametric model makes a prediction. They rely on heuristic approximations of causal influence and are sensitive to implementation choices, often failing faithfulness tests [a]. This issue is particularly prominent in NLP, where interactions between many tokens make it difficult for attribution methods to reliably interpret the importance of individual tokens.\n\nIn contrast, GIM is interpretable by design. It does not approximate influence but generates predictions through explicit retrieval of matched n-grams. Every unit of probability mass is directly traceable to specific spans in the input. This makes GIM not just a model with interpretable outputs, but a model with a transparent and causal prediction mechanism.\n\nWe agree that comparing the two approaches is meaningful, as GIM is designed to model a key mechanism in LLMs. Such a comparison helps assess whether GIM’s transparent behavior reflects how LLMs handle contextual repetition.\n\nWe compared GIM’s matched sequences with LXT [b] attributions from LLaMA-2-7B. For LXT, we display tokens with attribution scores ≥ 0.10 or in the top 10, annotated as _`WORD[ATTRIBUTION VALUE]`_. Non-italic tokens have omitted attribution values. For GIM (exact), we underline the matched n-gram; for GIM (fuzzy), we report the top-3 highest similarity sequences with their scores. Note that GIM provides explanations at the sequence level, not token-by-token. Below are two examples:\n\n### Example 1)\n**Query:** `\"... I said that they don’t hang\"`  -> **Prediction:** `\"about\"`\n- **LXT:** _`\"
[0.33] … \\n[0.59]`_`… she don’t`_`hang[0.10] about[0.36]`_`… she hangs`_`about[0.10]`_`… we don’t`_`hang[0.17] about[0.36]`_`… Billy`_`’[0.10]`_`s got …`_`they[0.18] don[0.14]’[0.10]t[0.20] hang[1.00]\"`_\n- **GIM (exact):**  The last sequence _`\"don’t hang\"`_ matched with `\"she`_`don't hang`_`about\"` and `\"we`_`don't hang`_`about\"`\n- **GIM (fuzzy):**  \n  - _`\"… saying that we don't hang [0.32]`_`about\"`\n  - _`\"… But she don't hang [0.30]`_`about\"`\n  - _`\"… And then, hang [0.21]`_`on\"`\n\n### Example 2)\n**Query:** `\"... They use it as an\"`  -> **Prediction:** `\"exc\"`\n- **LXT:** _`\"
[0.33]`_`… they’ve got`_`.[0.37]`_`… They’ll use it as`_`an[0.08] exc[0.22] use[0.21]. … They[0.08]`_`use it`_`as[0.26] an[1.0]\"`_\n- **GIM (exact):** The last sequence _`\"use it as an\"`_ matched with `\"They'll`_`use it as an`_`exc\"`\n- **GIM (fuzzy):**  \n  - _`\"… They'll use it as an [0.73]`_`exc\"`\n  - _`\"… Cos if you left that with the [0.01]`_`first\"`\n  - _`\"… point out to them of course that the [0.01]`_`second\"`\n\nAcross examples, both GIM and LXT tend to focus on repeated patterns from earlier in the context—consistent with known induction head behavior (Lines 109–118). In this sense, attribution maps partially recover the same inductive patterns in LLMs that GIM is explicitly designed to model. However, key differences remain: attribution maps may assign high importance to non-semantic tokens (e.g., `\\n`, `.`), and they require interpretation to infer causality. In contrast, GIM uses its retrieved sequences to construct next-token distributions, executing a faithful, copy-like mechanism.\n\nWe recognize the reviewer’s concern that similar explanations might be achievable through interpretability techniques on pretrained models, but these remain approximate and offer no guarantees of faithfulness. GIM, by contrast, constructs predictions directly from retrieved spans, making attribution causal, controllable, and repeatable. This is not just a theoretical distinction, it is essential in domains like clinical decision-making, scientific modeling, or regulatory auditing, where interpretability must be built into the model, not added afterward (see Introduction, Lines 15–19).\n\nWe will clarify this point in the revised Introduction (Lines 28–30), emphasizing that GIM is not a tool to interpret opaque systems, but an inherently interpretable system. We see these two approaches as complementary, not interchangeable.\n\n[a] Zhao. et.al. “Explainability for Large Language Models: A Survey,” ACM TIST (2024).\n\n[b] Achtibat. et.al. “AttnLRP: Attention-Aware Layer-Wise Relevance Propagation for Transformers,” ICML (2024).\n\n[c] Rudin. “Stop explaining black box machine learning models for high stakes decisions and use interpretable models instead,” Nature Machine Intelligence (2019).\n\n### W3. Clarity improvements\nWe respond to each question below.\n\n**Q1.** As described in Appendix A.2, the Fuzzy Matching Model contains 3 Transformer encoder layers (LLaMA-2 tokenizer) or 4 (GPT-2 tokenizer). It outputs embedding vectors from the final layer without projection. These embeddings are used for cosine similarity (Eq. 2).\n\n**Q2.** The notation $w_{:i−1}$ refers to the full input up to but excluding token $w_i$, which denotes the final token. Thus, $w_{:i−1}$ forms the query context for GIM matching.\n\n**Q3.** The discrepancy in token counts stems from tokenizer differences. OpenWebText contains ~10B tokens with LLaMA-2, and ~9.04B with GPT-2. We will clarify this in Lines 180–181.\n\n**Q4.** This refers to the evaluation on Pile-val with the LLaMA-2 tokenizer: GIM (exact) 32.6 vs. Infini-gram 27.1 (Δ = 5.5%p).\n\n**Q5.** The “Type” column refers to model type (e.g., GIM (exact), GIM (fuzzy)), not to the tokenizer. The table is divided into two blocks: the upper block uses the GPT-2 tokenizer, and the lower block uses the LLaMA-2 tokenizer. The table is split into two blocks: GPT-2 (upper) and LLaMA-2 (lower). Since tokenizers segment text differently, the same corpus yields different token counts (e.g., OpenWebText = 9.04B with GPT-2 vs. 10.3B with LLaMA-2), which can impact performance.\n\n**Q6.** The “next-token distribution” refers to the probability distribution over next tokens computed by GIM based on the story text the subject has heard up to each time point (TR). In the fMRI setting, we use this distribution as a proxy for internal model state: for a given TR t, we find a prior TR t* whose GIM distribution is most similar (by cosine similarity), and retrieve the change in fMRI response that followed hearing the text at t*. This allows us to link similar linguistic contexts, according to GIM, to similar brain responses. We have clarified this in the main text (Lines 238–239).\n\n**Q7.** The response variable refers to the fMRI recordings themselves: specifically, the BOLD signal measured across 95,556 voxels at each 2-second time repetition (TR). To reduce dimensionality and extract dominant patterns of neural activity, we compute principal components over the voxel responses across time. The changes in these PCs following matched linguistic contexts are then used to construct GIM-based features, capturing how the brain’s response evolves in response to similar text input. We will clarify this in the revised paper.\n\n**Q8.** GIM explicitly models induction heads—key mechanisms in LLMs for in-context learning—as interpretable modules. We believe more such modules can be designed for other known behaviors. We envision building a more general and transparent language model by composing such modules. For example, future work could incorporate circuits for indirect object identification as standalone modules, or introduce retrieval heads that explicitly fetch external knowledge when needed. These modules could be implemented using interpretable retrieval schemes or rule-based operations similar to GIM, with each module contributing a distinct, traceable component to the next-token distribution. As interpretability research progresses, GIM-style components may serve as the building blocks of transparent yet capable models.\n\n### L1. Symbolic reasoning limitations\nWe agree. Tasks requiring symbolic reasoning or variable binding are beyond GIM’s current scope. It lacks internal state tracking or arithmetic capacity. We discuss this limitation in the Discussion (Lines 300–302), and view this as a potential direction for future work."}, {"review_id": "Er7MrFCSig", "reviewer": "Reviewer_WRxr", "summary": "This 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.", "questions": "1. 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?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The clarification that the context-length limitation was a choice for evaluation rather than an inherent architectural constraint of GIM is a crucial point that resolves my primary reservation about long-range dependencies.\n\nFurthermore, the reframing of GIM not as a tool for post-hoc analysis of LLMs, but as a standalone, interpretable implementation of a specific mechanism (induction) makes the paper's contribution much clearer. The proposed additions to the discussion section will be valuable for readers. The new empirical results demonstrating the fuzzy matching component's robustness on out-of-domain data are also very convincing.", "rebuttal": "We thank the reviewer for their thoughtful and encouraging feedback. We appreciate the recognition of (1) its transparency through n-gram-based reasoning and (2) its potential for interpretable next-token prediction. Below, we respond to the reviewer’s concerns.\n\n### W1. Fixed-size window and long-range dependency\nWe agree that the use of fixed-size sliding windows limits long-range retrieval, but this is a design choice for evaluation, not an inherent constraint of the GIM framework. In fact, GIM (and its predecessor Infini-gram) is designed to construct n-gram and fuzzy-matching indices over the entire available context. For example, when using OpenWebText, we treat each full document as a single context and build an Infini-gram over the entire span. This enables GIM to effectively consider contexts that extend to millions or billions of tokens, limited only by hardware. Thus, the performance drop under short contexts reflects a general property of all context-based models, not a structural limitation of GIM itself.\n\n### W2. GIM’s scope and relation to LLM internals\nGIM is not intended to decode the full internal mechanisms of pretrained LLMs. Instead, as described in Lines 109–118, it reimplements a specific, well-studied behavior—induction via repetition—as a standalone, fully interpretable module. The goal is not to analyze LLMs post hoc, but to show that a core capability observed in them can be recreated faithfully and transparently.\n\nWhile GIM may offer a useful reference for understanding induction-like behavior in LLMs (pattern copying from prior input spans), it does not aim to explain their full predictive capabilities, which rely on many additional mechanisms such as compression, memorization, and reasoning. Rather than serving as a diagnostic tool, GIM demonstrates how targeted mechanisms can be isolated, audited, and extended through interpretable design. As discussed in the Related Work section, GIM builds on both classical n-gram modeling and recent findings in mechanistic interpretability, but is distinct in offering a fully transparent generation process for next-token prediction.\n\nWe will clarify this distinction more explicitly in the revised Discussion section, emphasizing that GIM is a composable interpretability unit, not a general-purpose LLM substitute, as follows:\n\nUpdates to the first paragraph of the Discussion (Lines 286-289):\n> GIM constitutes a significant step toward building mechanistically interpretable language models inspired by pre-trained LLMs. Unlike black-box models or partially interpretable approaches, GIM provides full transparency in next-token prediction while substantially narrowing the performance gap between interpretable and black-box architectures across two diverse domains. Importantly, GIM is not a general-purpose LLM or a tool to decode its internals; it isolates and reimplements a single observed capability, induction via repetition, as a fully interpretable module. This shows that high-performance behaviors implicitly learned by LLMs can be transparently reconstructed to advance performance in the interpretable modeling space.\n\nUpdates to the second paragraph of the Discussion (Lines 298-302):\n> GIM shows limited gains when the input context is short or uninformative. Its modular design enables the available context to be expanded through retrieval-augmented generation [76] or external memory. Like kNN-LMs [77], GIM’s n-gram-based reasoning also struggles with tasks requiring deeper reasoning. Future work may explore hybrid approaches that pair GIM with black-box models for better trade-offs; our speculative decoding setup, where GIM serves as a transparent draft generator verified by a larger LLM (Appendix A.5), illustrates one example in this direction.\n\n### W3. Robustness of fuzzy matching under distribution shift\n\nTo empirically evaluate semantic ambiguity in fuzzy matching under domain shift, we trained the Fuzzy Matching Model on Pile-of-Law[a], a domain-specific legal corpus that differs significantly in both style and content from the general-domain corpora (e.g., OpenWebText). When evaluated on the Pile-val benchmark, the performance drops by only 1.4%p.\n\n|  | OpenWebText | Pile of Law |\n|---|---|---|\n| GIM (fuzzy) | 43.2 | 41.8 |\n| Infini-gram + GIM | 49.8 | 48.7 |\n\nThis 1.4%p gap indicates that the learned similarity metric remains robust even under substantial distributional shift. Moreover, because GIM combines fuzzy matching with an exact n-gram component—which is distribution-agnostic—the overall prediction performance remains largely stable. These results suggest that the hybrid design (Infini-gram + GIM) preserves accuracy and interpretability even for out-of-domain inputs for neural similarity metric.\n\n[a] Henderson. et al. \"Pile of law: Learning responsible data filtering from the law and a 256gb open-source legal dataset.\" NeurIPS (2022).\n\n### Q1. Structural relation between GIM and LLM induction heads\nGIM directly implements the behavior of induction heads observed in LLMs in an explicit and auditable way. While induction heads in LLMs emerge implicitly through training, GIM constructs this mechanism by design, allowing for faithful, fully traceable attribution. GIM is not intended to serve as a tool for reverse-engineering LLM internals; instead, it distills one known mechanism into a standalone, interpretable module (as we mentioned in our response to W2 above).\n\nGIM may serve as a useful reference point for understanding induction-like behavior in black-box models. For example, one could compare the predictions of GIM and those of a large LLM on a controlled task to evaluate whether the LLM exhibits similar input-output dependencies, potentially offering a lightweight probe into the presence or strength of induction-head-like activity (for a detailed comparison, please refer to our response to Reviewer 3WB2, **W2: Comparison to attribution-based interpretability**). Still, our goal with GIM is not to interpret full LLM behavior, but to demonstrate that interpretable modules can replicate specific high-performance behaviors seen in large models.\n\n\n### Q2. Role of exact vs. fuzzy matching\nBoth n-gram (exact) and fuzzy (neural similarity) matching play important and complementary roles in GIM, but fuzzy matching tends to be more decisive for overall performance and robustness.\n\nAs shown in Table 1 (Lines 190–195), GIM (fuzzy) consistently achieves higher accuracy than GIM (exact) across datasets. Similarly, Table A3 (Appendix Lines 599–605) shows that removing the fuzzy component from the full model (i.e., Infini-gram + GIM (exact), the row for ‘w/o GIM (fuzzy)’) results in a greater performance drop than removing the exact component, indicating the stronger standalone utility of fuzzy matching.\n\nHowever, we view the hybrid design as key to GIM’s overall performance and adaptability. Exact n-gram matching provides high-precision retrieval when long repeated spans are available, enabling accurate token prediction based on exact recurrence. Fuzzy matching enhances robustness to minor variations such as typos, rephrasings, or paraphrases, especially when exact matches are sparse. As illustrated in Figure 4(b), fuzzy matching performs better when the effective n is small, i.e., when long exact matches are unavailable. In contrast, when high-effective-n matches exist, exact matching tends to be more effective. GIM dynamically selects between these estimates based on match quality and effective n (Equation (5)), enabling it to maintain high performance even under input perturbations or distribution shifts.\n\nThis hybrid approach ultimately allows GIM to balance accuracy and robustness, by accurately modeling repetition when it occurs while generalizing gracefully in its absence (e.g., under distribution shifts or in low-redundancy contexts).\n\n### Q3. Future directions for scaling context length\nWe appreciate the reviewer’s suggestions. Before addressing them directly, we would like to reiterate a key point discussed in our response to **W1 (Fixed-size window and long-range dependency)**: if the concern is about the length of the considered context, analogous to how LLMs use long prompts, then GIM does not face an inherent limitation. In our implementation, GIM is applied over the full available context (e.g., an entire OpenWebText document), which may span millions of tokens. Its context length is constrained only by available hardware, not by the model’s design.\n\nIf the concern instead pertains to modeling higher-level abstractions or long-range semantic dependencies beyond surface-level repetition, we agree this is an important direction for future work. Among the proposed options—dynamic window scaling, hierarchical context aggregation, and external memory integration—we view external memory as the most promising extension. GIM’s retrieval-based nature makes it a natural fit for augmentations like memory banks or document caches, which could expand its effective range while preserving transparency and token-level traceability. We already note in the Discussion (Line 299) that retrieval-augmented generation (RAG) is a natural future direction for GIM, as it would similarly expand context by integrating external knowledge sources.\n\nIf we have misunderstood the reviewer’s intent, we would be glad to revisit this question with further clarification."}]

Author rebuttal:

Meta-review: The authors introduce the Generalized Induction-Head Model (GIM) as an "interpretable" model. In particular, GIM generates next tokens as a function of completions to matched phrases, permitting explicit attribution to matched phrases.

While there were some reservations from certain reviewers about framing the experimental setup (e.g., 3WB2), all reviewers nonetheless found this work compelling: It is a nice example of drawing on insights from mech interp to inform model design. The empirical results are compelling, and the work is likely to be of relatively wide interest to the NeurIPs community.
Decision: Accept (poster)