Self-Attention in transformers, explained in 5 lines

Module 05 · Transformers · 19 concepts

In depth / 15 lines

Self-Attention, in more detail

Older models read a sentence one word at a time, left to right.

Everything they learned had to be carried forward step by step.

By the time they reached a far-away word, the earlier details had faded.

Self-attention lets every word look straight at every other word.

It does this in one step, however far apart the two words are.

Each word asks a question, and every word offers an answer.

The question is called a query and the offer is called a key.

Matching one query against all the keys gives a score for every word.

A high score means that word matters a lot right now.

Each word also carries a value, which is the content it passes on.

The output is a mix of all the values, weighted by those scores.

The word self means the questions and answers come from the same sentence.

Because every pair is one step apart, distant words are as easy as close ones.

Every position is worked out at the same time, which makes large training possible.

The cost is that doubling the sentence length makes the work four times bigger.

Want the full technical article? Wikipedia ↗

Deep dive

What is underneath Self-Attention

Same input, three times

Self-attention is scaled dot-product attention where the queries, keys and values are all projections of one sequence.

\mathrm{SelfAttn}(X) = \mathrm{softmax}\!\left( \frac{(XW_Q)(XW_K)^{\top}}{\sqrt{d_k}} \right) X W_V

Every position builds a query, compares it against every position's key, and takes a weighted average of every position's value.

The output has the same shape as the input, so these layers stack without any glue between them.

What it replaced, and why

A recurrent layer connects position $i$ to position $j$ through $|i - j|$ sequential steps.

\text{path length} \quad \text{RNN}: O(n) \qquad \text{self-attention}: O(1)

In self-attention every pair of positions is one multiplication apart, so a dependency across 500 tokens is as direct as one across 2.

That also removes the sequential dependency: an RNN must compute position $t$ before $t+1$ , while self-attention computes the whole sequence as one matrix multiply.

That is the property that made training on very large corpora possible, and it matters more than the accuracy improvement did.

It has no idea what order things are in

\mathrm{SelfAttn}(PX) = P \cdot \mathrm{SelfAttn}(X)

$P$ : any permutation matrix.

Permute the input and the output is permuted identically. Nothing inside the operation refers to position.

The set of tokens is all it sees, so "dog bites man" and "man bites dog" produce the same set of outputs in a different order, with no signal that the order changed.

This is not a flaw to be worked around, it is why positional encoding exists as a separate mechanism.

The causal mask

M_{ij} = \begin{cases} 0 & j \leq i \\ -\infty & j > i \end{cases}

A decoder must not let position $i$ read position $i+1$ , or predicting the next token would be a lookup.

Adding this matrix before the softmax sends every forbidden weight to exactly zero afterwards.

It is upper-triangular, and it is the only difference between the attention in an encoder and the attention in a decoder.

One layer, in order

Project the input into queries, keys and values.
Multiply queries by transposed keys.
Scale by the root of the head width, then add the mask.
Softmax each row into weights that sum to one.
Multiply by the values, and project the result back to the model width.

Test yourself

One question on this card. Every option explains itself once you pick, right or wrong.

Question 1

What does the word self mean in self-attention?

Transformers

19 concepts in this module, then a graded test with every answer explained.

Study this module →