# Positional kernels of attention heads

In this post, we present a decomposition of attention patterns into position-dependent and content-dependent components, under the assumption that the representation of identical content is approximately translation invariant up to a static layer-specific positional embedding. While this assumption cannot hold exactly due to layer normalization and other non-linearities, we find it useful for analysis. It can be viewed as a logical consequence of the Linear Representation Hypothesis.

## Attention decomposition

For a particular attention head `h`, consider an input sequence `x_i` at position `i` in the sequence. From the softmax probability formula:

\ = \frac{\text{pos}_{i,\text{embedding}} \cdot f_{x_i,\text{embedding}}}{\sum_{j=1}^n \text{pos}_{j,\text{embedding}} \cdot f_{\text{embedding},\text{embedding}}} \]

This is difficult to analyze because the denominator involves the entire previous sequence.

However, within a fixed context, we can model the sequence as drawn from i.i.d. representations according to some distribution. While nearby representations will correlate, distant representations should have low correlation within a fixed context.

This is a key place where we make use of the assumption that the `E` terms are translation-invariant up to static positional embedding. Without this assumption, the representations drawn from a fixed context can't be modeled as identically distributed across different positions.

Under these assumptions:

\} \cdot f_{\text{embedding},\text{embedding}}\right) = \left(\sum_{j=1}^{n-1} \text{pos}_{j,\text{embedding}}^2\right) \cdot \text{Var}(f_{\text{embedding},\text{embedding}}) \]

Two key factors determine this variance:

- Measures the spread of positional patterns. For uniform attention across `m` tokens, equals `1/m`.
- Quantifies variation in content-dependent component.

Heads with slow decay / global positional patterns have small values for `sum_{j=1}^n \text{pos}_{j,\text{embedding}}^2` because they are spread out. If they also have low content-dependent variance within a context, the softmax denominator will have low variance.

For a fixed `embedding`, this means the softmax denominator will concentrate around its expected value, effectively becoming a context-dependent constant.

If an attention head has low content-dependent variance across almost all contexts and `x_n` values, and a broad positional pattern (as measured by `sum_{i=1}^{n} \text{pos}_{i}^2`), we call it a "contextual attention head." Some have very small content-dependent components, appearing visually like fixed positional kernels averaging over previous positions, meaning that we can drop the context-dependent factor. Others are less well behaved, and for instance weigh keywords above stopwords, while still preserving the overall low content-dependent variance.

Contextual attention heads within the same layer as each other that have similar positional kernels are natural candidates for attention head superposition. Within each fixed context, each of the heads is effectively computing a positionally weighted, weakly content-modulated, linear summary of the text. We can combine together these linear summaries across contextual attention heads with similar positional kernels to form a large "contextual circuit."

## First layer contextual circuit

Recall the positional kernels of the first layer heads from above. Heads 0,1,2,6,8,9, and 10 all have very similar positional kernels. Head 1 is a duplicate token head, meaning there is a high content-dependent variance. The remaining heads empirically have low variance in their softmax denominator. In the first layer, we can get a visual for the variance in the softmax denominator for a fixed `embedding` by simply substituting the corresponding token into the residual stream at position `i` for `500 < i < 600` and plotting the position-normalised softmax denominator.

Below is a plot of the softmax denominators for a diverse variety of input texts with `embedding` substituted with `' the'` at each position `i`:

- Spanish Text
- Political Conspiracy ()
- Astronomy ()
- Commonwealth English vs American English (, Positive token contributions are British spellings and references, and negative token contributions are American spellings and references. Note top dataset activations are just spamming the "£" symbol because this strongly distinguishes British from American, but this neuron does not refer to "countries, regions, and large numerals for financial amounts"))
- Medieval text
- Cooking
- Bracket Matching (,+1 contribution for tokens containing '(', -1 contribution for tokens containing ')')

And many more contextual neurons.

## Metric for spread of positional pattern

The above analysis naturally suggests `sum_{i=1}^{n} \text{pos}_{i}^2` as a metric for the spread of positional patterns. Although, as previously mentioned, later layer attention heads often turn themselves off by attending to `<end-of-text>`, so we should exclude the first `5` or so positions and take the softmax over the remaining positions for this metric.

We define the Effective Token Count (ETC) of a positional kernel to be `1/sqrt(sum_{i=5}^{n} (\text{pos}_i)^2)`. If the positional kernel attends uniformly to `m` tokens, it will have an ETC of `m`, giving us a natural interpretation of this definition.

Now we expect local positional patterns to have a low ETC, as they attend to just the previous ~5 tokens. A slow positional decay will have a higher ETC, and a uniform positional kernel will have an ETC of `n`.

Language models aren't very creative with their positional kernels, so the ETC gives a good summary of the type of positional kernel at an attention head.

