The Attention Mechanism, Explained | Oh My Kode

The Attention Mechanism, Explained

12 Jul 2024

13 minutes read

Read this sentence : “The cat sat on the mat because it was tired.” Who was tired ? You knew instantly it was the cat, not the mat, because your brain paid attention to the right earlier word. The attention mechanism, introduced for translation by Bahdanau (Bahdanau et al., 2015) and then crowned as the single idea behind the Transformer (Vaswani et al., 2017), gives a neural network exactly this superpower : for every word, it decides which other words to look at, and how hard. This post builds attention from an everyday intuition up to the famous formula $\operatorname{softmax}!\big(QK^{\top}/\sqrt{d_k}\big)V$, one small step at a time.

1. The problem : words need context

A word on its own is ambiguous. “Bank” is a riverside or a place for money ; “it” points to something mentioned earlier ; “apple” could be a fruit or a company. To understand a word, a model must mix in information from the other words around it.

Older models (recurrent networks) read a sentence one word at a time, squeezing everything seen so far into a single memory vector. Like whispering a long story down a line of people, the beginning gets forgotten by the end. Attention throws that bottleneck away : it lets every word look directly at every other word, no matter how far apart, all at once.

Attention is a mechanism that computes, for each element of a sequence, a new representation as a weighted average of the other elements, where the weights (how much to "attend") are learned from how relevant each element is to it.

2. The big idea : looking things up

Here is the ELI5 picture. Imagine a tiny library search.

  • You walk in with a question — “something about a tired animal”. That is your query.
  • Every book on the shelf has a label describing what it is about — “cats”, “mats”, “weather”. Those are the keys.
  • Every book also has actual contents inside. Those are the values.

To answer your question you compare it against each label, decide how well each matches, and then read mostly the best-matching books, skimming the rest. Attention is precisely this : compare a query against all keys, turn the match scores into percentages, and blend the values accordingly.

Each element is turned into three vectors : a query $q$ ("what am I looking for ?"), a key $k$ ("what do I offer ?"), and a value $v$ ("what information do I carry ?"). These are produced by multiplying the element's embedding $x$ by three learned matrices, $q = W_{Q}\,x$, $k = W_{K}\,x$, $v = W_{V}\,x$.

3. Attention for a single word, step by step

Let us attend from one query $q$ to a set of keys $k_{1}, \dots, k_{n}$ and values $v_{1}, \dots, v_{n}$. Four small steps.

Step 1 — Score. How relevant is each word to my query ? Measure it with a dot product 1, which is large when two vectors point the same way :

\[\begin{equation} s_{j} = q \cdot k_{j} = \sum_{d} q_{d}\, k_{j,d}. \label{eq:score} \end{equation}\]

In plain words : a high score $s_j$ means “this word is a good match for what I’m looking for”.

Step 2 — Scale. Divide by $\sqrt{d_k}$ (with $d_k$ the vector length). We will justify this shortly ; for now, it just keeps the numbers from getting too big.

Step 3 — Softmax. Turn the scores into positive weights that add up to $1$ (percentages of attention) :

\[\begin{equation} \alpha_{j} = \frac{\exp(s_{j}/\sqrt{d_k})}{\sum_{\ell=1}^{n}\exp(s_{\ell}/\sqrt{d_k})}, \qquad \sum_{j}\alpha_{j}=1. \label{eq:softmax} \end{equation}\]

Step 4 — Blend. Read the values in those proportions :

\[\begin{equation} \text{output} = \sum_{j=1}^{n}\alpha_{j}\, v_{j}. \label{eq:blend} \end{equation}\]

In plain words : the answer is a mix of all the words’ contents, but dominated by the ones that matched. The animation below shows the query “it” scanning the sentence : the bars are the attention weights $\alpha_j$, and “cat” wins.

query : "it" → who / what ? The 3% cat 55% sat 10% because5% it 2% was 5% tired 20% attention weights αⱼ (they sum to 100%) — "it" attends mostly to "cat"
Figure 1 - Attention from the word "it". Each bar is the learned weight $\alpha_j$ from $\eqref{eq:softmax}$. The output for "it" becomes mostly the value of "cat", so the model now "knows" what it refers to.

4. All queries at once : the matrix formula

A sentence has many words, and each is a query. Instead of looping, we stack the query, key and value vectors as rows of matrices $Q$, $K$, $V$ and do it all with two matrix multiplications. This is the celebrated scaled dot-product attention (Vaswani et al., 2017) :

$$ \operatorname{Attention}(Q, K, V) = \operatorname{softmax}\!\left( \frac{Q K^{\top}}{\sqrt{d_k}} \right) V. $$

Let us read it left to right, because every piece is one of the four steps above.

  • $Q K^{\top}$ is the table of all scores at once : entry $(i,j)$ is query $i$ dotted with key $j$, exactly $\eqref{eq:score}$.
  • Dividing by $\sqrt{d_k}$ is the scaling of Step 2.
  • $\operatorname{softmax}$ (applied row by row) turns each row into attention weights that sum to $1$, exactly $\eqref{eq:softmax}$.
  • Multiplying by $V$ blends the values, exactly $\eqref{eq:blend}$, but for every word simultaneously.
