Feature Resolved Attention

Dictionary learning methods like Sparse Autoencoders and Crosscoders are meant to isolate largely independent features that function like the “atoms” of a model - combine independently to produce the output. This perspective has had notable success in understanding the MLP layers of a model. It is possible, however, for these features to act collectively. In fact, in Attention layers, which mix together features at different positions, this is unavoidabe. Moreover, model behaviour which relies on features interacting is likely to be important for deception, since the model must mix features representing its own goals, those of the user, and its action.

In earlier work (https://openreview.net/pdf?id=a98SfVkYRk, Section J.4) we outlined how Attention layers can be resolved from token-token maps into feature-feature maps. In this project, we propose to develop the resulting “Feature-Resolved Attention” maps into a practically useful tool for alignment researchers.

This project can be approached from a theoretical or applied angle. A theoretical approach would focus on deriving explicit formal bounds from feature interaction, as we do for the MLPs in (https://openreview.net/pdf?id=a98SfVkYRk, Section J.4). Specifically, we would:

1. Write down explicit equations for decompositions of attention into feature-interactions.
2. Construct a toy model in which we would measure the trade-off between computational cost and tightness of bound, as in Fig 6. of https://proceedings.neurips.cc/paper\_files/paper/2024/file/90e73f3cf1a6c84c723a2e8b7fb2b2c1-Paper-Conference.pdf
3. Demonstrate that we can control this trade-off via a loss penalty for crosscoder training, as in Fig 1b. of https://openreview.net/pdf?id=a98SfVkYRk.

A more applied approach would focus on measuring the impact of feature interactions on deceptive alignment, extending the work on MLPs in Section 4 of (https://openreview.net/pdf?id=a98SfVkYRk)

1. Construct a toy model for illustrating different possible choices for decomposing attention maps.
2. Measure the effect of interaction on canonical cases of models fine-tuned to elicit deceptive behaviour.
3. Use the resulting FRA attention maps to intervene on model behaviour, and compare to a baseline of single feature ablation.

Although it is more important to derive high quality results that we thoroughly understand, successfully completing either approach would be a strong candidate for a conference paper.

This project improves a standard tool in mechanistic interpretability - attention maps. The method proposed here is well suited to chain-of-thought rollouts, and could give us a valuable additional tool to defeat compressed chain-of-thought pressures.

Mentees will work collaboratively to derive theoretical results and design and implement experiments. I expect to spend time co-working together with mentees early on in the project (1-2 weeks), to craft a project plan, help run experiments, and work through the theory.

We will then collaboratively construct a more detailed project plan, and mentees will largely take ownership of the project although I will make suggestions about next steps and co-work on bottlenecks when needed.

• Good enough understanding of A Mathematical Framework for Transformer Circuits (https://transformer-circuits.pub/2021/framework/index.html) to understand the equations in J.4 of (https://openreview.net/pdf?id=a98SfVkYRk)

• Able to run simple but rigorous causal experiments (ablations, activation patching, etc...)

• Comfortable manipulating and understanding the equations for SAEs.

Prefer London so we can co-work in person at the MATS space and collaborate with other MATS scholars easier. I could also provide snacks. I am OK with working remote, though.