đź““ Notebook: gpt2things / gpt2_attentions.ipynb

I wanted to know what the attention blocks are actually doing before anything gets normalised away, so I picked the smallest pair of sentences I could think of that mean opposite things:

python
# positive_prompt mocks a user's desired text completion,
# negative_prompt mocks the default text completion
positive_prompt = "I love you"
negative_prompt = "I hate you"

Same length, same structure, and only one word apart. If the attention maps look different, that difference has to come from love versus hate.

GPT-2 isn't quite like the others ​

ℹ️ Info: apart from the number of layers and the training parameters, GPT-2's attention block uses a Conv1D layer instead of a linear layer to compute the attention matrices. The attention matrix it returns is a stochastic matrix, so each row sums to 1.

The Attendance class ​

A small helper that runs the model with output_attentions=True and stacks every block's attention into one tensor:

python
class Attendance:
    """Displays attention blocks for gpt2 specifically"""

    def run(self, inputs: str):
        tokens = self.tokenizer(inputs, add_special_tokens=True,
                                return_attention_mask=True, return_tensors="pt")
        with torch.no_grad():
            output = self.model(**tokens, output_attentions=True)
        # [num_attention_layers, batch_index, num_attention_heads, num_tokens, num_tokens]
        attn = torch.stack(output.attentions)
        return attn, tokens["input_ids"].squeeze(0)

That five-dimensional tensor is the whole point of the notebook. processAttention lets me slice it two ways:

  • view_by="block": pick one block (first or last of the 12) and plot every head.
  • view_by="layer": hold the head position fixed and walk through the blocks. Self note: I wasn't sure about the naming here; the heading in the notebook literally says "Viewing by third dimension or Layer (?)".

What came out ​

I plotted the first and last block for both prompts, then averaged the last four components with last_attention(start=4). For the three tokens the result is:

Ilove/hateyou
I1.000000
love/hate0.81040.18960
you0.74530.14340.1113

The upper triangle is all zero because GPT-2 is causal: a token can't attend to anything after it. Each row sums to 1, as expected. Most of the attention goes back to the first token, even from you.

Why I was doing this ​

This notebook was the groundwork for steering GPT-2's decoding two days later, and for cloning myself into a language model for Tom. Before I could push a model towards the reply I wanted, I needed to see where in the model the two replies come apart.