Attention sink
More evidence
Intro
I read the induction heads paper
While playing around with induction heads in GPT2, I thought to myself that "What if the input to induction heads isn't present, what do the induction heads pay attention to?" I thought this might be a good question to investigate and after a quick literature search, I stumbled on the attention sink paper and a bunch of other works that made fantastic attempts at answering the question.
Guo, et al., in "Active-Dormant Attention Heads"
While the question was sort of answered already, I thought it would still be a good exercise present the thought process I went through while attempting to answer the question.
In this remainder of this post, I briefly motivate what an attention head is doing, explain induction heads and how to look for them (with visualizations) and show what happens when the induction heads input isn't present.
Feel free to skip parts you're familiar with.
What is an attention head doing ?
Intuitive overview
In summary, attention heads move information between tokens!
The residual stream is the main object in the transformer. A way I
think of it is that it represents what the model currently thinks
about all the tokens in it's context, up to a particular layer. To
enrich and further refine the representation of the tokens in the
context, attention heads move information from earlier tokens in the
context to later tokens in the context
More concretely
The input to the attention layer is the residual stream
[batch_size, seq_len, d_model]. This input is linearly projected using three weight matrices:
W_Q,
W_K, and
W_V, each of shape
[d_model, d_model], to produce
the query (Q), key (K), and value (V) matrices.
In multi-head attention, Q,
K, and
V are split into
num_heads parts. Each head
processes a subspace of the input, with
Q and
K shaped as
[batch_size, num_heads, seq_len, d_k]
and V as
[batch_size, num_heads, seq_len, d_v], where
d_k = d_v = d_model / num_heads.
For each head, attention scores are computed as the
dot product of query and key vectors, scaled by
1/√d_k. The scores are passed
through a softmax to obtain the attention pattern,
which represents the importance of each token relative to others. This
pattern is then multiplied by the value vectors to produce the head's
output. The outputs of all heads are concatenated and projected using
a weight matrix W_O of shape
[d_model, d_model] to yield
the final attention output, shaped
[batch_size, seq_len, d_model].
From Vaswani et al. W_i^Q,
W_i^K, and
W_i^V are head-specific
projection matrices.
Further questions
The next logical question is, how does each attention head across all the layers know what sort of information to pay attention to? During pre-training, the goal is optimizing the next-token objective w.r.t the parameters of the model, over the language domain. It stands to reason that over the course of multiple steps of gradient descent, each attention head learns to pay attention to some pattern (semantic or syntactic) in the language data and this pattern, when learned, contributes to lower loss.
And indeed, numerous papers have explored this assumption.
In both decoder-only and encoder-decocder transformers, attention
heads have been discovered that specialize in attending to different
parts of speech, as well as other lingustic propertites such direct
objects of verbs, noun determiners, e.t.c.
Interesting mechanisms that further enable LLMs to act
autoregressively have been discovered, such as Copy Supression heads
A logical conclusion of the above paragraphs is that what attention heads pay attention to is input-specific. This begs the question : What does an attention head pay attention to, when it's input isn't present?
Induction heads
I'll present a super simplified explanation of Induction heads here,
but to better understand Induction heads mechanistically, Callum
McDougall wrote
a quite interesting explainer blog
which I invite readers to check out. The paper
Assume arbitrary tokens A, B.
Then assume a sequence of tokens with
A followed by
B and then some other
arbitrary tokens. The next time the model sees
A, i.e
[A B ... A], B turns out to be
one of the highly likely next tokens.
Anthropic researchers found these phenomenon in as little as 2L
transformer. One of the conclusions is that the model has learnt to
increase the logits on B if
the last token in the sequence is
A and indeed, it's theorized
that Induction heads is one of the mechanisms behind In-context
learning.
For this to be true, there has to be a previous-token head. This
ensures that the first occurrence of
[B] pays attention to the
first occurrence of [A] and
the \(W_V\) matrix copies A to
the subspace of B. Then when
A occurs in the context again,
for some head \(\hat{h}\), the second occurrence of
A pays attention to the first
occurrence of B, sees that
A is in the residual stream of
B and then copies
B to the residual stream of
the second occurrence of A and
increases it's logits. This new head \(\hat{h}\) is an induction head.
Identifying Induction heads in GPT2
Inputs
We sample N = 25 random tokens
from the vocabulary of a transformer language model and duplicate it
along it's axis. This becomes the input to the transformer.
Input to the transformer is a matrix of shape
[1, 2 * N + 1, d_model] i.e
batch is 1, sequence length is 2 * N + 1
Pass this sequence of randomly repeated tokens into GPT2 and cache the
activations. This can be done easily by loading the model with
transformer lens
and running
_, cache = model.run_with_cache(input_tokens)
Metric
Assume we have some head h at
some layer l, the attention
pattern is defined as, $$ \text{A}^{l, h} =
\text{softmax}\left(\frac{Q_{l, h}K^T_{l, h}}{\sqrt{d_k}}\right) $$
We define induction score for head
h in layer
l
as a measurement of how much attention a token in the second repeat
(at position i + N) pays to
its corresponding token in the first repeat (at position
i). It's represented as:
Identifying induction heads
Retrieve the attention pattern from the cache and for each head in each layer, calculate the induction score as defined above.
def induction_head_detector( cache, cfg, ) -> list:
induction_heads = []
for layer_idx in range(cfg.n_layers):
for head_idx in range(cfg.n_heads):
# fetch the attention pattern at some layer and some head
attn_pattern = cache["pattern",layer_idx][head_idx]
rand_tok_seq_len = (attn_pattern.shape[1] -1) // 2
# compute the induction score for the attention pattern
score = attn_pattern.diagonal(-rand_tok_seq_len + 1).mean()
# filter with threshold of 0.4
if score.item() >= 0.4:
induction_heads.append((layer_idx, head_idx))
return induction_heads
Results
Below is a visual map of the induction heads present in GPT2
Below is an interactive visualization of the attention patterns for the induction heads identified above.
What happens if the induction input isn't present?
Inputs
Load a tiny subset of the
10K pile dataset.For the purpose of this experiment, I used
batch = 1 and
sequence_length = 128.
Forward pass is also ran on this input and the activations are cached as in the case above as well.
For the induction heads that were identified in the section above, we simply visualize the attention pattern for these heads.
Results
As can be observed, these heads all pay an overwhelming amount of attention to the first token.Closing thoughts
Guo et. al.,
They showed further evidence of this phenomenon by confirming that the
value vectors of these tokens were much smaller than that of other
tokens
This however isn't the only explanation for the first-token/special
token phenomenon observed in attention heads. Federico