Decomposing attention into RoPE frequencies

Can we investigate how different parts of an LLM pay attention to high and low frequency parts of text?

Recently, I’ve been reading the Transformer Circuits blogs, and I’ve found them really cool. Philosophically, I find the field of interpretability really exciting, as having these models that we can easily probe and decompose helps us get a better understanding of how intelligence works. I also have a hope that interpretability work could help guide future machine learning architectures: we can try and understand how current LLMs perform their computations, and identify missing pieces that could make future computations a bit easier.

I also think that the field of ML is easier than ever before to get a splash into! One of the issues that used to be really frustrating in ML is the amount of time that one would need to spend on “plumbing” to even get started answering a basic research question; I remember the painful days of dealing with CUDA compatibility issues between who knows how many different libraries, or the frustration in working out how to load and clean a variety of different datasets. Now, the surrounding ecosystem is much more mature: I can pay a tiny amount of money and get access to a pretty solid GPU, with a tech stack that is guaranteed to work with it. I can ask AI to do the boring work of working out how to deal with datasets, or to extract a specific part of a model.

This is my first article (and one of my first personal projects) performed with AI assistance. In this article:

• all ideas are my own unless otherwise attributed
• all text + math is my own writing, but proofread by AI
• code + graphics are AI produced, with significant direction on structure from me
  Who knows what the future of thinking work will look like!

“Abstract”

If I had to, in a few lines, explain the mildly exciting things out of this article, they would be:

1. We have found simple metrics that let us ablate large portions of the attention computation, with relatively little model degradation
2. We have shown that the RoPE frequencies that are over/underutilized change heavily throughout the model’s layers, with some layers paying attention to local information and other layers only focusing on global information
3. We have shown that we can use these metrics to pick RoPE frequencies which improve model performance and decrease model training time by 4%
4. We have some human interpretable results on how ablating frequency ranges in different parts affect the performance of the LLM on different tasks, and use this to get some fuzzy idea of what the frequency bands are doing.

Hopefully that convinces you enough to read on!

Background and motivation

My initial pie-in-the-sky idea was that language is the kind of thing that can be split into frequency bands; the “low frequency” information describes the overall meaning and direction of a piece of text, which holds over a longer period of time, while “high frequency” information describes the precise sentence substructure.

A natural place to start looking is in the RoPE used in modern LLMs. This is a clever way of encoding positional information in a way that is closely coupled with the attention mechanism. To recap, RoPE is the following idea: when computing the attention between a \(d\)-dimensional query vector \(q\) and key vector \(k\), do the following:

• Pair up components of the query vector into \(d/2\) pairs.
• Interpreting the pairs as 2D vectors, rotate each pair \(i\) in the query vector by \(\theta_{i} \times p\), where \(p\) is the index/position of the query vector in the text.
• Do the same rotation of pairs for the key vector
• Compute attention on the rotated query and key vectors.

The idea is that the model can perform both “high frequency” attention by putting information into the parts of the vector that rotate quickly, and “low frequency” attention by putting information into the parts of the vector that rotate slowly.

In this article, we explore the attention mechanism on the level of these bands, including:

• how we can break down the attention given to a single (query, key) pair into different frequency components
• if different attention heads in the model attend to different frequencies (spoiler alert, they do)
• if we can ablate frequency bands that the model doesn’t seem to be utilizing
• whether adding capacity to saturated frequency bands helps the model overall
• how we can interfere with different frequency bands in the model, and see how this affects text generation

I hope that all of these various investigations shed just a tiny extra glimmer of light on the emergent behaviours of attention in LLMs!

Investigations

Decomposing the attention given to (query, key) pairs

First, I wanted to be able to try decomposing the attention to an individual (query, key) pair. For example, I wanted to be able to give an interpretable result along the lines of “the query for token X attended to token Y primarily through high/low frequencies”. But how do we say which frequency bands “caused” a token to be paid attention to?

Here’s the problem that we’re faced with. We have a single query vector \(q\) and \(N\) key vectors, \(k_{i}\). Each is \(d\)-dimensional, where each dimension corresponds to a different frequency band. We refer to the \(j\)th dimension of key \(i\) as \(k_{ij}\). Additionally, let \(x_{ij} = q_{j}k_{ij}\). Then, the dot product of our query and key vector is just \(\sum_{j} q_{j} k_{ij} = \sum_{j} x_{ij}\).

