Emergence of Minimal Circuits for Indirect Object Identification in Attention-Only Transformers

For an attention-only model, the state of the residual stream at layer $l$, denoted as $x^{(l)}$, is strictly a linear combination of the initial embeddings and the outputs of all preceding attention heads. Let $x_{embed}$ and $x_{pos}$ denote the token and positional embeddings, respectively. The residual stream just before unembedding in a single-layer model is formalized as follows.

where $\text{out}_{h}^{(0)}$ is the output vector written to the residual stream by head $h$ in layer $0$. The final logit prediction for any token $t$ is computed by projecting this residual stream onto the unembedding matrix $W_{U}$: $L_{t}=(x^{(1)})^{T}W_{U[:,t]}$.

To compute its output, each attention head $h$ reads from the residual stream using three projection matrices: the Query matrix ($W_{Q}^{h}$), the Key matrix ($W_{K}^{h}$), and the Value matrix ($W_{V}^{h}$). The queries and keys interact to determine the attention scores between tokens, while the values determine the information moved across the sequence, which is then projected back into the residual stream via an output matrix ($W_{O}^{h}$).

We can analyze the behavior of individual heads by decomposing these operations into two distinct circuits (Elhage et al., 2021). The $QK$ (Query-Key) circuit dictates the attention scores between a query token and a key token, represented by the end-to-end transition matrix $M_{QK}^{h}=W_{E}^{T}(W_{Q}^{h})^{T}W_{K}^{h}W_{E}$, where $W_{E}$ denotes the token embedding matrix. The $OV$ (Output-Value) circuit dictates how attending to a specific token updates the final output logits, represented by the transition matrix $M_{OV}^{h}=W_{U}W_{O}^{h}W_{V}^{h}W_{E}$. Because these exact matrices, $M_{QK}^{h}$ and $M_{OV}^{h}$, map directly from the vocabulary space back to the vocabulary space, they serve as the foundation for our spectral analysis in section˜4.2.3.

In multi-layer attention-only architectures, transformer heads develop functional hierarchies by composing information across layers (Elhage et al., 2021). Because the input to a head in a subsequent layer $j$ is the residual stream $x^{(j)}$, which contains the outputs of heads from an earlier layer $i$ ($i<j$), the projection matrices of layer $j$ directly read the information written by layer $i$. This interaction is formally categorized into three types of composition described below.

Models were trained from scratch on the symbolic IOI dataset using a cross-entropy loss to predict the token at the <MID> position. Each training batch contained all $60$ possible unique sequences ( $6\times 5=30$ ways of picking the names in the dependent clause and two ways of ordering them in the main clause). We used the AdamW (Loshchilov and Hutter, 2019) optimizer with the OneCycle (Smith and Topin, 2019) learning rate scheduler, with a maximum learning rate of $0.1$ and weight decay of $0.01$. Training and analyses were performed on a single NVIDIA A40 GPU using PyTorch (Paszke et al., 2019) and TransformerLens (Nanda and Bloom, 2022) libraries.

We observed that, for both the templates of our symbolic dataset, the first head consistently attends to the two names in the initial dependent clause (see figure˜2), indicating its role in identifying the relevant referents. The second head, however, always attends to the subject of the main clause and the non-repeated name in the dependent clause (see figure˜3). So, this head does most of the heavy lifting, finding out the unique set of tokens to attend to. Furthermore, it suggests that the second head is responsible for integrating the referential information with the context provided by the main clause to determine the correct output.

To understand how the model’s components contribute to the final prediction, we decompose the residual stream at the final token position (corresponding to <MID> token) into the contributions of those components. We then project these components onto directions in the embedding space corresponding to the correct and incorrect names, as well as their sum and difference, using LogitLens (nostalgebraist, 2020). Figure˜4 shows that the first head’s output is aligned closest with the sum direction, i.e., it represents the combined contribution of both the correct and incorrect names (correct + incorrect). On the other hand, the second head’s output aligns closest with the direction of the token difference, i.e., the contrast between the two name embeddings (correct – incorrect). Since the final logits are computed by adding the contribution of all the components, the logit component for the incorrect token roughly cancels out, and the direction corresponding to the correct token is amplified.

