Attribution methods for LLMs

Attribution is fundamental to mechanistic interpretability. We need to identify which model components (attention heads, MLP neurons, layers) are responsible for specific behaviors. Current methods face significant challenges.

Activation patching is the gold standard for causal attribution. This patches activations from a "clean" run into a "corrupted" run to identify which components carry critical information. However, it requires O(n) forward passes for n components, making it computationally expensive for large models. Attribution patching is a gradient-based approximation introduced by Nanda (2022) that uses one backward pass to estimate patching effects under a local linearity assumption. This is much faster but sacrifices some accuracy. One extension is logit-based gradient attribution which computies the gradient of a target logit (e.g., "Yes" token probability) with respect to each component. This is simple and fast, but not every behavior can be framed as a Yes/No question, models often prefer idiosyncratic first tokens (e.g., "Certainly," "Of course") rather than clean "Yes"/"No", and a single target logit may not capture the behavior of interest.

Can we do any better? We certainly can - join me and let's find out together. One quick idea: what if we train a probe for the behavior we're interested in, then take gradients with respect to the probe's output? Such an approach can allow attribution for any behavior that can be classified, not just Yes/No questions. It also has problems, e.g., the probe must be causal for the behavior.

Research questions:

1. Methodological. How do we best train probes for complex behaviors that translate cleanly into gradient attribution signals?
2. Comparative. How does different attribution methods compare to activation patching and attribution patching in terms of accuracy, efficiency, and the types of insights it enables?
3. Validation. Can our methods identify components that, when intervened upon, actually change the target behavior? (i.e., causal validation)
4. Generalization. How does different attribution methods work across model scales, architectures, and behavior types?

1. Better tools for understanding model internals. Ambitious mechanistic interpretability aims to reverse-engineer neural networks to understand their computations. Current attribution methods have significant limitations - devising better attribution tools can directly enable better understanding of model capabilities and risks.
2. Enabling attribution for safety-relevant behaviors. Many safety-relevant behaviors are difficult to frame as simple logit targets. We want to enable gradient-based attribution for complex behaviors exhibited by LLMs. This opens new avenues for studying how safety-relevant behaviors are implemented.
3. Efficient auditing at scale. Cheaper and more effective attribution methods can enable routine auditing of model components for concerning behaviors. This can help scale interpretability to match scaling of model capabilities.

Mentees will be primary researchers driving experimental work. They will implement attribution pipelines across different problems and domains, and run systematic comparisons against baseline methods. The mentees will be expected to propose and execute experiments autonomously. They will interpret the results, identify patterns, and draw conclusions. Mentees have significant autonomy in experimental choices within the overall research direction. We expect active intellectual contribution, e.g., proposing hypotheses, critiquing methods, and iterating on approaches. Deliverables include a paper/blog-post and developing clean and reusable implementations for the community, depending on the outcome.

Proficiency and experience with PyTorch is a must-have prerequisite.