Interpretable Language Modeling via Induction-head Ngram Models

Thank you for your thoughtful and positive feedback on our work. Your recognition of the simplicity and elegance of our proposed methods is truly encouraging. Below, we address your concern in more detail.

W1: Robustness across Tokenization Systems

1) Performance Variation Across Tokenizers

While the numerical performance of Induction-Gram differs between the GPT-2 and LLaMA-2 tokenizers, our analysis in Table 1 shows that the relative improvement from Infini-Gram to Induction-Gram is consistent across both systems. This suggests that the core mechanism of Induction-Gram operates reliably regardless of the tokenizer, even though absolute performance metrics may vary.

2) Factors Affecting Absolute Performance

The observed performance disparity can be explained by two key characteristics of the LLaMA-2 tokenizer: Increased Reference Tokens: The LLaMA-2 tokenizer generates more granular tokens, providing a richer set of references for the induction mechanism to utilize effectively. Reduced Prediction Space: With a smaller vocabulary size (32k vs. GPT-2's 50k), the LLaMA-2 tokenizer simplifies the next-token prediction task, which inherently improves accuracy.

These differences highlight the interaction between tokenization systems and model performance but do not diminish the generality or robustness of Induction-Gram’s core mechanisms.

Minor Corrections

Additionally, thank you for pointing out the minor corrections. We have incorporated these into the revised version of the paper, with the changes highlighted in blue.

Q1: Potential Integrations with Other Methods

There are many interesting potential directions involving bridging Induction-Gram with more powerful methods to improve performance in difficult input contexts. For example, Induction-Gram can struggle with very short or under-specified contexts (e.g., a single-sentence question), as the input does not provide enough contextualized information to search within.

One potential solution is to first use RAG to retrieve documents relevant to the question, incorporate the retrieved text into the Induction-Gram context, and then allow Induction-Gram to generate accurate, interpretable outputs based on this enriched context. This approach could significantly improve prediction performance at the cost of some efficiency (as the retrieval step may be computationally expensive).

Q2: Assessment of Induction-Gram on Diverse Setting

We conducted two additional assessments to address the performance of Induction-Gram on domain-specific datasets and with larger reference corpora.

1) Domain-specific dataset

We evaluated the model using the Pile of Law dataset [a], which contains English-language legal and administrative texts, such as court opinions, contracts, administrative rules, and legislative records. The table below summarizes the next-token prediction performance on 50k samples from Pile of Law, using the LLaMA-2 tokenizer. For Infini-Gram, Pile-train was used as the reference corpus.

These results demonstrate that Induction-Gram, with the introduction of induction-heads and fuzzy matching, outperforms Infini-Gram, consistent with the findings in our paper. Notably, Induction-Only (fuzzy) achieves significantly better performance than Induction-Only (exact). The complex and nuanced language in legal texts likely benefits more from the flexibility provided by fuzzy matching.

2) Impact of Larger Reference Corpora

Additionally, we explored the effect of using larger reference corpora on Induction-Gram’s performance. The table below presents next-token prediction accuracy on BabyLM-test (using the LLaMA-2 tokenizer) with an expanded reference corpus compared to the one used in the paper.

These results highlight two key points:

• Performance Scaling: The performance of Infini-Gram improves with a larger reference corpus, and this improvement translates directly to Induction-Gram.
• Competitiveness with LLMs: Induction-Gram achieves performance close to that of LLaMA-2-7B, demonstrating its ability to benefit from larger datasets without compromising its interpretability.

These findings demonstrate that our proposed strategies-induction-heads and fuzzy matching-remain effective even when scaling the corpus size, reinforcing their role in achieving the study’s objectives.

Q3: Impact of training data on the Fuzzy Matching Model

Regarding the training data for the Fuzzy Matching Model, we intentionally selected general datasets such as OpenWebText and Pile-train to ensure broad applicability. To explore the impact of different data distributions, we conducted an experiment using the Pile of Law dataset [a], which consists of English-language legal and administrative data, including court opinions, contracts, administrative rules, and legislative records. This dataset exhibits a significantly different distribution compared to OpenWebText and Pile-train.

For a fair comparison, we trained the Fuzzy Matching Model on a similar number of samples from Pile of Law as we used for OpenWebText. The results are summarized in the table below. Models trained on Pile of Law showed slightly lower performance in Induction-Only (fuzzy) and Induction-Gram compared to those trained on OpenWebText. This suggests that the language style and topic distribution of the training dataset influence the learned similarity metrics in the Fuzzy Matching Model. That said, the performance difference is relatively small, approximately 1%p, indicating that while dataset selection impacts performance, the Fuzzy Matching Model and Induction-Gram remain robust across different training distributions.

[a] Henderson, Peter, et al. "Pile of law: Learning responsible data filtering from the law and a 256gb open-source legal dataset." Advances in Neural Information Processing Systems 35 (2022): 29217-29234.