We can already find how much key \(i\) is attended to: this is just the standard softmax:

One way of thinking about this is to ask: how did \(a(i)\) come to be its value of \(a(i)\), and not the value from a uniform distribution (\(\frac{1}{N}\))? Indeed, if all our keys were 0, then the attention would be a uniform distribution, but deviations from all of our keys being 0 is the reason that attention is non-uniform.

Let’s define a vector \(s\) that we’ll use as a scale factor on each of the individual dimensions. Then, define

Let’s think about how this could look, and then formalize it into a working algorithm. Let’s keep track of a running total \(c_{ij}\), indicating how much dimension \(j\) contributed to the attention paid to key \(i\). Initially, when we start with \(s = 0\), all of our keys get \(\frac{1}{N}\) attention and we can equally attribute this to all of the dimensions by symmetry, so set \(c_{ij} = \frac{1}{Nd}\). Next, imagine that we nudge \(s_{j}\) for some fixed \(j\) by some \(\delta s\). This will cause \(a(i; s)\) to change by some \(\delta a_{i}\). We can attribute this change fully to dimension \(j\), and thus we can increment \(c_{ij}\) by \(\delta a_{i}\) for all \(i\). We can then nudge a different \(s_{j}\). Doing this in a round robin fashion until \(s = 1\), we will have our final set of \(c_{ij}\)s: an attribution of the attention paid to key \(i\) split among the different dimensions \(j\).

Note that the round robin part is load bearing. We could, for example, do all of our nudges on \(s_{0}\) from \(0\) to \(1\), first, then nudge \(s_{1}\) from \(0\) to \(1\). However, this introduces asymmetry between the attribution to different components, so the round robin framing is more natural.

Now, obviously all of this nudging and summing should scream some nice way to write this idea down in calculus, and so here it is:

Let’s try it! Here’s the attribution.
Experiments were performed on a LLama 3.1 8B model. We use 32 integration steps with Gauss-Legendre quadrature.

We can see some interesting things here! For example, take a look at head 2 in layer 1, which only seems to attend to the comma (“,”) and articles/prepositions (“a”, “to”). All of this attention is mostly in the high frequency space, and we can see how it diminishes as query tokens further away attempt to attend to these key tokens. On the other hand, head 1 in layer 10 seems to do most of its per token negative attribution in the low frequency space (maybe the higher frequencies could be used on a longer example).

This is a bit hard to interpret though. Most of the contributions graphs are quite spiky looking, which makes sense, as there’s no reason for two dimensions at similar frequencies to be communicating the same information. Nevertheless, looking through this data, it seems like some heads pay attention to different frequency components. Can we maybe formalize this idea?

Decomposing the attention’s entropy

