GIM: Improved Interpretability for Large Language Models

TL;DR

The paper addresses the problem of faithful interpretability in large language models by identifying a novel attention self-repair effect where softmax redistribution masks the true influence of important attention scores.It introduces Gradient Interaction Modifications (GIM), a gradient-based attribution framework that combines temperature-adjusted softmax gradients, Layernorm freeze, and gradient normalization to account for interactions that cause self-repair during backpropagation.Empirical results across Gemma, LLaMA, and Qwen families on six datasets show that GIM outperforms existing gradient-based circuit identification and attribution methods in faithfulness, with ablations confirming the contribution of each modification.The work provides practical insights into the inner mechanisms of transformers, improves explanation faithfulness, and offers a reproducible framework for evaluating mechanistic interpretability and safety-related aspects of LLMs.

Abstract

Ensuring faithful interpretability in large language models is imperative for trustworthy and reliable AI. A key obstacle is self-repair, a phenomenon where networks compensate for reduced signal in one component by amplifying others, masking the true importance of the ablated component. While prior work attributes self-repair to layer normalization and back-up components that compensate for ablated components, we identify a novel form occurring within the attention mechanism, where softmax redistribution conceals the influence of important attention scores. This leads traditional ablation and gradient-based methods to underestimate the significance of all components contributing to these attention scores. We introduce Gradient Interaction Modifications (GIM), a technique that accounts for self-repair during backpropagation. Extensive experiments across multiple large language models (Gemma 2B/9B, LLAMA 1B/3B/8B, Qwen 1.5B/3B) and diverse tasks demonstrate that GIM significantly improves faithfulness over existing circuit identification and feature attribution methods. Our work is a significant step toward better understanding the inner mechanisms of LLMs, which is crucial for improving them and ensuring their safety. Our code is available at https://github.com/JoakimEdin/gim.

Figures (12)

• Figure 1: Attention self-repair. When multiple values associated with large attention weights contain similar information, ablating one attention score has little effect on the output because the softmax activation function compensates by increasing the weight of other positions. This results in perturbation-based and gradient-based explanation methods underestimating the importance of components that contribute to the attention scores through the keys and queries.
• Figure 2: The attention self-repair effect and how temperature-adjusted softmax gradients approximate the joint ablation effect for LLAMA-3.2 1B on FEVER. The figures compare ablating the two largest attention scores jointly with a) ablating them separately, b) gradients, and c) TSG.
• Figure 3: Comparison of circuit identification methods for LLAMA-3.2 1B (95% CI). The top row depicts comprehensiveness per layer ($\uparrow$). The bottom row depicts sufficiency ($\downarrow$).
• Figure 4: Relative improvement of adding modifications to GradientXInput. We compare adding layernorm freeze and grad norm with also adding TSG. Each point represents the average statistics of a dataset. The y-axis shows the improvements over GradientXInput.
• Figure 5: Faithfulness per layer for LLAMA-3.2 1B. The top row depicts comprehensiveness per layer, where higher is better. The bottom row depicts sufficiency, where lower is better.