This analysis is also not foolproof because we can see that the second head also has some component in the direction of the correct token, as well as the sum direction. Nevertheless, we observe a clear division of labor between the two heads: one aggregates signals (additive), while the other suppresses the incorrect alternative (contrastive). Together, they form an additive-contrastive circuit to produce a clean, interpretable mechanism for generating the correct logits.

While random matrices typically exhibit symmetric eigenvalue distributions (Tarnowski, 2022), the QK and OV matrices in our two-head model display significant asymmetry, reinforcing their specialized functional roles (see figures˜6 and 5). Furthermore, on the top-right of each subfigures, we report the fraction of positive eigenvalues for each head calculated using the formula $\frac{\sum_{i}\lambda_{i}}{\sum_{i}\lvert\lambda_{i}\rvert}$, where $\lambda_{i}$ are the eigenvalues of the matrix and $\lvert\lambda_{i}\rvert$ are their magnitudes.

To isolate the model’s reliance on positional embeddings, we assign a single, averaged embedding to all name tokens, removing the model’s ability to differentiate them by identity. From figure˜7, we can consider the first head as a positional head that focuses on the positions of the names in the dependent clause, independent of their word embeddings, because the attention patterns in figures˜2(a) and 7(a) look the same.

The second head attends predominantly to the position of occurrence of the subject in the main clause. However, despite this positional focus, it attends to the name in the dependent cause, not repeated in the main clause. This indicates that the second head is responsible for integrating positional as well as contextual information to determine the correct output.

To study how the model utilizes positional information, we train the same model architecture without any positional embeddings. The model achieves an accuracy of $\approx 70\%$ on the IOI task, with $\approx 67\%$ probability on the correct token, indicating that positional embeddings are not strictly necessary for the model to learn the task. Unlike the ones with positional embeddings, the heads trained in this manner exhibit similar attention patterns, managing to focus primarily on the correct non-repeated name (see figure˜8). This suggests the model can fall back on purely semantic contextual relationships, though explicit positional cues drastically simplify the optimization landscape, allowing it to reach $100\%$ accuracy and perfectly decouple into additive-contrastive roles.

Figures˜9 and 10 show the attention heatmap for a two-layer, one-head model, averaged across “BAAB” and “BABA” templates, respectively. We observe that the attention patterns of both layers change depending on the template. From figures˜9(a) and 10(a), we observe that the <MID> token in the first layer still attends to both the name tokens in the dependent clause for both templates, similar to the first head of the one-layer two-heads model. However, the S2 token in the first layer changes its attention pattern based on the template; it attends more to the IO token than the S1 token for the “BAAB” template and more to the S1 than the IO token. So, the first head is not solely positional, but aggregates information to S2 token to be used by the latter head. Although the attention pattern of the second layer for the <MID> token (see figures˜9(b) and 10(b)) seems almost similar to the attention pattern of the second head of the one-layer two-heads model (see figure˜3), this time it attends to the aggregated information from the first layer.

Similar to section˜4.2.4, we analyze the attention pattern of the heads when the distinction among the names is removed (see figure˜11). We observe that the attention pattern in figure˜11(a) is very similar to figure˜7(a), indicating that it has a strong positional focus. However, although the <MID> token in the second layer (see figure˜11(b)) attends primarily to the token before it, it changes its attention pattern based on the context provided by the first layer. So, the second head is not solely positional, but integrates positional as well as contextual information to determine the correct output.

To study the type of composition that the model is performing (Q, K, or V), we ablate them one by one by subtracting the output of the first layer from the corresponding input of the Q, K, and V matrices. We observe a drop in accuracy in the following order: Q-composition ($\approx 100\%$ drop), V-composition ($\approx 93.33\%$ drop), and K-composition ($\approx 26.67\%$ drop). This indicates that the model is heavily relying on the Q and V-compositions to perform the task. So, we can conclude that the two-layer one-head model is indeed performing some composition to solve the IOI task, different from the one-layer two-heads model. This hints that finding a circuit capable of solving a given task using composition is an easier task for the optimizer than building two orthogonal subspaces in the residual stream.