What we developed above isn’t the right tool to ask over the whole attention matrix, “how much of this attention distribution is due to high vs low frequency bands”? Indeed, if you try summing \(\sum_{i} a'(i, j)\), you will find that you just get \(\frac{1}{d}\). This kind of makes sense: all frequency bands get the same opportunity to contribute to the final attention total of \(1\), and nudging any individual band’s \(s_{j}\) does not affect the attention total. The intuition is more that a frequency band with high variance in its contributions is one that “shapes” the underlying attention distribution more.

If you think about this a bit, a more natural question is how can we attribute the entropy of the attention distribution to different frequency bands? We can imagine a model that isn’t able to pay attention to anything having a uniform entropy (\(\log N\)) over the attention distribution. However, when we have a finishing entropy of \(S\), we can interpret this as the model saying “I originally didn’t know what to pay attention to, but I was able to extract \(\log N - S\) bits of information out of this distribution; this is the amount of information contained in my understanding of where the attention should go”.

Let’s formalize this! The entropy of the attention is

One question to ask is, do different parts of the model attribute their attention distribution to different frequency bands? To do this, we go through 10,000,000 tokens worth of text from a web dataset, and for each head in each layer, we keep track of the total amount of entropy coming from each band. Note that in my analysis, I only looked at text samples with > 2048 tokens, and I clipped each sample to 2048 tokens. The tiny size is more just a consequence of the small, single GPU I am using, but note that this means some of the results in the mid band of frequencies will not be properly explored. The results should still hold on high frequencies and very low frequencies (which don’t meaningfully change in a single context window).

Experiments were again performed on a LLama 3.1 8B model. We use 32 integration steps with Gauss-Legendre quadrature.

Take a look at the results below!

There are many interesting observations that we can make from this data:

• Different layers of the network attend to different frequency bands. For example, across all heads in layer 1, the attention is quite evenly spread across most frequency bands, but in middle layers, we get much lower entropy attention distributions, primarily attributed to lower frequencies. High frequency contributions only slightly pick up again in the very last layer.
  • We can intuitively see why this could be the case. Earlier layers of the model might need to pick up on local information in the text, e.g. to connect adjectives with the nouns that they refer to. Middle layers of the model could be understanding the overall structure of the text, and final layers could need to perform higher frequency processing again in order to get the sentence structure of the prediction correct.
• Many heads are dominated by just a few frequency bands. For example, look at head 1 in layer 10. This is attributing most of its entropy to a single peak frequency.

Can we ablate out some frequency bands without harming model performance?

The entropy data above hints that individual attention heads, and individual layers, don’t use all frequency bands equally. Maybe, models could simply avoid computing pre-attention logits for some frequency bands? But how should we decide which frequency bands to remove? After all, entropy is quite a local quantity, and there are many ways in which our entropy attribution could be a poor proxy for overall importance of individual frequency bands in the final model.

To test our ideas, we try a couple of different strategies to guide which frequency bands we should remove:

• (Control baseline): removing random frequency bands
• Removing bands by their total activation
• Removing bands by how much they contribute to entropy, in the above decomposition
• Removing bands by their Fisher information

The general structure of each of these removal methods is to assign each band an “importance score”. Then, we run the following experiment:

1. For a variety of fractions \(f\):
   1. Select the lowest \(f\%\) of bands for ablation. Remove these bands from the model (i.e. zero out these bands from the query matrices)
   2. On a couple (~300k) FineWeb tokens, compute the total NLL predicted by the model.

Our study then boils down to how we choose to score frequency bands:
Random: simply assign frequency bands random scores.
Total activation: Across some training tokens, sum \(|q_{*b}k_{*b}|\) for each frequency band \(b\): in other words, find the total contribution of each frequency band to the pre-softmax logits. The importance of a band \(b\) is the ratio of this total activation over the maximum total activation of all bands in the given head. Doing this in a memory efficient manner is a little tricky with pure torch primitives, so we write a custom kernel for this (methodology in an appendix).
Entropy: We use the total entropy attribution from above. The importance of a band \(b\) is the ratio of the absolute value of the entropy attribution of this band over the maximum entropy attribution of all bands in the head.
Fisher information: this one requires a little bit more exposition, so I’ve put the methodology in an appendix. TL;DR though, we try and investigate how sensitive the model’s output distribution is to each frequency band of each attention head, and use this sensitivity. We say that the importance of each band is the ratio of the sensitivity of a band, divided by the maximum sensitivity of a band in the given head.

Take a look at the results!

This is pretty solid evidence that the total activation is the best metric for determining importance of individual bands (which is nice, as it’s also the easiest to compute!): we can remove 40% of the bands of the model at a cost of 0.25 nats of per token NLL. It also turns out that activation importance is (somewhat unsurprisingly) decently correlated with entropy importance (a linear model gets an R^2 of 0.877), so these metrics are related, but empirically, the best for determining end model performance is total activation.

Can we use this information to train better models?

Here’s where we take a bit of a stab in the dark. I had a hypothesis that information was somehow bottlenecked by the specific RoPE frequencies available.

Going back to the entropy interpretation, we see that our middle layers do not attribute much of their attention entropy to high frequency bands. We know from other work that LLMs typically develop middle workspace layers, which already have a decent high level understanding of the text at hand (or the task at hand, in a model fine tuned to act as an agent). These layers most likely do not need to use high frequency bands, as earlier layers have most likely used their higher frequency bands to connect together the meanings of nearby words. This means that these middle layers, which would normally have a full 64 dimensional space in which to project queries and keys, can only reasonably access a smaller dimensional space formed by the lower frequencies.

We can envision a world where we want a bit more of a uniform entropy attribution across bands per head. This means that more of the bands are contributing to the final entropy attribution, and thus potentially more of the bands are carrying meaningful information.

We try testing this hypothesis by using the total activation statistics as a proxy to determine what frequency bands should be present in the final model, and we see how changing these affects pretraining. The exact methodology we use is as follows:

1. Pretrain a “bootstrap” LLM with a standard ladder of RoPE frequencies
2. Use our total activation metric to determine, for each frequency band of each head of each layer, the total activation present; so \(a(f, h, l)\) for each frequency band, head and layer.
3. Next, use the total activations to create an improved set of RoPE frequencies for a new model
   1. Clamp the measured total activations to the 5th and 95th percentiles: call this \(\hat{a}(f, h, l)\).
   2. Across each layer, aggregate the total activation across all heads, to produce a metric of total activation of each frequency band of each layer.
   3. Next for our new model, we sample the RoPE frequencies for each layer with a probability proportional to the total activation of each frequency band; formally, the probability of a frequency band in layer \(l\) getting the frequency \(f\) is \(\left( \sum_{h} \hat{a}(f, h, l) \right) / \left( \sum_{h, f} \hat{a}(f, h, l) \right)\). The motivation behind this is that the expected number of bands of frequency \(f\) is proportional to its activation, and this means that if activation is split equally among the bands, each band will get in expectation, an about equal amount of activation
4. Train the LLM with the improved RoPE frequencies
5. Train an LLM with the original RoPE frequencies, with the same seed as in step 4.

We repeat steps 2-5 a few times to get reproducible results, using the LLM trained in step 5 as the bootstrap LLM for the next iteration. Note that we are only training a 100M parameter model with the Chinchilla token count (~2B tokens), which only costs me about \(\$3.50\) to train each model.

The model architecture that we use has a GPT2 tokenizer (50k vocab size) with tied embeddings (embedding and unembedding matrices are transposes of each other). The model has 10 layers, each with 12 attention heads. The hidden dimension is 768, and FFN layers use SwiGLU with a 2048 hidden dimension. The model was trained with a 6e-4 learning rate with an AdamW optimizer, providing weight decay on all weights except the embeddings.
Take a look at the results!

Our models with the sampled RoPE frequencies consistently perform better than models with the default RoPE frequencies, on all seeds. The improvement on the NLL at the end of the training process is relatively small, of on average 0.0089 nats per token. Overall, this results in an average 4.0% reduction in training time, which is pretty neat! These are also super similar results to the group behind LeRoPE who tried making RoPE frequencies learnable parameters, but used the same frequencies throughout the whole model.

However, note some caveats with these results:

• We are only training on really tiny LLMs. It’s possible that behaviour on larger models is pretty different.
• The improvement, while consistent, is at the moment within the seed to seed variance.
• We are also not getting to train to convergence, so we can’t really say if the extra capacity in these bands is needed, or if this just makes training more efficient.

Let’s dive a little deeper into our results. First, let’s see what our sampled RoPE frequencies actually were:

Our sampled RoPE frequencies are able to increase the number of bands in a mid-low frequency range, while reducing the number of bands in the high frequency range, without too much variation across our repeats. Interestingly, this differs from learned frequencies in prior work; the team behind the LeRoPE paper found that high frequency bands were basically completely unchanged, and low frequencies were generally pushed even higher. Interestingly, where this reweighting happens changes by layer.
I was too busy to make a nice graphic for this, I’ll fix it soon!

We can see that later layers tend to get more lower frequencies, while the first layer gets a slight bias for higher frequencies.

Now, let’s investigate whether or not these extra channels in certain frequencies actually allowed the model to more evenly distribute its attention across its channels. We look at how the total activation is different in the models with the original RoPE frequencies, against the models with the sampled RoPE frequencies. We use the importance scores from earlier, and sum these across all heads in a layer to get a per-layer, per-channel importance score.

Nice! From these results, we observe that our distributions of activations are pretty consistent across frequency bands, and that the models with the sampled RoPE frequencies are able to more evenly distribute their attention across frequency bands.

Bonus: can we interpret what frequency bands are broadly doing?

The neuroscientists of ye olde trying to study brain structure used brain lesions as a primitive, low tech way to study what different parts of the brain do. Our synthetic brain unfortunately doesn’t have this nice spatial nature that we can exploit, but it is really easy to artificially inject “lesions”! So to wrap up, we qualitatively look at how models (we use LLama 3.1 8B Instruct instead of the base models this time) perform on tasks when noise is injected in different locations.

To do this, we inject noise into either the 12 highest frequencies or the 12 lowest frequencies of the model, one layer at a time, to create a “lesioned” model. The layers we inject into are:

• layers 0, 1, 2, 3 (hopefully sensory layers, which might interpret the text)
• layers 23, 24, 25, 26 (hopefully workspace layers (supported by some prior work), which could allow for more complex reasoning and refinement of the initial understanding)
• layers 28, 29, 30, 31 (hopefully motor layers, which could be responsible for producing a coherent textual output)
  We iteratively increase the variance of the noise we are injecting, until the model begins to produce too nonsensical of an output. We determine how nonsensical the output is by measuring the mean per-token NLL of the lesioned model’s output, where we compute the log likelihood under the base model’s probability distribution on the same sequence. We save all of the outputs, and then we look at them to try and see if we can glean out any qualitative patterns in the results.

For reference, here are the prompts for the tasks that the model was asked to complete. I first generated this fake story, which we call WONDERLAND (full text in an appendix).

The model was then asked to complete the three following tasks:

• Repetition: "Repeat the following text:\n" + WONDERLAND
• Continuation: "Continue the following text:\n" + WONDERLAND
• Coding: Write a Python function extract(cfg) that reads this nested structure and returns the list of enabled plugin names: cfg = {"services": {"web": {"plugins": [{"name": "auth", "enabled": True}, {"name": "cache", "enabled": False}]}, "db": {"plugins": []}}}

Once again, like the neuroscientists of ye olde, we also ask the models to introspect on their own thoughts after they respond by asking, Does your previous output make sense to you?.

We put all of the runs that we performed below for you to peruse!

This is a lot of data to qualitatively make sense of! As I was reading the results, I made a bunch of notes, tucked away in an appendix. I then summarized these notes to produce this picture of the model, mostly ordered by layer.

• Sensory layers (0-3)
  • Layer 0 is relatively more resistant to damage.
  • Otherwise, damage to both high and low frequency components look similar; the model quickly fails to understand the task at hand. The model also is not aware of the fact that it is unable to understand the task, and is convinced that it is operating correctly.
• Workspace layers (23-26)
  • Low frequency bands seem to relate to the overall meaning of the input. Interventions here can cause nouns to be replaced with other versions of the noun consistently through the text, like character names being misspelled throughout the whole text.
  • High frequency bands, on the other hand, may seem to carry some style information, whether that is in the output format, or the style of prose in the continuation (it seems to have no impact on the repetition task)
  • Errors in these bands can be noticed by the model and corrected. This echoes other work where models were trained to remove weights used in certain tasks, but they were able to circumvent this damage with prompting. In other words, LLM brains have multiple redundant pathways for the processing of the same information.
• Motor layers (28-31)
  • High frequency interventions do almost nothing, but low frequency interventions create more random typos throughout the text, which the model can recognize and correct. These differ from the workspace layer typos, which will replicate a typo throughout the whole text.

One important caveat: these results are tenuous at best. These are just one person looking through all of the responses in the view above, and trying to spot patterns. However, it slightly nudges me to build this creaky mental model of LLMs that looks like this:

• Sensory layers use all frequency bands to make sense of the text. If this area is damaged, the model can output not much past incoherent babble that it doesn’t recognize as babble (similar to damage in our Wernicke’s area!)
• Workspace layers might store style information in areas attended to by high frequency bands, and definitely store the higher level meaning of the text (the constituent characters in a story or overarching logic in a program) in these layers. Errors in these areas can be detected and corrected.
• Motor layers seem to have almost all information in low frequency bands, and these contain the final identities of the tokens to be outputted. Errors in these areas can also be detected and corrected. (I secretly hoped that these could be seen as analogous to Broca’s area, but the sentences seem to always remain grammatically correct…)

Conclusion

I’m pretty happy with how this work turned out! I think the thing that is most valuable that I will personally take away from this is a slightly sharper picture on how the circuitry of am LLM works. Honestly, I don’t think there is much new here past the already well documented division of the role of layers in an LLM, but the contextualization into high and low frequency bands allows the work of these layers to be further subdivided. If I really only got one thing from this, it would be that high frequency positional information doesn’t seem to be that relevant after the early layers. It does seem like high frequency information is used only in early layers to connect nearby words, but once the overall meaning of the text is extracted, most “work” seems to be done in spaces that are attended to by lower frequency bands.

There’s only so much I can fit into one article, but some things I want to try include:

• If I had more compute, I would extend this work to longer context windows. I think some of the big problems with attention are with these massive, 100k+ context windows, which is why many modern models use NoPE in later layers (also kind of what this work has rediscovered).
• Following this understanding that most information lives in this low frequency space without positional information, I want to dive deeper into understanding how this translates to the production of actual characters; if low frequency information can be interpreted as semantic and high level, where does the information of what the word before this token was live?
  • To expand a bit more on this, two adjacent tokens have pretty much the same high level understanding of the text in their hidden activations, but still end up getting mapped to very different next character predictions. Where does this difference between “high level understanding” and “next character” live?
• Trying to see if there’s a way to automatically split a network up into “areas” that we can get very different results on after ablation. We got frequency bands as a nice, interpretable split of the attention mechanism from RoPE, but I’m sure that there are many other splits that could be interesting to investigate.
  • I like how our neuroscientists get to play with a physical brain and ablate out parts living in a 3D space. We are kind of restricted to ablations that are baked into a hardcoded model design.

Any one of these could be an article in the future, so watch this space! I really enjoyed this and I’m sure that (time permitting) I’ll be doing lots more in the future :)
(and hopefully I don’t get beaten to the punch on some of these cool ideas ;) )

