4️⃣ Reverse-engineering induction circuits

Learning Objectives

• Understand the difference between investigating a circuit by looking at activation patterns, and reverse-engineering a circuit by looking directly at the weights
• Use the factored matrix class to inspect the QK and OV circuits within an induction circuit
• Perform further exploration of induction circuits: composition scores, and targeted ablations

In previous exercises, we looked at the attention patterns and attributions of attention heads to try and identify which ones were important in the induction circuit. This might be a good way to get a feel for the circuit, but it's not a very rigorous way to understand it. It would be better described as feature analysis, where we observe that a particular head seems to be performing some task on a certain class of inputs, without identifying why it does so.

Now we're going to do some more rigorous mechanistic analysis - digging into the weights and using them to reverse engineer the induction head algorithm and verify that it is really doing what we think it is.

Refresher - the induction circuit

Before we get into the meat of this section, let's refresh the results we've gotten so far from investigating induction heads. We've found:

• When fed repeated sequences of tokens, heads 1.4 and 1.10 have the characteristic induction head attention pattern of a diagonal stripe with offset seq_len - 1.
  • We saw this both from the CircuitsVis results, and from the fact that these heads had high induction scores by our chosen metric (with all other heads having much lower scores).
• We also saw that head 0.7 strongly attends to the previous token in the sequence (even on non-repeated sequences).
• We performed logit attribution on the model, and found that the values written to the residual stream by heads 1.4 and 1.10 were both important for getting us correct predictions in the second half of the sequence.
• We performed zero-ablation on the model, and found that heads 0.7, 1.4 and 1.10 all resulted in a large accuracy degradation on the repeated sequence task when they were ablated.

Based on all these observations, try and summarise the induction circuit and how it works, in your own words. You should try and link your explanation to the QK and OV circuits for particular heads, and describe what type (or types) of attention head composition are taking place.

You can use the dropdown below to check your understanding.

Refresher - QK and OV circuits

Before we start, a brief terminology note. I'll refer to weight matrices for a particular layer and head using superscript notation, e.g. $W_Q^{1.4}$ is the query matrix for the 4th head in layer 1, and it has shape [d_model, d_head] (remember that we multiply with weight matrices on the right). Similarly, attention patterns will be denoted $A^{1.4}$ (remember that these are activations, not parameters, since they're given by the formula $A^h = x W_{QK}^h x^T$, where $x$ is the residual stream (with shape [seq_len, d_model]).

As a shorthand, I'll often have $A$ denote the one-hot encoding of token A (i.e. the vector with zeros everywhere except a one at the index of A), so $A^T W_E$ is the embedding vector for A.

Lastly, I'll refer to special matrix products as follows:

• $W_{OV}^{h} := W_V^{h}W_O^{h}$ is the OV circuit for head $h$, and $W_E W_{OV}^h W_U$ is the full OV circuit.
• $W_{QK}^h := W_Q^h (W_K^h)^T$ is the QK circuit for head $h$, and $W_E W_{QK}^h W_E^T$ is the full QK circuit.

Note that the order of these matrices are slightly different from the Mathematical Frameworks paper - this is a consequence of the way TransformerLens stores its weight matrices.

Question - what is the interpretation of each of the following matrices?

There are quite a lot of questions here, but they are conceptually important. If you're confused, you might want to read the answers to the first few questions and then try the later ones.

In your answers, you should describe the type of input it takes, and what the outputs represent.

$W_{OV}^{h}$

$W_E W_{OV}^h W_U$

$W_{QK}^{h}$

$W_E W_{QK}^h W_E^T$

$W_{pos} W_{QK}^h W_{pos}^T$

$W_E W_{OV}^{h_1} W_{QK}^{h_2} W_E^T$

where $h_1$ is in an earlier layer than $h_2$.

Before we start, there's a problem that we might run into when calculating all these matrices. Some of them are massive, and might not fit on our GPU. For instance, both full circuit matrices have shape $(d_\text{vocab}, d_\text{vocab})$, which in our case means $50278\times 50278 \approx 2.5\times 10^{9}$ elements. Even if your GPU can handle this, it still seems inefficient. Is there any way we can meaningfully analyse these matrices, without actually having to calculate them?

Factored Matrix class

In transformer interpretability, we often need to analyse low rank factorized matrices - a matrix $M = AB$, where M is [large, large], but A is [large, small] and B is [small, large]. This is a common structure in transformers.

For instance, we can factorise the OV circuit above as $W_{OV}^h = W_V^h W_O^h$, where $W_V^h$ has shape [768, 64] and $W_O^h$ has shape [64, 768]. For an even more extreme example, the full OV circuit can be written as $(W_E W_V^h) (W_O^h W_U)$, where these two matrices have shape [50278, 64] and [64, 50278] respectively. Similarly, we can write the full QK circuit as $(W_E W_Q^h) (W_E W_K^h)^T$.

The FactoredMatrix class is a convenient way to work with these. It implements efficient algorithms for various operations on these, such as computing the trace, eigenvalues, Frobenius norm, singular value decomposition, and products with other matrices. It can (approximately) act as a drop-in replacement for the original matrix.

This is all possible because knowing the factorisation of a matrix gives us a much easier way of computing its important properties. Intuitively, since $M=AB$ is a very large matrix that operates on very small subspaces, we shouldn't expect knowing the actual values $M_{ij}$ to be the most efficient way of storing it!

Exercise - deriving properties of a factored matrix

To give you an idea of what kinds of properties you can easily compute if you have a factored matrix, let's try and derive some ourselves.

Suppose we have $M=AB$, where $A$ has shape $(m, n)$, $B$ has shape $(n, m)$, and $m > n$. So $M$ is a size-$(m, m)$ matrix with rank at most $n$.

Question - how can you easily compute the trace of $M$?

Question - how can you easily compute the eigenvalues of $M$?