Q × Kᵀ QKᵀ/ √dₖ softmax attention weights A × V output scores → weights (each row sums to 1)
Figure 2 - The pipeline of scaled dot-product attention. Scores are computed ($QK^{\top}$), scaled, softmaxed into a weight table $A$, and used to average the values $V$.
Why divide by $\sqrt{d_k}$ ? If queries and keys have independent, unit-variance entries, their dot product $\eqref{eq:score}$ is a sum of $d_k$ random terms, so its variance grows like $d_k$. For large $d_k$ the scores become huge, the softmax saturates (one weight $\approx 1$, the rest $\approx 0$), and gradients vanish. Dividing by $\sqrt{d_k}$ rescales the variance back to $\approx 1$, keeping the softmax in its sensitive, learnable range.

5. Self-attention : a sentence looking at itself

When the queries, keys and values all come from the same sentence ($Q$, $K$, $V$ are three projections of the same words), we call it self-attention. Every word refreshes itself by mixing in the words most relevant to it. The full weight table $A = \operatorname{softmax}(QK^{\top}/\sqrt{d_k})$ is an $n \times n$ heatmap : row $i$ says where word $i$ looks.

keys (looked at) → queries (doing the looking) → Thecatsatonmat Thecatsatonmat darker = more attention · the moving box scans one query (row) at a time
Figure 3 - A self-attention heatmap for "The cat sat on mat". Each row is one word's attention distribution over all words. Notice the diagonal (words attend to themselves) plus meaningful off-diagonal links like sat→cat.

6. Many spotlights : multi-head attention

One attention pattern can only capture one kind of relationship. Real language has several at once : who did what (subject–verb), what modifies what (adjective–noun), what refers to what (pronoun–antecedent). So Transformers run several attentions in parallel, called heads, each with its own $W_Q, W_K, W_V$, and then concatenate the results :

\[\begin{equation} \operatorname{MultiHead}(Q,K,V) = \operatorname{Concat}(\text{head}_1, \dots, \text{head}_h)\, W_{O}, \quad \text{head}_i = \operatorname{Attention}(QW_Q^i, KW_K^i, VW_V^i). \label{eq:multihead} \end{equation}\]

In plain words : each head is a different pair of “reading glasses”. One head might track grammar, another long-range references ; stacking their views gives a richer representation.

7. Why attention changed everything

Two properties made attention the backbone of modern AI.

  • Parallelism. Every word is processed at the same time, not one after another, so training uses modern hardware fully. This is the practical reason it dethroned recurrent networks.
  • Direct long-range links. Any two words are one step apart, so a pronoun can look back at a noun fifty words earlier with no signal decay.

Attention is the engine. To build a full model, we stack it into a repeating block, add a way to encode word order, and sprinkle in some plumbing. That assembled machine is the Transformer, the subject of a companion post. And why can a stack of these blocks represent essentially any function we need ? That question is answered by the Universal Approximation Theorem.

8. Conclusion

Attention is, at heart, a soft, differentiable lookup table : compare a query to keys, softmax the matches into weights, and average the values. Everything else, the scaling by $\sqrt{d_k}$, self-attention, multiple heads, is refinement on that one idea. It replaced the fragile “remember everything in one vector” strategy with a simple instruction that a network can learn : for each word, look at the words that matter. From that sentence-sized intuition grew the models that now write, translate and reason.

References

  1. Bahdanau, D., Cho, K., & Bengio, Y. (2015). Neural Machine Translation by Jointly Learning to Align and Translate. International Conference on Learning Representations (ICLR).
    @inproceedings{Bahdanau2014,
      author = {Bahdanau, Dzmitry and Cho, Kyunghyun and Bengio, Yoshua},
      title = {Neural Machine Translation by Jointly Learning to Align and Translate},
      booktitle = {International Conference on Learning Representations (ICLR)},
      year = {2015}
    }
    
  2. Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., & Polosukhin, I. (2017). Attention Is All You Need. Advances in Neural Information Processing Systems (NeurIPS), 30, 5998–6008.
    @inproceedings{Vaswani2017,
      author = {Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N. and Kaiser, {\L}ukasz and Polosukhin, Illia},
      title = {Attention Is All You Need},
      booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
      volume = {30},
      pages = {5998--6008},
      year = {2017}
    }
    
  1. The dot-product form used here is sometimes called multiplicative attention. Bahdanau’s original 2014 formulation used an additive variant with a small feed-forward network to score query–key pairs ; the two perform comparably, but the multiplicative form is far faster as a matrix product. ↩

who am i

Hi! I am a Data Scientist by profession, an Emacs devotee and an untalented bassist. I intend to use this space for writing about things that I think I have understood well in the hope that they may be helpful to others, including my future self.

what is this

OhMyKode is an opportunity to share knowledge about mathematics, computer science, machine learning and algorithmic beauty, which allows us to improve our skills and learn in depth. It is a sharing place to learn the how and the why.

© MMXVIII - MMXX by Maâmra Youcef - معامره يوسف
Content available under Creative Commons (BY-NC-SA) unless otherwise noted.
This site is hosted at Github Pages and powered by Jekyll & Papyrus.
“We can't skip Math forever !”