Thank you for your thoughtful and encouraging feedback on our work. We deeply appreciate your recognition of the methodological contributions and practical implications of the Induction-Gram. Below, we provide additional clarifications and present experimental results to further address your points and concerns. Some changes have also been incorporated into the revised manuscript, highlighted in blue for your convenience.

W1 (Q2): Comparison with Other models

First, we would like to clarify our rationale for the selected baselines and the scope of our comparisons:

As outlined in the paper:

"... there is a considerable gap between the performance of interpretable models and black-box LLMs in next- token prediction. There is a considerable gap between the performance of interpretable models and black-box LLMs in next-token prediction."

Our primary goal in this work is to develop an inherently interpretable model that offers competitive performance. In the current landscape of language modeling, interpretability is often addressed through post-hoc methods that aim to analyze or explain black-box models after training. While these approaches can provide insights into model behavior, they do not inherently build transparency into the prediction process, limiting their ability to offer meaningful rationales.

Given this context, few models are designed to prioritize both interpretability and performance. We chose Infini-Gram as our primary baseline because it is a state-of-the-art interpretable ngram model and directly aligns with our focus on inherently interpretable modeling. Comparisons with black-box models (e.g., GPT-2, LLaMA) were included to illustrate the performance gap between interpretable and black-box models. Models focusing solely on efficiency or black-box optimization techniques (e.g., pruning, distillation) were excluded, as they do not align with our primary objective of interpretability.

Relative Advantages and Disadvantages of Induction-Gram

The key advantage of Induction-Gram lies in its inherent interpretability. Unlike black-box LLMs, Induction-Gram provides transparent predictions by leveraging ngram properties, ensuring that each predicted token is grounded in a data-driven, explainable process. This approach adheres to the principles of interpretable modeling and enhances trust in the model’s outputs.

In addition, Induction-Gram addresses major limitations of Infini-Gram, as noted in our paper:

"Infini-Gram struggles with adapting to new contexts and with matching queries that cannot be found exactly within a reference dataset (e.g., typos or rephrasings)."

To overcome these challenges, we introduced induction-head mechanisms and fuzzy matching, which improve adaptability and robustness without compromising interpretability. As shown in Table 1, these enhancements lead to significant performance improvements over Infini-Gram, particularly in handling complex linguistic patterns.

While Induction-Gram introduces clear advantages, it also comes with trade-offs. The primary cost is increased resource usage, as building and running the induction mechanism requires additional computational resources. We acknowledge this limitation and view it as an area for future optimization.

W2: Ablation Study on Training the Fuzzy Matching Model

We included details of Fuzzy Matching Model in Section A.1 of the Appendix. Additionally, we conducted an ablation study to evaluate the effects of key components, including transformer block types (Relative Position Encoding, RPE vs. Sinusoidal) and the two loss terms (reverse KLD and CE). The results, summarized in the table below, illustrate the impact of these components on the performance of the Induction-Only (fuzzy) model on the BabyLM-test dataset (using LLaMA-2 tokenizer). Here, Fuzzy Matching Models were trained on the OpenWebText dataset.

Key observations are:

• Combining reverse KLD loss and CE loss leads to significant performance improvements, demonstrating their complementary roles in optimizing the Fuzzy Matching Model.
• Reverse KLD is more effective than forward KLD, suggesting its suitability for this context.
• Relative Position Encoding (RPE) outperforms the Sinusoidal transformer block, likely due to its ability to better capture contextual relationships and positional information.

This ablation study and detailed analysis have been included in the Appendix.

Thank you for your thoughtful observation and for recognizing the transparency of our proposed model. Below, we provide further clarification and present additional experimental results to address your concerns. Some changes have also been incorporated into the revised manuscript, highlighted in blue for your convenience.

W1. Performance Compared to LLMs

Below, we clarify our study’s scope, objectives, and the relevance of our findings.

1) Scope and Objectives of the Study

As mentioned in the paper:

"... there is a considerable gap between the performance of interpretable models and black-box LLMs in next-token prediction."

Our work aims to address this gap by improving the performance of interpretable models without sacrificing their transparency. Rather than directly competing with LLMs, our focus lies in enhancing interpretable language modeling. Specifically, we identified limitations of Infini-Gram, a state-of-the-art scalable n-gram model, such as difficulties in adapting to new contexts or handling mismatched queries (e.g., typos or rephrasings). To address these issues, we proposed two strategies: induction-heads and fuzzy matching. These techniques improve performance while preserving interpretability, as demonstrated throughout the paper. Thus, the significance of our work lies in demonstrating how transparent modeling can bridge the gap between interpretability and performance. By leveraging these strategies, Induction-Gram achieves substantial improvements over Infini-Gram and demonstrates the potential to approach the performance of black-box LLMs.

2) Impact of Larger Reference Corpora