(As you'll see in later exercises, eigenvalues are very important for evaluating matrices, for instance we can assess the copying scores of an OV circuit by looking at the eigenvalues of $W_{OV}$.)

Question (hard) - how can you easily compute the SVD of $M$?

If you're curious, you can go to the FactoredMatrix documentation to see the implementation of the SVD calculation, as well as other properties and operations.

Now that we've discussed some of the motivations behind having a FactoredMatrix class, let's see it in action.

Basic Examples

We can use the basic class directly - let's make a factored matrix directly and look at the basic operations:

A = t.randn(5, 2)
B = t.randn(2, 5)
AB = A @ B
AB_factor = FactoredMatrix(A, B)
print("Norms:")
print(AB.norm())
print(AB_factor.norm())

print(f"Right dim: {AB_factor.rdim}, Left dim: {AB_factor.ldim}, Hidden dim: {AB_factor.mdim}")

We can also look at the eigenvalues and singular values of the matrix. Note that, because the matrix is rank 2 but 5 by 5, the final 3 eigenvalues and singular values are zero - the factored class omits the zeros.

print("Eigenvalues:")
print(t.linalg.eig(AB).eigenvalues)
print(AB_factor.eigenvalues)

print("\nSingular Values:")
print(t.linalg.svd(AB).S)
print(AB_factor.S)

print("\nFull SVD:")
print(AB_factor.svd())

We can multiply a factored matrix with an unfactored matrix to get another factored matrix (as in example below). We can also multiply two factored matrices together to get another factored matrix.

C = t.randn(5, 300)
ABC = AB @ C
ABC_factor = AB_factor @ C

print(f"Unfactored: shape={ABC.shape}, norm={ABC.norm()}")
print(f"Factored: shape={ABC_factor.shape}, norm={ABC_factor.norm()}")
print(f"\nRight dim: {ABC_factor.rdim}, Left dim: {ABC_factor.ldim}, Hidden dim: {ABC_factor.mdim}")

If we want to collapse this back to an unfactored matrix, we can use the AB property to get the product:

AB_unfactored = AB_factor.AB
t.testing.assert_close(AB_unfactored, AB)

Reverse-engineering circuits

Within our induction circuit, we have four individual circuits: the OV and QK circuits in our previous token head, and the OV and QK circuits in our induction head. In the following sections of the exercise, we'll reverse-engineer each of these circuits in turn.

• In the section OV copying circuit, we'll look at the layer-1 OV circuit.
• In the section QK prev-token circuit, we'll look at the layer-0 QK circuit.
• The third section (K-composition) is a bit trickier, because it involves looking at the composition of the layer-0 OV circuit and layer-1 QK circuit. We will have to do two things:
  1. Show that these two circuits are composing (i.e. that the output of the layer-0 OV circuit is the main determinant of the key vectors in the layer-1 QK circuit).
  2. Show that the joint operation of these two circuits is "make the second instance of a token attend to the token following an earlier instance.

The dropdown below contains a diagram explaining how the three sections relate to the different components of the induction circuit. You might have to open it in a new tab to see it clearly.

After this, we'll have a look at composition scores, which are a more mathematically justified way of showing that two attention heads are composing (without having to look at their behaviour on any particular class of inputs, since it is a property of the actual model weights).

[1] OV copying circuit

Let's start with an easy parts of the circuit - the copying OV circuit of 1.4 and 1.10. Let's start with head 4. The only interpretable (read: privileged basis) things here are the input tokens and output logits, so we want to study the matrix:

(and same for 1.10). This is the $(d_\text{vocab}, d_\text{vocab})$-shape matrix that combines with the attention pattern to get us from input to output.

We want to calculate this matrix, and inspect it. We should find that its diagonal values are very high, and its non-diagonal values are much lower.

Question - why should we expect this observation? (you may find it helpful to refer back to the previous section, where you described what the interpretation of different matrices was.)

Answer

If our repeating sequence is A B ... A B, then:

$$ B^T W_E W_{OV}^{1.4} W_U $$

is the vector of logits which gets moved from the first B token to the second A token, to be used as the prediction for the token following the second A token. It should result in a high prediction for B, and a low prediction for everything else. In other words, the (B, X)-th element of this matrix should be highest for X=B, which is exactly what we claimed.

If this still seems confusing, the diagram below might help:

Exercise - compute OV circuit for 1.4

This is the first of several similar exercises where you calculate a circuit by multiplying matrices. This exercise is pretty important (in particular, you should make sure you understand what this matrix represents and why we're interested in it), but the actual calculation shouldn't take very long.

You should compute it as a FactoredMatrix object.

Remember, you can access the model's weights directly e.g. using model.W_E or model.W_Q (the latter gives you all the W_Q matrices, indexed by layer and head).

head_index = 4
layer = 1

# YOUR CODE HERE - complete the `full_OV_circuit` object

tests.test_full_OV_circuit(full_OV_circuit, model, layer, head_index)

full_OV_circuit = FactoredMatrix(W_E @ W_V, W_O @ W_U)

OV_circuit = FactoredMatrix(W_V, W_O)
full_OV_circuit = W_E @ OV_circuit @ W_U

head_index = 4
layer = 1
W_O = model.W_O[layer, head_index]
W_V = model.W_V[layer, head_index]
W_E = model.W_E
W_U = model.W_U
OV_circuit = FactoredMatrix(W_V, W_O)
full_OV_circuit = W_E @ OV_circuit @ W_U

Now we want to check that this matrix is the identity. Since it's in factored matrix form, this is a bit tricky, but there are still things we can do.

First, to validate that it looks diagonal-ish, let's pick 200 random rows and columns and visualise that - it should at least look identity-ish here! We're using the indexing method of the FactoredMatrix class - you can index into it before returning the actual .AB value, to avoid having to compute the whole thing (we take advantage of the fact that A[left_indices, :] @ B[:, right_indices] is the same as (A @ B)[left_indices, right_indices]).

indices = t.randint(0, model.cfg.d_vocab, (200,))
full_OV_circuit_sample = full_OV_circuit[indices, indices].AB

imshow(
    full_OV_circuit_sample,
    labels={"x": "Logits on output token", "y": "Input token"},
    title="Full OV circuit for copying head",
    width=700,
    height=600,
)

Exercise - compute circuit accuracy

When you index a factored matrix, you get back another factored matrix. So rather than explicitly calculating A[left_indices, :] @ B[:, left_indices], we can just write AB[left_indices, left_indices].

You should observe a pretty distinct diagonal pattern here, which is a good sign. However, the matrix is pretty noisy so it probably won't be exactly the identity. Instead, we should come up with a summary statistic to capture a rough sense of "closeness to the identity".

Accuracy is a good summary statistic - what fraction of the time is the largest logit on the diagonal? Even if there's lots of noise, you'd probably still expect the largest logit to be on the diagonal a good deal of the time.

If you're on a Colab or have a powerful GPU, you should be able to compute the full matrix and perform this test. However, it's better practice to iterate through this matrix when we can, so that we avoid CUDA issues. We've given you a batch_size argument in the function below, and you should try to only explicitly calculate matrices of size batch_size * d_vocab rather than the massive matrix of d_vocab * d_vocab.

def top_1_acc(full_OV_circuit: FactoredMatrix, batch_size: int = 1000) -> float:
    """
    Return the fraction of the time that the maximum value is on the circuit diagonal.
    """
    raise NotImplementedError()


print(f"Fraction of time that the best logit is on diagonal: {top_1_acc(full_OV_circuit):.4f}")

def top_1_acc(full_OV_circuit: FactoredMatrix, batch_size: int = 1000) -> float:
    """
    Return the fraction of the time that the maximum value is on the circuit diagonal.
    """
    total = 0
for indices in t.split(t.arange(full_OV_circuit.shape[0], device=device), batch_size):
        AB_slice = full_OV_circuit[indices].AB
        total += (t.argmax(AB_slice, dim=1) == indices).float().sum().item()
return total / full_OV_circuit.shape[0]

This should return about 30.79% - pretty underwhelming. It goes up to 47.73% for top-5, but still not great. What's up with that?

Exercise - compute effective circuit

Now we return to why we have two induction heads. If both have the same attention pattern, the effective OV circuit is actually $W_E(W_V^{1.4}W_O^{1.4}+W_V^{1.10}W_O^{1.10})W_U$, and this is what matters. So let's re-run our analysis on this!

# YOUR CODE HERE - compute the effective OV circuit, and run `top_1_acc` on it

W_O_both = einops.rearrange(model.W_O[1, [4, 10]], "head d_head d_model -> (head d_head) d_model")
W_V_both = einops.rearrange(model.W_V[1, [4, 10]], "head d_model d_head -> d_model (head d_head)")
W_OV_eff = W_E @ FactoredMatrix(W_V_both, W_O_both) @ W_U
print(f"Fraction of the time that the best logit is on the diagonal: {top_1_acc(W_OV_eff):.4f}")

[2] QK prev-token circuit

The other easy circuit is the QK-circuit of L0H7 - how does it know to be a previous token circuit?

We can multiply out the full QK circuit via the positional embeddings:

to get a matrix pos_by_pos of shape [max_ctx, max_ctx] (max ctx = max context length, i.e. maximum length of a sequence we're allowing, which is set by our choice of dimensions in $W_\text{pos}$).

Note that in this case, our max context window is 2048 (we can check this via model.cfg.n_ctx). This is much smaller than the 50k-size matrices we were working with in the previous section, so we shouldn't need to use the factored matrix class here.

Exercise - interpret full QK-circuit for 0.7

The code below plots the full QK circuit for head 0.7 (including a scaling and softmax step, which is meant to mirror how the QK bilinear form will be used in actual attention layers). You should run the code, and interpret the results in the context of the induction circuit.

layer = 0
head_index = 7

# Compute full QK matrix (for positional embeddings)
W_pos = model.W_pos
W_QK = model.W_Q[layer, head_index] @ model.W_K[layer, head_index].T
pos_by_pos_scores = W_pos @ W_QK @ W_pos.T

# Mask, scale and softmax the scores
mask = t.tril(t.ones_like(pos_by_pos_scores)).bool()
pos_by_pos_pattern = t.where(mask, pos_by_pos_scores / model.cfg.d_head**0.5, -1.0e6).softmax(-1)

# Plot the results
print(f"Avg lower-diagonal value: {pos_by_pos_pattern.diag(-1).mean():.4f}")
imshow(
    utils.to_numpy(pos_by_pos_pattern[:200, :200]),
    labels={"x": "Key", "y": "Query"},
    title="Attention patterns for prev-token QK circuit, first 100 indices",
    width=700,
    height=600,
)

Click to see the expected output (and interpretation)

Avg lower-diagonal value: 0.9978

The full QK circuit $W_\text{pos} W_{QK}^{0.7} W_\text{pos}^T$ has shape [n_ctx, n_ctx]. It is a bilinear form, with the $(i, j)$-th element representing the attention score paid by the $i$-th token to the $j$-th token. This should be very large when $j = i - 1$ (and smaller for all other values of $j$), because this is a previous head token. So if we softmax over $j$, we should get a lower-diagonal stripe of 1.

Why is it justified to ignore token encodings? In this case, it turns out that the positional encodings have a much larger effect on the attention scores than the token encodings. If you want, you can verify this for yourself - after going through the next section (reverse-engineering K-composition), you'll have a better sense of how to perform attribution on the inputs to attention heads, and assess their importance.

[3] K-composition circuit

We now dig into the hard part of the circuit - demonstrating the K-Composition between the previous token head and the induction head.

Splitting activations

We can repeat the trick from the logit attribution scores. The QK-input for layer 1 is the sum of 14 terms (2+n_heads) - the token embedding, the positional embedding, and the results of each layer 0 head. So for each head $\text{H}$ in layer 1, the query tensor (ditto key) corresponding to sequence position $i$ is:

where $e$ stands for the token embedding, $pe$ for the positional embedding, and $x^\text{0.h}$ for the output of head $h$ in layer 0 (and the sum of these tensors equals the residual stream $x$). All these tensors have shape [seq, d_model]. So we can treat the expression above as a sum of matrix multiplications [seq, d_model] @ [d_model, d_head] -> [seq, d_head].

For ease of notation, I'll refer to the 14 inputs as $(y_0, y_1, ..., y_{13})$ rather than $(e, pe, x^\text{0.h}, ..., x^{h.11})$. So we have:

with each $y_i$ having shape [seq, d_model], and the sum of $y_i$s being the full residual stream $x$. Here is a diagram to illustrate:

Exercise - analyse the relative importance

We can now analyse the relative importance of these 14 terms! A very crude measure is to take the norm of each term (by component and position).

Note that this is a pretty dodgy metric - q and k are not inherently interpretable! But it can be a good and easy-to-compute proxy.

Fill in the functions below:

def decompose_qk_input(cache: ActivationCache) -> Float[Tensor, "n_heads+2 posn d_model"]:
    """
    Retrieves all the input tensors to the first attention layer, and concatenates them along the
    0th dim.

    The [i, :, :]th element is y_i (from notation above). The sum of these tensors along the 0th
    dim should be the input to the first attention layer.
    """
    raise NotImplementedError()


def decompose_q(
    decomposed_qk_input: Float[Tensor, "n_heads+2 posn d_model"],
    ind_head_index: int,
    model: HookedTransformer,
) -> Float[Tensor, "n_heads+2 posn d_head"]:
    """
    Computes the tensor of query vectors for each decomposed QK input.

    The [i, :, :]th element is y_i @ W_Q (so the sum along axis 0 is just the q-values).
    """
    raise NotImplementedError()


def decompose_k(
    decomposed_qk_input: Float[Tensor, "n_heads+2 posn d_model"],
    ind_head_index: int,
    model: HookedTransformer,
) -> Float[Tensor, "n_heads+2 posn d_head"]:
    """
    Computes the tensor of key vectors for each decomposed QK input.

    The [i, :, :]th element is y_i @ W_K(so the sum along axis 0 is just the k-values)
    """
    raise NotImplementedError()


# Recompute rep tokens/logits/cache, if we haven't already
seq_len = 50
batch_size = 1
(rep_tokens, rep_logits, rep_cache) = run_and_cache_model_repeated_tokens(
    model, seq_len, batch_size
)
rep_cache.remove_batch_dim()

ind_head_index = 4

# First we get decomposed q and k input, and check they're what we expect
decomposed_qk_input = decompose_qk_input(rep_cache)
decomposed_q = decompose_q(decomposed_qk_input, ind_head_index, model)
decomposed_k = decompose_k(decomposed_qk_input, ind_head_index, model)
t.testing.assert_close(
    decomposed_qk_input.sum(0),
    rep_cache["resid_pre", 1] + rep_cache["pos_embed"],
    rtol=0.01,
    atol=1e-05,
)
t.testing.assert_close(
    decomposed_q.sum(0), rep_cache["q", 1][:, ind_head_index], rtol=0.01, atol=0.001
)
t.testing.assert_close(
    decomposed_k.sum(0), rep_cache["k", 1][:, ind_head_index], rtol=0.01, atol=0.01
)

# Second, we plot our results
component_labels = ["Embed", "PosEmbed"] + [f"0.{h}" for h in range(model.cfg.n_heads)]
for decomposed_input, name in [(decomposed_q, "query"), (decomposed_k, "key")]:
    imshow(
        utils.to_numpy(decomposed_input.pow(2).sum([-1])),
        labels={"x": "Position", "y": "Component"},
        title=f"Norms of components of {name}",
        y=component_labels,
        width=800,
        height=400,
    )

positional_embedding_type="shortformer"

def decompose_qk_input(cache: ActivationCache) -> Float[Tensor, "n_heads+2 posn d_model"]:
    """
    Retrieves all the input tensors to the first attention layer, and concatenates them along the
    0th dim.
    The [i, :, :]th element is y_i (from notation above). The sum of these tensors along the 0th
    dim should be the input to the first attention layer.
    """
    y0 = cache["embed"].unsqueeze(0)  # shape (1, seq, d_model)
    y1 = cache["pos_embed"].unsqueeze(0)  # shape (1, seq, d_model)
    y_rest = cache["result", 0].transpose(0, 1)  # shape (12, seq, d_model)
return t.concat([y0, y1, y_rest], dim=0)
def decompose_q(
    decomposed_qk_input: Float[Tensor, "n_heads+2 posn d_model"],
    ind_head_index: int,
    model: HookedTransformer,
) -> Float[Tensor, "n_heads+2 posn d_head"]:
    """
    Computes the tensor of query vectors for each decomposed QK input.
    The [i, :, :]th element is y_i @ W_Q (so the sum along axis 0 is just the q-values).
    """
    W_Q = model.W_Q[1, ind_head_index]
return einops.einsum(decomposed_qk_input, W_Q, "n seq d_model, d_model d_head -> n seq d_head")
def decompose_k(
    decomposed_qk_input: Float[Tensor, "n_heads+2 posn d_model"],
    ind_head_index: int,
    model: HookedTransformer,
) -> Float[Tensor, "n_heads+2 posn d_head"]:
    """
    Computes the tensor of key vectors for each decomposed QK input.
    The [i, :, :]th element is y_i @ W_K(so the sum along axis 0 is just the k-values)
    """
    W_K = model.W_K[1, ind_head_index]
return einops.einsum(decomposed_qk_input, W_K, "n seq d_model, d_model d_head -> n seq d_head")
# Recompute rep tokens/logits/cache, if we haven't already
seq_len = 50
batch_size = 1
(rep_tokens, rep_logits, rep_cache) = run_and_cache_model_repeated_tokens(
    model, seq_len, batch_size
)
rep_cache.remove_batch_dim()
ind_head_index = 4
# First we get decomposed q and k input, and check they're what we expect
decomposed_qk_input = decompose_qk_input(rep_cache)
decomposed_q = decompose_q(decomposed_qk_input, ind_head_index, model)
decomposed_k = decompose_k(decomposed_qk_input, ind_head_index, model)
t.testing.assert_close(
    decomposed_qk_input.sum(0),
    rep_cache["resid_pre", 1] + rep_cache["pos_embed"],
    rtol=0.01,
    atol=1e-05,
)
t.testing.assert_close(
    decomposed_q.sum(0), rep_cache["q", 1][:, ind_head_index], rtol=0.01, atol=0.001
)
t.testing.assert_close(
    decomposed_k.sum(0), rep_cache["k", 1][:, ind_head_index], rtol=0.01, atol=0.01
)
# Second, we plot our results
component_labels = ["Embed", "PosEmbed"] + [f"0.{h}" for h in range(model.cfg.n_heads)]
for decomposed_input, name in [(decomposed_q, "query"), (decomposed_k, "key")]:
    imshow(
        utils.to_numpy(decomposed_input.pow(2).sum([-1])),
        labels={"x": "Position", "y": "Component"},
        title=f"Norms of components of {name}",
        y=component_labels,
        width=800,
        height=400,
    )

This tells us which heads are probably important, but we can do better than that. Rather than looking at the query and key components separately, we can see how they combine together - i.e. take the decomposed attention scores.

This is a bilinear function of q and k, and so we will end up with a decomposed_scores tensor with shape [query_component, key_component, query_pos, key_pos], where summing along BOTH of the first axes will give us the original attention scores (pre-mask).

Exercise - decompose attention scores

Implement the function giving the decomposed scores (remember to scale by sqrt(d_head)!) For now, don't mask it.

Help - I'm confused about what we're doing / why we're doing it.

Remember that each of our components writes to the residual stream separately. So after layer 0, we have:

We're particularly interested in the attention scores computed in head 1.4, and how they depend on the inputs into that head. We've already decomposed the residual stream value $x$ into its terms $e$, $pe$, and $x^ 0$ through $x^{11}$ (which we've labelled $y_0, ..., y_{13}$ for simplicity), and we've done the same for key and query terms. We can picture these terms being passed into head 1.4 as:

So when we expand attn_scores out in full, they are a sum of $14^2 = 196$ terms - one for each combination of (query_component, key_component).

---

#### Why is this decomposition useful?

We have a theory about a particular circuit in our model. We think that head 1.4 is an induction head, and the most important components that feed into this head are the prev token head 0.7 (as key) and the token embedding (as query). This is already supported by the evidence of our magnitude plots above (because we saw that 0.7 as key and token embeddings as query were large), but we still don't know how this particular key and query work together; we've only looked at them separately.

By decomposing attn_scores like this, we can check whether the contribution from combination (query=tok_emb, key=0.7) is indeed producing the characteristic induction head pattern which we've observed (and the other 195 terms don't really matter).

def decompose_attn_scores(
    decomposed_q: Float[Tensor, "q_comp q_pos d_head"],
    decomposed_k: Float[Tensor, "k_comp k_pos d_head"],
    model: HookedTransformer,
) -> Float[Tensor, "q_comp k_comp q_pos k_pos"]:
    """
    Output is decomposed_scores with shape [query_component, key_component, query_pos, key_pos]

    The [i, j, 0, 0]th element is y_i @ W_QK @ y_j^T (so the sum along both first axes are the
    attention scores)
    """
    raise NotImplementedError()


tests.test_decompose_attn_scores(decompose_attn_scores, decomposed_q, decomposed_k, model)

def decompose_attn_scores(
    decomposed_q: Float[Tensor, "q_comp q_pos d_head"],
    decomposed_k: Float[Tensor, "k_comp k_pos d_head"],
    model: HookedTransformer,
) -> Float[Tensor, "q_comp k_comp q_pos k_pos"]:
    """
    Output is decomposed_scores with shape [query_component, key_component, query_pos, key_pos]
    The [i, j, 0, 0]th element is y_i @ W_QK @ y_j^T (so the sum along both first axes are the
    attention scores)
    """
    return einops.einsum(
        decomposed_q,
        decomposed_k,
        "q_comp q_pos d_head, k_comp k_pos d_head -> q_comp k_comp q_pos k_pos",
    ) / (model.cfg.d_head**0.5)

Once these tests have passed, you can plot the results:

# First plot: attention score contribution from (query_component, key_component) = (Embed, L0H7), you can replace this
# with any other pair and see that the values are generally much smaller, i.e. this pair dominates the attention score
# calculation
decomposed_scores = decompose_attn_scores(decomposed_q, decomposed_k, model)

q_label = "Embed"
k_label = "0.7"
decomposed_scores_from_pair = decomposed_scores[
    component_labels.index(q_label), component_labels.index(k_label)
]

imshow(
    utils.to_numpy(t.tril(decomposed_scores_from_pair)),
    title=f"Attention score contributions from query = {q_label}, key = {k_label}<br>(by query & key sequence positions)",
    width=700,
)


# Second plot: std dev over query and key positions, shown by component. This shows us that the other pairs of
# (query_component, key_component) are much less important, without us having to look at each one individually like we
# did in the first plot!
decomposed_stds = einops.reduce(
    decomposed_scores, "query_decomp key_decomp query_pos key_pos -> query_decomp key_decomp", t.std
)
imshow(
    utils.to_numpy(decomposed_stds),
    labels={"x": "Key Component", "y": "Query Component"},
    title="Std dev of attn score contributions across sequence positions<br>(by query & key comp)",
    x=component_labels,
    y=component_labels,
    width=700,
)

Note that plots like the ones above are often the most concise way of presenting a summary of the important information, and understanding what to plot is a valuable skill in any model internals-based work. However, if you want to see the "full plot" which the two plots above are both simplifications of in some sense, you can run the code below, which gives you the matrix of every single pair of components' contribution to the attention scores. So the first plot above is just a slice of the full plot below, and the second plot above is just the plot below after reducing over each slice with the standard deviation operation.

(Note - the plot you'll generate below is pretty big, so you'll want to clear it after you're done with it. If your machine is still working slowly when rendering it, you can use fig.show(config={"staticPlot": True}) to display a non-interactive version of it.)

decomposed_scores_centered = t.tril(
    decomposed_scores - decomposed_scores.mean(dim=-1, keepdim=True)
)

decomposed_scores_reshaped = einops.rearrange(
    decomposed_scores_centered,
    "q_comp k_comp q_token k_token -> (q_comp q_token) (k_comp k_token)",
)

fig = imshow(
    decomposed_scores_reshaped,
    title="Attention score contributions from all pairs of (key, query) components",
    width=1200,
    height=1200,
    return_fig=True,
)
full_seq_len = seq_len * 2 + 1
for i in range(0, full_seq_len * len(component_labels), full_seq_len):
    fig.add_hline(y=i, line_color="black", line_width=1)
    fig.add_vline(x=i, line_color="black", line_width=1)

fig.show(config={"staticPlot": True})

Interpreting the full circuit

Now we know that head 1.4 is composing with head 0.7 via K composition, we can multiply through to create a full circuit:

and verify that it's the identity. (Note, when we say identity here, we're again thinking about it as a distribution over logits, so this should be taken to mean "high diagonal values", and we'll be using our previous metric of top_1_acc.)

Question - why should this be the identity?

Answer

This matrix is a bilinear form. Its diagonal elements $(A, A)$ are:

$$ A^T \, W_E\, W_{QK}^{1.4}\, W_{OV}^{0.7}\, W_E^T \, A = \underbrace{(A^T W_E W_Q^{1.4})}_{\text{query}} \underbrace{(A^T W_E W_{OV}^{0.7} W_K^{1.4})^T}_{\text{key}} $$

Intuitively, the query is saying "I'm looking for a token which followed $A$", and the key is saying **"I am a token which folllowed $A$"** (recall that $A^T W_E W_{OV}^{0.7}$ is the vector which gets moved one position forward by our prev token head 0.7).

Now, consider the off-diagonal elements $(A, X)$ (for $X \neq A$). We expect these to be small, because the key doesn't match the query:

$$ A^T \, W_E\, W_{QK}^{1.4}\, W_{OV}^{0.7}\, W_E^T \, X = \underbrace{(\text{I'm looking for a token which followed A})}_\text{query} \boldsymbol{\cdot} \underbrace{(\text{I am a token which followed X})}_{\text{key}} $$

Hence, we expect this to be the identity.

An illustration:

Exercise - compute the K-comp circuit

Calculate the matrix above, as a FactoredMatrix object.

FactoredMatrix(M1.AB @ M2.A, M2.B)

FactoredMatrix(M1.A, M1.B @ M2.AB)

def find_K_comp_full_circuit(
    model: HookedTransformer, prev_token_head_index: int, ind_head_index: int
) -> FactoredMatrix:
    """
    Returns a (vocab, vocab)-size FactoredMatrix, with the first dimension being the query side
    (direct from token embeddings) and the second dimension being the key side (going via the
    previous token head).
    """
    raise NotImplementedError()


prev_token_head_index = 7
ind_head_index = 4
K_comp_circuit = find_K_comp_full_circuit(model, prev_token_head_index, ind_head_index)

tests.test_find_K_comp_full_circuit(find_K_comp_full_circuit, model)

print(f"Token frac where max-activating key = same token: {top_1_acc(K_comp_circuit.T):.4f}")

def find_K_comp_full_circuit(
    model: HookedTransformer, prev_token_head_index: int, ind_head_index: int
) -> FactoredMatrix:
    """
    Returns a (vocab, vocab)-size FactoredMatrix, with the first dimension being the query side
    (direct from token embeddings) and the second dimension being the key side (going via the
    previous token head).
    """
    W_E = model.W_E
    W_Q = model.W_Q[1, ind_head_index]
    W_K = model.W_K[1, ind_head_index]
    W_O = model.W_O[0, prev_token_head_index]
    W_V = model.W_V[0, prev_token_head_index]
Q = W_E @ W_Q
    K = W_E @ W_V @ W_O @ W_K
    return FactoredMatrix(Q, K.T)

You can also try this out for our other induction head ind_head_index=10, which should also return a relatively high result. Is it higher than for head 1.4 ?

Further Exploration of Induction Circuits

I now consider us to have fully reverse engineered an induction circuit - by both interpreting the features and by reverse engineering the circuit from the weights. But there's a bunch more ideas that we can apply for finding circuits in networks that are fun to practice on induction heads, so here's some bonus content - feel free to skip to the later bonus ideas though.

Composition scores

A particularly cool idea in the paper is the idea of virtual weights, or compositional scores. (Though I came up with it, so I'm deeply biased!). This is used to identify induction heads.

The key idea of compositional scores is that the residual stream is a large space, and each head is reading and writing from small subspaces. By default, any two heads will have little overlap between their subspaces (in the same way that any two random vectors have almost zero dot product in a large vector space). But if two heads are deliberately composing, then they will likely want to ensure they write and read from similar subspaces, so that minimal information is lost. As a result, we can just directly look at "how much overlap there is" between the output space of the earlier head and the K, Q, or V input space of the later head.

We represent the output space with $W_{OV}=W_V W_O$. Call matrices like this $W_A$.

We represent the input space with $W_{QK}=W_Q W_K^T$ (for Q-composition), $W_{QK}^T=W_K W_Q^T$ (for K-Composition) or $W_{OV}=W_V W_O$ (for V-Composition, of the later head). Call matrices like these $W_B$ (we've used this notation so that $W_B$ refers to a later head, and $W_A$ to an earlier head).

How do we formalise overlap? This is basically an open question, but a surprisingly good metric is $\frac{\|W_AW_B\|_F}{\|W_B\|_F\|W_A\|_F}$ where $\|W\|_F=\sqrt{\sum_{i,j}W_{i,j}^2}$ is the Frobenius norm, the square root of the sum of squared elements. (If you're dying of curiosity as to what makes this a good metric, you can jump to the section immediately after the exercises below.)

Exercise - calculate composition scores

Let's calculate this metric for all pairs of heads in layer 0 and layer 1 for each of K, Q and V composition and plot it.

We'll start by implementing this using plain old tensors (later on we'll see how this can be sped up using the FactoredMatrix class). We also won't worry about batching our calculations yet; we'll just do one matrix at a time.

We've given you tensors q_comp_scores etc. to hold the composition scores for each of Q, K and V composition (i.e. the [i, j]th element of q_comp_scores is the Q-composition score between the output from the ith head in layer 0 and the input to the jth head in layer 1). You should complete the function get_comp_score, and then fill in each of these tensors.

def get_comp_score(W_A: Float[Tensor, "in_A out_A"], W_B: Float[Tensor, "out_A out_B"]) -> float:
    """
    Return the composition score between W_A and W_B.
    """
    raise NotImplementedError()


tests.test_get_comp_score(get_comp_score)

def get_comp_score(W_A: Float[Tensor, "in_A out_A"], W_B: Float[Tensor, "out_A out_B"]) -> float:
    """
    Return the composition score between W_A and W_B.
    """
    W_A_norm = W_A.pow(2).sum().sqrt()
    W_B_norm = W_B.pow(2).sum().sqrt()
    W_AB_norm = (W_A @ W_B).pow(2).sum().sqrt()
return (W_AB_norm / (W_A_norm * W_B_norm)).item()

Once you've passed the tests, you can fill in all the composition scores. Here you should just use a for loop, iterating over all possible pairs of W_A in layer 0 and W_B in layer 1, for each type of composition. Later on, we'll look at ways to batch this computation.

# Get all QK and OV matrices
W_QK = model.W_Q @ model.W_K.transpose(-1, -2)
W_OV = model.W_V @ model.W_O

# Define tensors to hold the composition scores
composition_scores = {
    "Q": t.zeros(model.cfg.n_heads, model.cfg.n_heads).to(device),
    "K": t.zeros(model.cfg.n_heads, model.cfg.n_heads).to(device),
    "V": t.zeros(model.cfg.n_heads, model.cfg.n_heads).to(device),
}

# YOUR CODE HERE - fill in values of the `composition_scores` dict, using `get_comp_score`

# Plot the composition scores
for comp_type in ["Q", "K", "V"]:
    plot_comp_scores(model, composition_scores[comp_type], f"{comp_type} Composition Scores")

for i in tqdm(range(model.cfg.n_heads)):
    for j in range(model.cfg.n_heads):
        composition_scores["Q"][i, j] = get_comp_score(W_OV[0, i], W_QK[1, j])
        composition_scores["K"][i, j] = get_comp_score(W_OV[0, i], W_QK[1, j].T)
        composition_scores["V"][i, j] = get_comp_score(W_OV[0, i], W_OV[1, j])

Exercise - Setting a Baseline

To interpret the above graphs we need a baseline! A good one is what the scores look like at initialisation. Make a function that randomly generates a composition score 200 times and tries this. Remember to generate 4 [d_head, d_model] matrices, not 2 [d_model, d_model] matrices! This model was initialised with Kaiming Uniform Initialisation:

W = t.empty(shape)
nn.init.kaiming_uniform_(W, a=np.sqrt(5))

(Ideally we'd do a more efficient generation involving batching, and more samples, but we won't worry about that yet.)

def generate_single_random_comp_score() -> float:
    """
    Write a function which generates a single composition score for random matrices
    """
    raise NotImplementedError()


n_samples = 300
comp_scores_baseline = np.zeros(n_samples)
for i in tqdm(range(n_samples)):
    comp_scores_baseline[i] = generate_single_random_comp_score()

print("\nMean:", comp_scores_baseline.mean())
print("Std:", comp_scores_baseline.std())

hist(
    comp_scores_baseline,
    nbins=50,
    width=800,
    labels={"x": "Composition score"},
    title="Random composition scores",
)

def generate_single_random_comp_score() -> float:
    """
    Write a function which generates a single composition score for random matrices
    """
    W_A_left = t.empty(model.cfg.d_model, model.cfg.d_head)
    W_B_left = t.empty(model.cfg.d_model, model.cfg.d_head)
    W_A_right = t.empty(model.cfg.d_model, model.cfg.d_head)
    W_B_right = t.empty(model.cfg.d_model, model.cfg.d_head)
for W in [W_A_left, W_B_left, W_A_right, W_B_right]:
        nn.init.kaiming_uniform_(W, a=np.sqrt(5))
W_A = W_A_left @ W_A_right.T
    W_B = W_B_left @ W_B_right.T
return get_comp_score(W_A, W_B)

We can re-plot our above graphs with this baseline set to white. Look for interesting things in this graph!

baseline = comp_scores_baseline.mean()
for comp_type, comp_scores in composition_scores.items():
    plot_comp_scores(model, comp_scores, f"{comp_type} Composition Scores", baseline=baseline)

Some interesting things to observe:

The most obvious thing that jumps out (when considered in the context of all the analysis we've done so far) is the K-composition scores. 0.7 (the prev token head) is strongly composing with 1.4 and 1.10 (the two attention heads). This is what we expect, and is a good indication that our composition scores are working as intended.

Another interesting thing to note is that the V-composition scores for heads 1.4 and 1.10 with all other heads in layer 0 are very low. In the context of the induction circuit, this is a good thing - the OV circuits of our induction heads should be operating on the embeddings, rather than the outputs of the layer-0 heads. (If our repeating sequence is A B ... A B, then it's the QK circuit's job to make sure the second A attends to the first B, and it's the OV circuit's job to project the residual vector at that position onto the embedding space in order to extract the B-information, while hopefully ignoring anything else that has been written to that position by the heads in layer 0). So once again, this is a good sign for our composition scores.

Theory + Efficient Implementation

So, what's up with that metric? The key is a cute linear algebra result that the squared Frobenius norm is equal to the sum of the squared singular values.

So if $W_A=U_AS_AV_A^T$, $W_B=U_BS_BV_B^T$, then $\|W_A\|_F=\|S_A\|_F$, $\|W_B\|_F=\|S_B\|_F$ and $\|W_AW_B\|_F=\|S_AV_A^TU_BS_B\|_F$. In some sense, $V_A^TU_B$ represents how aligned the subspaces written to and read from are, and the $S_A$ and $S_B$ terms weights by the importance of those subspaces.

Exercise - batching, and using the FactoredMatrix class

We can also use this insight to write a more efficient way to calculate composition scores - this is extremely useful if you want to do this analysis at scale! The key is that we know that our matrices have a low rank factorisation, and it's much cheaper to calculate the SVD of a narrow matrix than one that's large in both dimensions. See the algorithm described at the end of the paper (search for SVD).

So we can work with the FactoredMatrix class. This also provides the method .norm() which returns the Frobenium norm. This is also a good opportunity to bring back batching - this will sometimes be useful in our analysis. In the function below, W_As and W_Bs are both >2D factored matrices (e.g. they might represent the OV circuits for all heads in a particular layer, or across multiple layers), and the function's output should be a tensor of composition scores for each pair of matrices (W_A, W_B) in the >2D tensors (W_As, W_Bs).

def get_batched_comp_scores(W_As: FactoredMatrix, W_Bs: FactoredMatrix) -> Tensor:
    """
    Computes the compositional scores from indexed factored matrices W_As and W_Bs.

    Each of W_As and W_Bs is a FactoredMatrix object which is indexed by all but its last 2
    dimensions, i.e.:
        W_As.shape == (*A_idx, A_in, A_out)
        W_Bs.shape == (*B_idx, B_in, B_out)
        A_out == B_in

    Return: tensor of shape (*A_idx, *B_idx) where the [*a_idx, *b_idx]th element is the
    compositional score from W_As[*a_idx] to W_Bs[*b_idx].
    """
    raise NotImplementedError()


W_QK = FactoredMatrix(model.W_Q, model.W_K.transpose(-1, -2))
W_OV = FactoredMatrix(model.W_V, model.W_O)

composition_scores_batched = dict()
composition_scores_batched["Q"] = get_batched_comp_scores(W_OV[0], W_QK[1])
composition_scores_batched["K"] = get_batched_comp_scores(
    W_OV[0], W_QK[1].T
)  # Factored matrix: .T is interpreted as transpose of the last two axes
composition_scores_batched["V"] = get_batched_comp_scores(W_OV[0], W_OV[1])

t.testing.assert_close(composition_scores_batched["Q"], composition_scores["Q"])
t.testing.assert_close(composition_scores_batched["K"], composition_scores["K"])
t.testing.assert_close(composition_scores_batched["V"], composition_scores["V"])
print("Tests passed - your `get_batched_comp_scores` function is working!")

W_As.shape == (A1A2...Am, 1, A_in, A_out)
W_Bs.shape == (1, B1B2...Bn, B_in, B_out)

def get_batched_comp_scores(W_As: FactoredMatrix, W_Bs: FactoredMatrix) -> Tensor:
    """
    Computes the compositional scores from indexed factored matrices W_As and W_Bs.
    Each of W_As and W_Bs is a FactoredMatrix object which is indexed by all but its last 2
    dimensions, i.e.:
        W_As.shape == (A_idx, A_in, A_out)
        W_Bs.shape == (B_idx, B_in, B_out)
        A_out == B_in
    Return: tensor of shape (A_idx, B_idx) where the [a_idx, b_idx]th element is the
    compositional score from W_As[a_idx] to W_Bs[b_idx].
    """
    # Flatten W_As into (single_A_idx, 1, A_in, A_out)
    W_As = FactoredMatrix(
        W_As.A.reshape(-1, 1, W_As.A.shape[-2:]),
        W_As.B.reshape(-1, 1, W_As.B.shape[-2:]),
    )
    # Flatten W_Bs into (1, single_B_idx, B_in(=A_out), B_out)
    W_Bs = FactoredMatrix(
        W_Bs.A.reshape(1, -1, W_Bs.A.shape[-2:]),
        W_Bs.B.reshape(1, -1, W_Bs.B.shape[-2:]),
    )
# Compute the product, with shape (single_A_idx, single_B_idx, A_in, B_out)
    W_ABs = W_As @ W_Bs
# Compute the norms, and return the metric
    return W_ABs.norm() / (W_As.norm() * W_Bs.norm())

Targeted Ablations

We can refine the ablation technique to detect composition by looking at the effect of the ablation on the attention pattern of an induction head, rather than the loss. Let's implement this!

Gotcha - by default, run_with_hooks removes any existing hooks when it runs. If you want to use caching, set the reset_hooks_start flag to False.

seq_len = 50


def ablation_induction_score(prev_head_index: int | None, ind_head_index: int) -> float:
    """
    Takes as input the index of the L0 head and the index of the L1 head, and then runs with the
    previous token head ablated and returns the induction score for the ind_head_index now.
    """

    def ablation_hook(v, hook):
        if prev_head_index is not None:
            v[:, :, prev_head_index] = 0.0
        return v

    def induction_pattern_hook(attn, hook):
        hook.ctx[prev_head_index] = attn[0, ind_head_index].diag(-(seq_len - 1)).mean()

    model.run_with_hooks(
        rep_tokens,
        fwd_hooks=[
            (utils.get_act_name("v", 0), ablation_hook),
            (utils.get_act_name("pattern", 1), induction_pattern_hook),
        ],
    )
    return model.blocks[1].attn.hook_pattern.ctx[prev_head_index].item()


baseline_induction_score = ablation_induction_score(None, 4)
print(f"Induction score for no ablations: {baseline_induction_score:.5f}\n")
for i in range(model.cfg.n_heads):
    new_induction_score = ablation_induction_score(i, 4)
    induction_score_change = new_induction_score - baseline_induction_score
    print(f"Ablation score change for head {i:02}: {induction_score_change:+.5f}")

Bonus

Looking for Circuits in Real LLMs

A particularly cool application of these techniques is looking for real examples of circuits in large language models. Fortunately, there's a bunch of open source ones you can play around with in the TransformerLens library! Many of the techniques we've been using for our 2L transformer carry over to ones with more layers.

This library should make it moderately easy to play around with these models - I recommend going wild and looking for interesting circuits!

Some fun things you might want to try:

• Look for induction heads - try repeating all of the steps from above. Do they follow the same algorithm?
• Look for neurons that erase info
  • i.e. having a high negative cosine similarity between the input and output weights
• Try to interpret a position embedding.

• Look for heads with interpretable attention patterns: e.g. heads that attend to the same word (or subsequent word) when given text in different languages, or the most recent proper noun, or the most recent full-stop, or the subject of the sentence, etc.
  • Pick a head, ablate it, and run the model on a load of text with and without the head. Look for tokens with the largest difference in loss, and try to interpret what the head is doing.
• Try replicating some of Kevin's work on indirect object identification.
• Inspired by the ROME paper, use the causal tracing technique of patching in the residual stream - can you analyse how the network answers different facts?

Note: I apply several simplifications to the resulting transformer - these leave the model mathematically equivalent and doesn't change the output log probs, but does somewhat change the structure of the model and one change translates the output logits by a constant.

Training Your Own Toy Models

A fun exercise is training models on the minimal task that'll produce induction heads - predicting the next token in a sequence of random tokens with repeated subsequences. You can get a small 2L Attention-Only model to do this.

Interpreting Induction Heads During Training

A particularly striking result about induction heads is that they consistently form very abruptly in training as a phase change, and are such an important capability that there is a visible non-convex bump in the loss curve (in this model, approx 2B to 4B tokens). I have a bunch of checkpoints for this model, you can try re-running the induction head detection techniques on intermediate checkpoints and see what happens. (Bonus points if you have good ideas for how to efficiently send a bunch of 300MB checkpoints from Wandb lol)

Further discussion / investigation

Anthropic has written a post on In-context Learning and Induction Heads which goes into much deeper discussion on induction heads. The post is structured around six different points of evidence for the hypothesis that induction heads are the main source of in-context learning in transformer models, even large ones. Briefly, these are:

1. Transformers undergo a "phase change" where they suddenly become much better at in-context learning, and this is around the same time induction heads appear.
2. When we change the transformer's architecture to make it easier for induction heads to form, we get a corresponding improvement in in-context learning.
3. When we ablate induction heads at runtime, in-context learning gets worse.
4. We have specific examples of induction heads performing more complex in-context learning algorithms (you'll have the opportunity to investigate one of these later - indirect object identification).
5. We have a mechanistic explanation of induction heads, which suggests natural extensions to more general forms of in-context learning.
6. In-context learning-related behaviour is generally smoothly continuous between small and large models, suggesting that the underlying mechanism is also the same.

Here are a few questions for you:

• How compelling do you find this evidence? Discuss with your partner.
  • Which points do you find most compelling?
  • Which do you find least compelling?
  • Are there any subset of these which would be enough to convince you of the hypothesis, in the absence of others?
• In point 3, the paper observes that in-context learning performance degrades when you ablate induction heads. While we measured this by testing the model's ability to copy a duplicated random sequence, the paper used in-context learning score (the loss of the 500th token in the context, minus the loss on the 50th token).
  • Can you see why this is a reasonable metric?
  • Can you replicate these results (maybe on a larger model than the 2-layer one we've been using)?
• In point 4 (more complex forms of in-context learning), the paper suggests the natural extension of "fuzzy induction heads", which match patterns like [A*][B*]...[A][B] rather than [A][B]...[A][B] (where the * indicates some form of linguistic similarity, not necessarily being the same token).
  • Can you think of any forms this might take, i.e. any kinds of similarity which induction heads might pick up on? Can you generate examples?