Appendices

Appendix: memory efficient total activation

As is usual with pretty much any attention based work, we’re primarily trying to avoid materializing the full attention matrix, which will take \(O(T^2)\) space, leading to us at best wasting a lot of bandwidth in memory writes and at worst OOM-ing on the GPU.

If we didn’t have the absolute value problem, this computation is super easy, and can be pretty trivially done in \(O(N)\): consider a single band \(b\). We would just be trying to sum

However, we in fact need to compute the absolute values of each of these products, which is a bit trickier.

@triton.jit
def attention_activations_kernel(
    q1_ptr, q2_ptr, k1_ptr, k2_ptr, out_ptr,
    # stride arguments
    T, D,
    BLOCK_Q: tl.constexpr,
    BLOCK_D: tl.constexpr,
):
    b = tl.program_id(0)
    h = tl.program_id(1)
    qb = tl.program_id(2)

    # fold b, h into the base -> everything below is a 2D [T, D] problem
    bh = b * stride_b + h * stride_h
    q_lo = qb * BLOCK_Q

    q1_bp = tl.make_block_ptr(...
        offsets=(q_lo, 0), block_shape=(BLOCK_Q, BLOCK_D)
    )
    q2_bp = tl.make_block_ptr(...
        offsets=(q_lo, 0), block_shape=(BLOCK_Q, BLOCK_D)
    )

    q1t = tl.load(q1_bp, boundary_check=(0, 1), padding_option="zero").to(tl.float32)
    q2t = tl.load(q2_bp, boundary_check=(0, 1), padding_option="zero").to(tl.float32)

    offs_q = q_lo + tl.arange(0, BLOCK_Q)
    row_base = bh + offs_d * stride_d

    acc = tl.zeros((BLOCK_Q, BLOCK_D), dtype=tl.float32)

    k_hi = tl.minimum(q_lo + BLOCK_Q, T)
    for k in range(0, k_hi):
        k_off = row_base + k * stride_t
        k1r = tl.load(k1_ptr + k_off, ..., other=0.0).to(tl.float32)
        k2r = tl.load(k2_ptr + k_off, ..., other=0.0).to(tl.float32)

        p = q1t * k1r[None, :] + q2t * k2r[None, :]
        acc += tl.where((offs_q >= k)[:, None], tl.abs(p), 0.0)

    res = tl.sum(acc, axis=0)  # [BLOCK_D] -- q is summed away entirely
    
    tl.store(out_ptr + out_off, res, mask=d_mask)

