Charts/Matrix charts
Attention Map
<AttentionMap /> · 10 examples
A grid of weights between queries and keys, one cell per pair, darker where the weight is higher.
Overview
When a language model reads "the animal didn't cross the street because it was too tired", something inside it has to decide what "it" refers to. That decision is made by attention: for every word, the model computes how much to look at every other word. Those numbers form a table, one row per word doing the looking and one column per word being looked at. An attention map is that table, coloured.
The picture is simple and, since the transformer architecture arrived in 2017, it has become the first thing people draw when they want to look inside a model. Row by row, you can see which tokens a token leaned on. A bright diagonal means "I mostly look at myself". A bright column means "everybody looks at this one word", which is frequently the first token or a full stop, to the endless amusement of people who hoped it would be profound.
The chart is useful beyond language models: any matrix of non-negative weights between two lists of labels fits. Translation alignments, co-occurrence tables, routing weights in a mixture of experts. If it is rows of things that look at columns of things, this is how you see it.
Props
API
<AttentionMap tokens={tokens} weights={weights} />| Prop | Type | Default |
|---|---|---|
tokens | Vec<SharedString>: the keys, along the bottom | required |
queries | Vec<SharedString>: the rows, down the side | the same as tokens |
weights | Vec<f64>, row by row: queries.len() × tokens.len() values | required |
scheme | HeatmapScheme | Blues |
show_values | bool: write each weight in its cell | true up to 12 columns |
class |
The first query is at the top, as in a matrix. Missing weights are left empty.
Run them
10 examples
Press Run to load an example into the live render at the top of the page. Open one to read its code.
How to read it
- 01Rows are queries, columns are keys. Read along a row to see where that token looked.
- 02A dark diagonal means tokens attend to themselves. A dark band just below it means they attend to the previous token.
- 03A dark column means one key is a magnet for everyone else. Check whether it is the first token or a delimiter before you get excited.
- 04In a causal model the cells above the diagonal are empty or zero, because tokens cannot look at the future. A triangle is the correct picture.
- 05Hover a cell for its query, key and weight.
When to use it, and when not to
Reach for it when
- You want to see which input tokens a model attended to while producing each output or hidden state.
- You have a matrix of relationships between two lists of labelled items, for instance source and target words in a translation alignment.
- You are debugging a single head or a short sequence, up to a few dozen tokens, where each cell is still readable.
- You teach, or write, about transformers and want a picture of attention that is honest about being a table.
Skip it when
- The sequence has hundreds of tokens. Cells shrink to dots and the labels overlap. Average over heads, pool, or pick a window.
- You need to say what the model relies on. Attention weights are not an explanation of the output, and the literature has argued about this for years. Treat the map as a hint, not as evidence.
- The rows are not comparable to each other. A matrix of raw scores with different scales per row will light up the largest row and hide the rest. Normalise first.
- The matrix has negative entries that matter. The colour scheme is sequential, so sign is not shown. Use a diverging scheme in a heatmap.
The math behind it
Where the numbers come from
The chart does not compute attention. It draws the numbers you give it. In a standard transformer layer, a token's query is compared with every key by a scaled dot product, and a softmax turns each row of scores into weights that are positive and sum to 1. For a head with key dimension that gives the matrix below. The chart's weights parameter is that matrix, flattened row by row. Row (a query) is drawn at the top when it is first, as in a matrix. Column (a key) is along the bottom.
What gets drawn
With queries and tokens, weights must hold values: the cell at row and column is entry . A missing weight (the list is too short) is left empty rather than drawn as zero. If you leave out queries the rows use the same labels as tokens, which is the self-attention case.
Each cell is a unit square centered at integer coordinates, separated from its neighbours by a 1 pixel gap, and coloured by the weight through a sequential scheme (Blues by default. Pass scheme for another). The axes are labelled Query and Key.
Printing the values
Weights are written in their cells with two decimals when there are 12 or fewer columns, because beyond that the cells are too small for text. Pass show_values={true} to force them on or false to force them off.
A useful sanity check: every row of a softmax attention map sums to 1, so the sum of a row's printed numbers should be close to 1 up to rounding. If it is not, you are drawing pre-softmax scores or a column-wise normalisation, and the picture says something different.
Real life usage
The pronoun problem
Joaquín taught the first week of a machine-learning bootcamp, which meant he had forty students, one projector and the unenviable task of explaining attention in a morning. His previous attempt had involved a matrix equation and a lot of silence.
This time he wanted something with a punchline. He took a small pretrained model, fed it the sentence "The trophy didn't fit in the suitcase because it was too big", extracted the attention weights from one head in the middle layers, and prepared to show that the model knew what "it" meant.
The first picture he drew was a table of numbers with 11 columns and 11 rows. It was true. It was not what anyone in the room could read, and a student in the second row asked whether the red cells meant a bug.
He needed the same matrix as a picture a nervous beginner could scan: tokens down the side, tokens along the bottom, darker where the weight is higher, and each weight in its cell for the one person who wants to check. That is a very specific shape, and he had it already.
He passed the token list and the flat weights to an <AttentionMap> in zopra-charts, with tokens and weights and nothing else. The "it" row had its darkest cell under "trophy", at 0.41, and a weaker one under "suitcase", at 0.22. The room got it immediately. A hand went up: "so it picked the right one?"
He said yes, and then spent ten minutes on the more interesting part. The same head, on the sentence with "small" instead of "big", did not flip. Attention was one ingredient, not the whole decision. The students filled the board with questions, and the survey at the end of the week gave the session an average of 4.7 out of 5, the best of the course, ahead of the free lunch.
Joaquín keeps the first table of numbers in his slides, as a warning. The picture is the same data, he tells the next class, just better dressed.
Pitfalls
- Attention is not explanation. A high weight says the model mixed that token's value in strongly, not that it caused the answer. Be careful what you claim from a map.
- Which head, which layer? Models have dozens of heads, and each map is one of them. Averaging them is common and hides the sharp heads. Say which one you are showing.
- Rows and columns swapped. Weights go row by row, one row per query. If your matrix is transposed the picture still draws, but the story it tells is the reverse.
- Long sequences with values on. Beyond 12 columns the numbers are hidden for a reason. Forcing them on with
show_valuesgives you an unreadable grid.
From around the web