Improving Neuron-level Interpretability with White-box Language Models (2410.16443v2)

Abstract: Neurons in auto-regressive LLMs like GPT-2 can be interpreted by analyzing their activation patterns. Recent studies have shown that techniques such as dictionary learning, a form of post-hoc sparse coding, enhance this neuron-level interpretability. In our research, we are driven by the goal to fundamentally improve neural network interpretability by embedding sparse coding directly within the model architecture, rather than applying it as an afterthought. In our study, we introduce a white-box transformer-like architecture named Coding RAte TransformEr (CRATE), explicitly engineered to capture sparse, low-dimensional structures within data distributions. Our comprehensive experiments showcase significant improvements (up to 103% relative improvement) in neuron-level interpretability across a variety of evaluation metrics. Detailed investigations confirm that this enhanced interpretability is steady across different layers irrespective of the model size, underlining CRATE's robust performance in enhancing neural network interpretability. Further analysis shows that CRATE's increased interpretability comes from its enhanced ability to consistently and distinctively activate on relevant tokens. These findings point towards a promising direction for creating white-box foundation models that excel in neuron-level interpretation.

• The paper's main contribution is introducing crate, a transformer that embeds sparse coding to boost neuron-level interpretability.
• The methodology adapts Multi-Head Subspace Self-Attention with a causal mask and over-parameterizes ISTA to reduce reconstruction noise.
• Experimental results demonstrate up to a 103% improvement in interpretability scores, paving the way for safer and more controllable AI models.

Improving Neuron-level Interpretability with White-box LLMs

The research paper presents significant strides in enhancing neuron-level interpretability of LLMs through a novel approach—integrating sparse coding directly into the model architecture. Unlike traditional post-hoc sparse coding techniques, this paper proposes embedding sparsity within the LLM itself, yielding a more interpretable and robust architecture called Coding Rate TransformEr (crate). This method holds promise for understanding neuron behaviors and improving the transparency of LLMs, which is critical for addressing issues such as bias, hallucination, and catastrophic forgetting.

Architectural Innovation and Methodology

The core contribution of this work lies in the development of crate—a white-box transformer-like architecture designed to capture sparse, low-dimensional structures in data. The architecture strategically integrates sparse coding principles within a causal LLM context, aiming to enhance interpretability from the ground up. Key changes include the adaptation of the Multi-Head Subspace Self-Attention (MSSA) with a causal mask and the over-parameterization of the Iterative Shrinkage-Thresholding Algorithm (ISTA) block. These adjustments ensure the architecture is well-suited for LLMing tasks while maintaining the interpretability benefits derived from sparsity.

The architectural adjustments focus on creating a LLM that inherently supports neuron-level interpretation by aligning token representations with distinct semantic axes. This approach stands in contrast to the often lossy and computationally intensive nature of traditional sparse auto-encoders (SAEs), which are applied post-hoc and add reconstruction noise.

Empirical Validation

The experimental results indicate that crate significantly enhances interpretability across various evaluation metrics when compared with models like GPT-2. For instance, the interpretability scores show up to a 103% relative improvement under specific conditions. The results underscore crate's ability to consistently and selectively activate neurons on relevant tokens, a property not as evident in traditional architectures where neurons often fire on semantically unrelated tokens.

Crate's interpretability was further validated using automated evaluation methods from both OpenAI and Anthropic, showing robust performance on metrics that assess model capability to distinguish semantically relevant tokens. The use of interpretability evaluation models like Mistral-7B and LLaMA-2 further strengthens the reliability of these findings.

Theoretical and Practical Implications

The integration of white-box modeling principles into LLMs opens new avenues for understanding and controlling the inner workings of deep learning models, especially those trained on complex tasks like language processing. By embedding sparsity into the model architecture, crate promises not only improved interpretability but also aids in efficient and accurate model editing and control.

This research suggests a shift towards building LLMs that balance performance and interpretability, addressing longstanding issues with traditional black-box approaches. The inherent sparsity provides a pathway for designing models with reduced interpretability challenges, paving the way for safer and more accountable AI systems.

Future Prospects

While crate presents an exciting direction for LLM design, future work should explore the potential trade-offs between model performance and interpretability. Investigations into the scalability and domain-specific adaptations of crates, possibly incorporating hybrid approaches that balance built-in sparsity with elements of traditional transformer architectures, could yield further insights.

In conclusion, this paper offers a promising method for enhancing neuron-level interpretability through innovative architectural design, suggesting a transformative shift in how LLMs are developed and understood. By leveraging sparsity from the outset, crate positions itself as a foundational step towards more transparent and controllable AI.