TL;DR, we simply have each CTA own a block of queries. Each CTA then loops over the keys that are casually relevant to this query, computes the product, and accumulates it in registers. The CTA then sums along the query dimension, returning an intermediate result. The launcher then sums all intermediate results from all CTAs to get a final number.

Appendix: Fisher information metric

Recall that the pre-attention logit for a (query, key) pair is \(q^Tk\), for a query and key vector. We create a vector \(\theta\) per attention head (this is similar in role to our \(s\) vector) and compute attention logits as \((q \odot \theta)^Tk\) in our model. Of course, if we set \(\theta\) to ones, we recover the original attention logits. Next, we can try computing the second derivative of the loss with respect to \(\theta\). This is possible for LLMs as all functions do have meaningful second derivatives! However, the full Hessian is enormous, so we restrict our attention to the diagonal of this Hessian, i.e. the elements \(\frac{d^2L}{d\theta_{i}^2}\) for each \(i\). This is a bit difficult to analyse, so instead of looking at the true loss, we assume that the true distribution is perfectly predicted by the model at \(\hat{\theta} = [1, 1, \dots, 1]\). Let \(f(s, \theta)\) denote the probability distribution of the token following tokens \(s\) at a given theta, we let the cross entropy \(h(\theta) = H(f(s, \hat{\theta}), f(s, \theta))\), and we try and determine \(\frac{d^2h(\theta)}{d\theta_{i}^2}\). This is the same as our original quantity if our model is accurately predicting the probability distribution of the next token, which for a trained model will not be too far from the truth. We can recognize this quantity as the Fisher information, which can also be computed as \(\mathbb{E}\left[ \left( \frac{d}{d\theta_{i}} \log f(s, \theta)[t] \right)^2 \right]\), where \(f(s, \theta)[t]\) is the probability determined by the model of seeing token \(t\) next, and the expectation is taken over sampling \(t\) from the distribution \(f(s, \hat{\theta})\). As the model predicts an independent distribution for each token, we can also observe that