Additionally, we explored the effect of using larger reference corpora on Induction-Gram’s performance. The table below presents next-token prediction accuracy on BabyLM-test (using the LLaMA-2 tokenizer) with an expanded reference corpus compared to the one used in the paper.

These results highlight two key points:

• Performance Scaling: The performance of Infini-Gram improves with a larger reference corpus, and this improvement translates directly to Induction-Gram.
• Competitiveness with LLMs: Induction-Gram achieves performance close to that of LLaMA-2-7B, demonstrating its ability to benefit from larger datasets without compromising its interpretability.

These findings demonstrate that our proposed strategies-induction-heads and fuzzy matching-remain effective even when scaling the corpus size, reinforcing their role in achieving the study’s objectives.

W2: Efficiency Comparison

We have included the acceptance rate in Table 2 to provide a more comprehensive evaluation. The updated table highlights that while the acceptance rate of our proposed method is lower compared to larger LLMs due to slightly reduced alignment with their outputs, it still achieves notable speed improvements during speculative decoding. This supports our primary objective of balancing speed and performance in practical, real-world applications.

Regarding the baseline choice, we selected a bare Transformer to focus on the interpretability and impact of our specific contributions. While comparisons with other compute-efficient models, such as TinyLLaMA, are valuable, they fall beyond the scope of this work. That said, we have included additional comparisons with TinyLLaMA in an H100 environment in the revised version of the paper to further address this concern.

Thanks for providing thorough comments and additional information for my review. They cralified some of my questions and I would increase my recommendation by 1 as it sounds enough to express the advantage of the proposed method.

I still recommend to split the fMRI part of the paper into another because it shrinks the space of the paper for analysing the method itself (e.g., I guessed the additional content in the appendix looks more reasonable to be included into the main text), but I wouldn't claim that this is critical.

We sincerely thank the reviewers for their constructive feedback and for acknowledging the effectiveness of our proposed model in next-token prediction. Your comments on the identified weaknesses and thoughtful questions have provided valuable directions for improvement. Below, we address each point in detail.

W1. Effectiveness in Various Settings

Our primary goal in this study is to improve upon the state-of-the-art for interpretable, next-token prediction (Infini-Gram) by adaptively incorporating input context. We show promising results in two domains: language modeling and fMRI encoding tasks. While we acknowledge the model’s constraints with shorter contexts, one potential direction could involve integrating retrieval-augmented generation (RAG) methods. This approach would allow Induction-Gram to retrieve and incorporate more representative texts from larger corpora, enhancing its contextual understanding. Although this might introduce additional computational steps, it could significantly improve accuracy and interpretability in short-context scenarios.

For reasoning tasks, we agree that this remains a broader challenge, even for advanced black-box LLMs, which have yet to demonstrate consistent reasoning abilities [a, b, c]. Induction-Gram is designed to bridge interpretability and contextualization; advancing its reasoning capabilities goes beyond the scope of this paper. Future research could explore tailoring the model to address reasoning-specific challenges.

[a] Wang, Siyuan, et al. "Can llms reason with rules? logic scaffolding for stress-testing and improving llms." arXiv preprint arXiv:2402.11442 (2024).

[b] Valmeekam, Karthik, et al. "Large language models still can't plan (a benchmark for LLMs on planning and reasoning about change)." NeurIPS 2022 Foundation Models for Decision Making Workshop. 2022.

[c] Li, Zhiming, et al. "Llms for relational reasoning: How far are we?." Proceedings of the 1st International Workshop on Large Language Models for Code. 2024.

Q1: Discussion on Larger Models

We interpret "larger models" in two possible ways: expanding the reference corpus for Infini-Gram or scaling up large language models (LLMs). Below, we address both interpretations.

1) Larger Reference Corpus for Infini-Gram

Infini-Gram serves as the foundation of our baseline, where scaling corresponds to increasing the size of the reference corpus. To explore this, we conducted additional experiments using a larger reference corpus released by the authors of Infini-Gram. The table below shows next-token prediction accuracy on BabyLM-test with the LLaMA-2 tokenizer.

As seen in the results, Infini-Gram’s performance improves with a larger reference corpus, and Induction-Gram also shows a corresponding performance gain. This demonstrates that the inductive mechanism of Induction-Gram remains effective regardless of the size of the reference corpus, emphasizing its continued relevance even as Infini-Gram scales. Therefore, our conclusions remain consistent.

2) Larger LLM

If the comment refers to comparisons with large-scale LLMs such as LLaMa3.1-405B, we omit these as they do not affect our main goal: improving the state-of-the-art for interpretable next-token prediction. Thus, we compare to SOTA baselines for interpretable next-token prediction, which still trail the black-box LLMs we compare against in performance. While larger LLMs often achieve superior raw performance, they do so at the expense of interpretability and computational cost. Our study prioritizes models that strike a balance between interpretability and performance. Induction-Gram excels in this niche, delivering meaningful performance improvements over other interpretable language models while maintaining a lightweight design.