Reducing the spread to a single summary statistic allows us to produce a single graph giving an idea for the positional kernels across all heads and layers of a single language model.

Below is the heatmap of `sqrt(ETC)` across all layers and heads of GPT-2 and Tiny stories. I found it best to plot `sqrt(ETC)` for visualization purposes. Orange-Yellow corresponds to uniform positional patterns. Pink-Purple corresponds to slow positional decay. And Blue corresponds to local positional patterns.

For reference, compare the first column of the GPT2-Small heatmap with the plots of the positional patterns given in Section 3. Heads 3, 4, and 7 are in blue because they are local positional patterns. Heads 0,1,2,6,8,9,10 are in magenta as they have a slow positional decay. And Heads 5 and 11 are yellow because they are close to uniform.

We can visually observe interesting things about the layers of GPT2-Small. For instance, notice the density of local positional patterns in the third and fourth layers. Potentially this is from the model extracting local grammatical structure from the text.

On the other hand, the second layer has more slow positional decay / uniform positional patterns. In fact, on closer inspection, the second layer has lots of attention heads which act purely as fixed positional kernels, falling into the category of "contextual attention heads" discussed earlier. This suggests the model builds an initial linear summary of the surrounding text, and then begins to build more symbolic representations in the third and fourth layer. We observe a similar pattern in GPT2-Medium, and to some extent in GPT2-XL.

Layer 5 is known for having many induction heads: Heads 5.0, 5.1, and 5.5 are known to be induction heads. These stick out visually as having close to uniform positional patterns, which validates the intuition that induction heads tend not to care about position.

The fact that there are so many local positional patterns in layers 2-4 gives a potential explanation for the small number of interesting specialized heads found in these layers. Attention heads with uniform positional kernels like induction heads feel more likely to be selected for "interesting behaviour", than heads which attend only locally.

## Conclusion

It seems like positional kernels are a useful notion to look at when first assessing attention heads, and they suggest many different lines of inquiry. One interesting piece of future work could be looking at how these positional kernels develop over the course of training.

However, the assumption made at the start of the post has not been validated, and it'd be important to look at this in future work.

This  contains the code required to reproduce the results found here.

## Appendix

### Alternative positional schemes

There are many different positional schemes used for different LLMs, but the three most common alternatives to additive positional embeddings are:

- T5: Adds a learnt relative bias to the pre-softmax attention scores based on the relative distance between the positions of the current token and the token being attended to.
- ALiBi: Similar to T5, adds a linear bias to the pre-softmax attention scores proportional to the relative distance between the current token and the token being attended to. The slope of this bias is hard-coded, not learnt.
- RoPE: Applies a complex rotation of `θ = 10000^{\frac{-2i}{rot_{dim}}}` to the `2i`th and `(2i+1)`th dimensions of the query and key vectors, up to the `rot_{dim}`th index. `rot_{dim}` varies depending on implementation: some implementations take `rot_{dim} = d_{head}`, where `d_{head}` is the dimension of the query and key vectors. The most common model which uses RoPE is the Pythia line of models, which uses `rot_{dim} = \frac{d_{head}}{4}` for efficiency purposes.

T5 and ALiBi are pretty straightforward to handle under this framework. The relative bias terms can simply be added to the EQKP+PQKP terms before taking the softmax to compute the positional kernel. We can't assume that `P` is constant, because models are still able to develop emergent positional embeddings, as we see in NoPE models, which they may use to overcome the limitations of these positional schemes.

RoPE is more difficult to work with. The main trouble with analyzing RoPE is that while it allows for relative positional indexing, the rotations of queries and keys can potentially disrupt semantic attention (as noted in Round and Round we go). Round and Round we Go suggested p-RoPE as a solution to this problem, where only a percentage `p` of the highest frequency query and key dimensions are rotated, with the remaining dimensions left without rotation. They found that taking `p=0.75`, models perform better on long-context tasks than for `p=1`. Interestingly, they found that `p=0.25`, corresponding to Pythia, performs worse than `p=1`. The analysis we give here is based on empirical results from Pythia models, but the ideas should apply for `p<1` p-RoPE models.

The motivation for studying p-RoPE models is that these models can perform better than vanilla RoPE models, and they have concrete semantic channels given by the non-rotating dimensions, allowing us to attempt a separation of content and position.

One might hope that p-RoPE models would use the non-rotating dimensions exclusively as semantic channels, and it seems like they do this. But, they also use the rotating key dimensions in non-trivial ways that seem like they would require a project of their own.

For now, note simply that if the rotating key bias dominated over the keys, making the rotating key dimensions close to constant, we could recover a positional kernel. Some attention heads seem like they do this but it's not that clean.
```