đź““ 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:
# 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
Conv1Dlayer 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:
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:
| I | love/hate | you | |
|---|---|---|---|
| I | 1.0000 | 0 | 0 |
| love/hate | 0.8104 | 0.1896 | 0 |
| you | 0.7453 | 0.1434 | 0.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.