Interpretable Language Modeling via Induction-head Ngram Models (2411.00066v1)

Abstract: Recent LLMs have excelled across a wide range of tasks, but their use in high-stakes and compute-limited settings has intensified the demand for interpretability and efficiency. We address this need by proposing Induction-head ngram models (Induction-Gram), a method that builds an efficient, interpretable LM by bolstering modern ngram models with a hand-engineered "induction head". This induction head uses a custom neural similarity metric to efficiently search the model's input context for potential next-word completions. This process enables Induction-Gram to provide ngram-level grounding for each generated token. Moreover, experiments show that this simple method significantly improves next-word prediction over baseline interpretable models (up to 26%p) and can be used to speed up LLM inference for large models through speculative decoding. We further study Induction-Gram in a natural-language neuroscience setting, where the goal is to predict the next fMRI response in a sequence. It again provides a significant improvement over interpretable models (20% relative increase in the correlation of predicted fMRI responses), potentially enabling deeper scientific investigation of language selectivity in the brain. The code is available at https://github.com/ejkim47/induction-gram.

• The paper introduces Induction-Gram, a new language model framework combining interpretable ngrams with a neural induction head for context-sensitive prediction.
• This approach utilizes fuzzy matching and a distilled model to significantly enhance next-token prediction accuracy while preserving interpretability.
• Induction-Gram demonstrates practical utility in speculative decoding for faster inference and improves fMRI response modeling in natural language neuroscience.

Interpretable LLMing via Induction-head Ngram Models

The paper introduces Induction-head ngram models (Induction-Gram), a new paradigm that bridges the gap between interpretable and neural LLMs. The authors present a framework that leverages the efficiency and interpretability of ngram models, augmented with a hand-engineered ``induction head,'' to produce LLMs that provide improved next-token prediction. This induction head employs a custom neural similarity metric to search within the model's input context for potential next-token completions, thus facilitating ngram-level grounding and context-sensitive prediction.

Core Contributions

The primary contributions of the paper are as follows:

1. Framework Development for Induction-Gram: The proposed Induction-Gram enhances traditional ngram models with an induction head, effectively integrating context-driven prediction methods. This novel approach bolsters the model's performance without sacrificing interpretability.
2. Efficiency and Applicability: The authors present evidence of the computational efficiency of Induction-Gram, demonstrating its capability to significantly improve next-token prediction accuracy. This improvement is observed in diverse textual settings, with performance gains reaching up to 26 percentage points over baseline interpretability-focused models.
3. Speculative Decoding: The model's application in speculative decoding showcases its utility in scenarios requiring rapid inference, boasting substantive inference speed improvements.
4. Natural-language fMRI Domain Application: The paper extends the applicability of Induction-Gram into the domain of natural language neuroscience. A notable 20% increase in the correlation of predicted fMRI responses over interpretable models is reported, presenting a potent mechanism for exploring language selectivity in the brain.

Key Methodological Insights

The Induction-Gram model remains entirely interpretable, a distinct departure from the black-box nature of contemporary LLMs. By utilizing fuzzy matching within input contexts, the induction head efficiently mines relevant patterns, akin to "induction heads" in transformer architectures. This similarity-based mechanism not only fine-tunes token predictions but also aids in maintaining the interpretability of prediction pathways.

Furthermore, the introduction of a small-scale Fuzzy Matching Model for similarity scoring—trained via knowledge distillation—improves computational efficiency by substituting large-scale LLM inference with lightweight correlation scoring.

Experimental Validation

Comprehensive experiments validate the induction head's role in enhancing prediction accuracy across multiple benchmark textual datasets. For instance, in the BabyLM and Pile datasets, the Induction-Gram with fuzzy matching consistently outperforms traditional ngram methods and even matches the performance of full-scale LLMs in various settings.

The speculative decoding patterns evince up to twice the typical inference speed improvement when pairing the Induction-Gram as the draft model, underpinning its practical utility in low-compute environments.

In the fMRI application, the precision in modeling neural activations further establishes the framework's versatility. Comparative models attest to the superior performance of the Induction-Gram approach, predominantly in narrative comprehension contexts—a testament to its adept contextual mapping and predictive accuracy.

Implications and Future Prospects

The implications of Induction-Gram stretch across the theoretical and practical domains of artificial intelligence. Its success unveils pathways for the development of interpretable LLMs capable of efficient, context-driven inferences while retaining transparency.

The model's application in neuroscientific research could catalyze novel explorations into natural language processing mechanisms within the human brain, providing interpretable mappings and insights into cerebral language processing pathways.

Future developments could explore integrating additional mechanistic elements discovered within transformer models, potentially enhancing cognitive tasks that necessitate intricate reasoning capabilities. Moreover, the framework highlights prospects for evolving hybrid interpretive models that combine both rule-based methodologies and state-of-the-art neural mechanisms. As a direction, augmenting the model with retrieval-augmented-generation techniques may further expand its competence in matching contexts across diverse repositories.

Conclusively, Induction-Gram exemplifies a stride towards high-efficiency, interpretable AI systems, balancing the dichotomy of explanation and efficacy in a field where AI's interpretability garners increasing demand.