Lensa ML

Self-attention in 1 picture: every word looks at every word

Self-attention is the step inside a transformer, the model design behind large language models, where each word gathers context from the rest of the sentence. Every word gives every word, itself included, a weight, and the weights one word gives add up to 1.

2 min read · Lensa ML
The sentence The cat sat on the mat, written twice. Lines run from the query word cat in the bottom row to all six words in the top row, labeled 20%, 22%, 16%, 10%, 19% and 13%, with the strongest line going to cat itself. Below, a 6 by 6 grid shades the attention weight of every word pair, one row per query word.
  1. The picture uses the sentence "The cat sat on the mat". The query, the word doing the looking, is "cat" in the bottom row. Lines run from it to all six words in the top row, called the keys, and each percentage is the weight "cat" gives that word. The green line goes to "cat" itself at 22%, then "The" at 20% and "the" at 19%, down to "on" at 10%. The six weights sum to 100%.
  2. Each word's embedding, a list of numbers that stands for it, becomes three lists called vectors: a query, a key, and a value, the content a word passes on. A pair's score is the dot product of one word's query vector and the other's key vector (multiply matching numbers, then add), divided by √3 since each vector holds 3 numbers. Softmax turns the scores into weights that sum to 1, and the output for "cat" is the weighted average of all six values.
  3. The grid repeats this for every word: each row is one query, each column one key, and brighter cells mean bigger weights. The purple row is the "cat" row drawn above. The brightest gray cell is "sat" looking at "cat", at 24%. These toy 3-number vectors give fairly flat weights, from 10% to 24%. Trained models often use 64 or 128 numbers per attention head, one of several copies of this step running side by side, and draw much sharper patterns.

Takeaway

Self-attention lets every word use the whole sentence in one step, so a model can tie a pronoun to its noun however far apart they sit. The cost is the grid: one cell per pair of words, so doubling the words quadruples the scores. A language model that writes text also hides every cell above the diagonal, so no word sees the words after it.

Go deeper · interactive

How Self-Attention Works

Click each word in the query row and watch its attention lines and its row of the grid change.

Open the explainer

More bites

References