1. Select a document, and use the model to produce, for each prefix of the token sequence, the probability distribution for the next token, \(f(s_{1}, \hat{\theta}), f(s_{2}, \hat{\theta}), \dots\)
2. Sample a token for each position in the sequence from the computed probability distributions to get \(\hat{t_{1}}, \hat{t_{2}}, \dots\).
3. Compute \(\sum_{j} \log f(s_{j}, \theta)[\hat{t_{j}}]\) at \(\theta = \hat{\theta}\), and use PyTorch’s backpropagation to compute \(\frac{d}{d\theta_{i}} \sum_{j} \log f(s_{j}, \theta)[\hat{t_{j}}]\) for all \(\theta_{i}\) in a single backwards pass.
4. Accumulate \(\left( \frac{d}{d\theta_{i}} \sum_{j} \log f(s_{j}, \theta)[t_{j}] \right)^2\) across several documents. We use this quantity as our importance score for the Fisher information metric.

Appendix (rough): notes as I read the “lesioned” outputs

• Repetition task
  • Low frequency bands:
    • Sensory layers:
      • Low frequency intervention in layer 0 surprisingly is relatively less damaging than in other sensory layers, which could be explained by the purpose of this layer being to connect words close together, in order for later layers to understand what is actually going on.
      • However, in layers 1-3, the model very quickly degrades. First, it starts producing a grammatically coherent response that fails to acknowledge the task. In one run, the model begins to summarize the text. The model is unaware of its degradation when asked to introspect. Eventually, the model devolves into repeating nonsense.
    • Workspace layers:
      • Relatively little degradation, even on large magnitudes of variance injection. Interestingly, some seeds (which presumably place noise on certain weights) seem to produce consistent errors in the reproduction, such as seed 1 of layer 25 replacing “Willoughby” with “Will 0” and losing “door” in “door-sized”. This could mean that some of the heads that were intervened in caused information on what is the original prompt to be lost. Interestingly, here the model is able to recognize its mistake and can even correct it.
    • Motor layers:
      • The model begins to produce some random typos with some seeds, like introducing in numbers before paragraphs, or misspelling words. The model is almost always able to recognize that it has made a mistake, but continues to make similar minor errors. Surprisingly, the errors are consistent for given seeds, which maybe implies that the model is, through the attention mechanism, able to give certain generation positions high confidence that the generated token should be something else.
  • High frequency bands:
    • Sensory bands:
      • High frequency intervention in layer 0 is again less damaging than in other sensory layers, which is evidence to refute the hypothesis that layer 0 is primarily dependent on high frequency information.
      • The model degrades much more quickly and starts performing the wrong task. The degradation pattern is otherwise similar to the low frequency bands.
    • Workspace layers
      • There are far fewer consistent typos. Often times, the high frequency degradation does almost nothing.
    • Motor layers:
      • Even fewer typos from this intervention, compared to the low frequency case! Once again, these typos are quickly caught.
