This post aims to quickly break down and explain the dominant mental models interpretability researchers currently use when thinking about how transformers work.

In my view, the focus of transformer interpretability research is teleological: we care about the functions each component in the transformer performs and how those functions interact to yield the capabilities we see in language models today. From a functional understanding of transformer internals, we then hope to be able to answer other important interpretability questions such as, “Where/How is knowledge stored?,” “Do transformers have beliefs?”, and “How do transformers reason?” As such, the mental models described in this post will be functional theories about how researchers think about transformer internals, and not hypotheses about other interpretability questions, like the ones mentioned above.

There are three main components at work in the transformer: the attention heads, the MLP, and the additive residual stream that connects all the layers. The functions that each play aren’t clear and are the subject of much research, and the mental models described below will be different ways of thinking about each main component. None of these mental models are mutually exclusive, and in reality, transformer internals probably look like a messy combination of many of these models.

This post assumes a working understanding of transformers. For a primer on the transformer architecture, I recommend The Illustrated Transformer, Transformers from Scratch, and The Annotated Transformer. Other works that this post draws greatly from are the logit lens, Transformer Feed-Forward Layers Are Key-Value Memories, ROME, and the Anthropic papers.

The Models

Residual Stream as Output Accumulation

The residual stream is simply the accumulation of all the stuff the transformer wants to say at the end of its inference step.

Clearly, the last residual hidden state is the transformer’s prediction (before it gets projected to the vocabulary), and all this mental model is suggesting is that the prior hidden states are less-refined versions of the final residual hidden state. The strongest version of this model states that the middle hidden states are the model’s nascent predictions, and the weaker version just states that the middle hidden states&nbs...

Posts

Wiki Contributions

Comments