AtP*: An efficient and scalable method for localizing LLM behaviour to components (2403.00745v1)

Abstract: Activation Patching is a method of directly computing causal attributions of behavior to model components. However, applying it exhaustively requires a sweep with cost scaling linearly in the number of model components, which can be prohibitively expensive for SoTA LLMs. We investigate Attribution Patching (AtP), a fast gradient-based approximation to Activation Patching and find two classes of failure modes of AtP which lead to significant false negatives. We propose a variant of AtP called AtP*, with two changes to address these failure modes while retaining scalability. We present the first systematic study of AtP and alternative methods for faster activation patching and show that AtP significantly outperforms all other investigated methods, with AtP* providing further significant improvement. Finally, we provide a method to bound the probability of remaining false negatives of AtP* estimates.

• The paper introduces AtP*, an enhanced attribution patching method for efficiently localizing LLM behavior.
• It employs a recalibrated softmax and dropout in the backward pass to reduce false negatives and improve reliability.
• The evaluation demonstrates that AtP* outperforms traditional methods in speed without sacrificing interpretability or accuracy.

Enhancing Interpretability of LLMs through Advanced Attribution Patching Techniques

Introduction

Understanding the internal mechanics of LLMs is paramount as their role in digital systems continue to expand. A critical aspect of this endeavor is attributing model behavior causally to its components, a task that is not only intellectually fascinating but also crucial for enhancing model transparency, reliability, and control. Despite the obvious necessity, tracing behaviors back to specific elements within state-of-the-art LLMs presents a considerable challenge due to their immense complexity. Activation Patching has been a preferred approach for its ability to directly compute causal attributions by intervening on model components. However, its linearly scaling cost with the number of components in SoTA models makes exhaustive investigation impractical. This paper addresses this scalability issue by exploring Attribution Patching (AtP), introducing a refined variant AtP* for improved efficiency and accuracy, and establishing a systematic comparison with alternative methods.

Attribution Patching and its Pitfalls

Attribution Patching (AtP), as an approximation to Activation Patching, offers significant speedups but is not without its limitations. Notably, it encounters two main classes of failure modes: false negatives arising from attention saturation and brittle false negatives resulting from cancellations of positive and negative effects. These failures substantially reduce AtP's reliability, potentially overlooking crucial components involved in model behavior.

Introducing AtP*

To counter these deficiencies, this paper proposes AtP*, an enhanced version of AtP, incorporating two key modifications. Firstly, it adopts a recalculated softmax operation for queries and keys to tackle gradient approximation issues in saturated attention scenarios. Secondly, it introduces dropout in the backward pass to mitigate cancellations, thereby reducing brittle false negatives. These adjustments considerably preserve AtP's scalability while significantly curtailing its proneness to false negatives.

Systematic Evaluation of Patching Methods

Through exhaustive comparisons with both brute force Activation Patching and various approximations, including novel alternatives, the paper establishes that AtP significantly outperforms other methods in speed without compromising accuracy. With the introduction of AtP*, there’s a noticeable further improvement in performance, validating the proposed enhancements. Additionally, the authors present a diagnostic method to statistically bound the probability of any remaining false negatives in AtP* estimates, adding a layer of reliability in its application.

Implications and Future Directions

Aside from demonstrating superior performance, the refinement and validation of AtP* have broader implications:

• Theoretical Advancements: AtP* contributes to the understanding of interpretability in LLMs, highlighting both the potential and limitations of gradient-based approximations.
• Practical Applications: By offering a scalable method for causal attribution, AtP* aids researchers in dissecting LLM behavior, paving the way for more interpretable and controllable models.
• Future Research: The findings invite further exploration into other components like layer normalization and extensions to edge attribution and coarser nodes, suggesting an expansive horizon for future investigations in LLM interpretability.

Conclusion

This paper's journey into refining Attribution Patching underscores the intricate balance between computationally feasible methods and the fidelity of causal attributions in LLMs. AtP*, with its careful adjustments, represents a significant step forward in this domain, offering a viable path for rigorously unraveling the mechanisms underpinning LLM behaviors. By leaning into the complexities and addressing the subtleties head-on, we inch closer to the larger goal of creating transparent, interpretable, and reliable AI systems.