• Continuation task
  • Low frequency intervention: We see similar results to the repetition task in the sensory and motor layers. In the workspace layer, at least in this nonsensical story, all outputs seem to make sense! There is a noticeable change in style though; the model output begins to feel less coherent.
  • High frequency intervention: In workspace layers, the original direction may be preserved better; the model seems to prefer its initial, unintervened generation slightly more. However, it seems to do a slightly worse job at reflecting the original style. The motor layers do not seem to produce any typos in this case.
• Coding task:
  • This might not have been the best task.
  • In the workspace layers, low frequency interventions seem to affect logic, while high frequency interventions seem to affect style. For example, high frequency interventions can get the model to always insert headers into its responses, while low frequency interventions can get the model to more commonly try recursive approaches or make more logical errors.

Appendix (rough): WONDERLAND text

Prudence Willoughby-Fane was seven and three-quarters, and she had strong opinions about Tuesdays. Tuesdays, she felt, were the least trustworthy of days — not honest like Monday, not cheerful like Friday, but sly, and prone to hiding things behind the umbrella stand. It was on just such a Tuesday, while looking for her missing left shoe, that she found the Thimble Door.

It was not, properly speaking, door-sized. It was thimble-sized, set into the skirting board behind the umbrella stand, painted the exact shade of brass that makes grown-ups say “how curious” and then immediately forget about it. Prudence, who had never once forgotten a curious thing in her life, knocked on it with her smallest fingernail.

“Password,” said the door, in a voice like a kettle deciding whether to boil.

“I don’t have one,” said Prudence. “I only just found you.”

“That,” said the door, “is the most common password there is. Do come in, but mind you come in sideways — everyone insists on trying to come in forwards, and it never works, and then they blame the door, which is terribly rude.”

Prudence came in sideways. It worked at once.

She found herself in a hallway that ran uphill in both directions, wallpapered in a pattern of teacups that watched her walk past with mild interest. At the far end sat a creature that was mostly hat, with a small amount of badger arranged underneath it for support.

“You’re new,” said the Hat-Badger, without looking up from the letter it was writing to itself. “I can always tell. New people walk as though the floor might change its mind, which it frequently does, so that’s actually quite sensible of you.”

“I’m looking for my shoe,” said Prudence.

“Everyone’s looking for something,” said the Hat-Badger. “I’m looking for the end of this sentence, which keeps moving whenever I get close to it. It’s exhausting. Would you like to help me chase it? Two of us have a much better chance, statistically, of cornering a sentence.”

“How does one corner a sentence?”

“Very carefully, and preferably near a comma, where they like to stop and catch their breath.” The Hat-Badger stood, and several letters fell out of the hat’s brim, all addressed to Whom It May Concern, which seemed to concern nobody at all. “Come on, then. It went this way, or possibly that way, or — oh, bother, it’s gone and become a question.