Lecture 9: Transformers
Intro to Machine Learning
[video edited from 3b1b]
Recap: Word embedding
this enables "soft" dictionary look-up:
dict_en2fr = {
"apple" : "pomme",
"banana": "banane",
"lemon" : "citron"}Good word-embeddings space is equipped with semantically meaningful vector arithmetic
Key
Value
apple
pomme
\(:\)
banane
banana
\(:\)
citron
lemon
\(:\)
Python would complain. 🤯
orange
Query
Output
???
query = "orange"
output = dict_en2fr[query]But we can probably see the rationale behind something like this:
Query
Key
Value
Output
orange
apple
\(:\)
pomme
banana
\(:\)
banane
lemon
\(:\)
citron
dict_en2fr = {
"apple" : "pomme",
"banana": "banane",
"lemon" : "citron"}
query = "orange"
output = dict_en2fr[query]0.1
pomme
0.1
banane
0.8
citron
+
+
0.1
pomme
0.1
banane
0.8
citron
+
+
Query
Key
Value
Output
orange
apple
\(:\)
pomme
0.1
pomme
0.1
banane
0.8
citron
banana
\(:\)
banane
lemon
\(:\)
citron
+
+
orange
orange
pomme
banane
citron
+
+
0.8
0.1
0.1
pomme
banane
citron
+
+
- output mixed values:
- compare query and key (dot-product similarity) for merging percentages:
apple
banana
lemon
orange
,
,
orange
orange
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
=[\quad \quad \quad ]
0.1
0.1
0.8
The attention mechanism in transformers mimics this:
We picked these percentages by hand; let's see how a transformer comes up with them.
Outline
- Intuition and overview
- Attention mechanism
- Multi-head attention
- Attention beyond language
- Intuition and overview
- Attention mechanism
- Multi-head attention
- Attention beyond language
Large Language Models (LLMs) are trained in this self-supervised way:
- Scrape the internet for plain texts.
- Cook up “labels” from these texts.
- Use the unhidden part to predict the hidden label.
"To date, the cleverest thinker of all time was Isaac."
feature
label
To date, the
cleverest
To date, the cleverest
thinker
To date, the cleverest thinker
was
\dots
\dots
To date, the cleverest thinker of all time was
Isaac
auto-regressive prediction
[video edited from 3b1b]
\(n\)
\underbrace{\hspace{5.98cm}}
\left\{
\begin{array}{l}
\\
\\
\\
\\
\\
\\
\\
\end{array}
\right.
\(d\)
input embedding
[video edited from 3b1b]
Minimizing cross-entropy loss drives the weights to assign higher probability to the correct next token
word embedding to a softmax distribution over the vocabulary
Transformer
"To date, the cleverest [thinker] of all time was Isaac.
push for Prob("date") to be high
push for Prob("the") to be high
push for Prob("cleverest") to be high
push for Prob("thinker") to be high
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
distribution over the vocabulary
\(\dots\)
\(\dots\)
\(\dots\)
\(\dots\)
To
date
the
cleverest
input embedding
\(\dots\)
\(\dots\)
\(\dots\)
transformer block
transformer block
transformer block
\left\{
\begin{array}{l}
\\
\\
\\
\\
\\
\end{array}
\right.
\(L\) blocks
\(\dots\)
\(\dots\)
output embedding
To
date
the
cleverest
\(\dots\)
transformer block
transformer block
transformer block
x_1
x_2
x_3
x_4
A sequence of \(n\) tokens, each token in \(\mathbb{R}^{d}\)
\(\dots\)
\(\dots\)
\(\dots\)
\(\dots\)
output embedding
input embedding
input embedding
To
date
the
cleverest
[video edited from 3b1b]
input embedding
\(\dots\)
transformer block
output embedding
x_1
x_2
x_3
x_4
\(\dots\)
\(\dots\)
\(\dots\)
\(\dots\)
transformer block
transformer block
each of the \(n\) tokens transformed, block by block
within a shared \(d\)-dimensional word-embedding space.
To
date
the
cleverest
attention layer
MLP
x_1
x_2
x_3
x_4
\(\dots\)
\(\dots\)
\(\dots\)
\(\dots\)
neuron weights
\nabla{\mathcal{L}}
\dots
W_k
W_v
W_q
W_o
input embedding
transformer block
output embedding
To
date
the
cleverest
attention layer
x_1
x_2
x_3
x_4
W_o
attention output projection
W_k
W_v
W_q
\((qkv)\) projection
attention mechanism
word space
\(\mathbb{R}^{d}\)
word space
\(\mathbb{R}^{d}\)
working space
\(\mathbb{R}^{d_k}\)
working space
\(\mathbb{R}^{d_k}\)
1.
2.
e.g., for GPT-2, \(d = 768\), \(d_k = 64\), both hyperparameters
working space \(d_k \ll\) word space \(d\)
W_k
W_v
W_q
x_1
q_1
v_1
k_1
1. (query, key, value) projection — why do we need anything beyond the word space?
To
\(\leftarrow\) word space \(\mathbb{R}^{d}\)
\(\leftarrow\) working space \(\mathbb{R}^{d_k}\)
\(x\) has to play three roles:
- \(q\): what it looks for (\(W_q\) learns how to ask)
- \(k\): what it offers (\(W_k\) learns how to listen)
- \(v\): what it hands over (\(W_v\) learns how to speak)
\(\leftarrow W_q, W_k, W_v \in \mathbb{R}^{d \times d_k}\)
W_k
W_v
W_q
x_2
W_k
W_v
W_q
x_3
W_k
W_v
W_q
x_4
W_k
W_v
W_q
q_2
v_2
k_2
q_3
v_3
k_3
q_4
v_4
k_4
x_1
q_1
v_1
k_1
1. (query, key, value) projection — why do we need anything beyond the word space?
To
date
the
cleverest
\(x\) has to play three roles:
- \(q\): what it looks for (\(W_q\) learns how to ask)
- \(k\): what it offers (\(W_k\) learns how to listen)
- \(v\): what it hands over (\(W_v\) learns how to speak)
same \(W_q, W_k, W_v\) for every token: weight sharing, all in parallel
x_1
q_1
v_1
k_1
W_k
W_v
W_q
x_2
x_3
x_4
q_2
v_2
k_2
q_3
v_3
k_3
q_4
v_4
k_4
- Attention mechanism turns the projected \((q,k,v)\) into \(z\)
- Each \(z\) is context-aware: a mixture of everyone's values, weighted by relevance
W_k
W_v
W_q
W_k
W_v
W_q
W_k
W_v
W_q
2. Attention mechanism
attention mechanism
z_1
z_2
z_3
z_4
To
date
the
cleverest
Outline
- Intuition and overview
- Attention mechanism
- Multi-head attention
- Attention beyond language
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
q_1
q_1
q_1
q_1
k_1
k_2
k_3
k_4
attention mechanism
x_1
x_2
x_3
x_4
z_1
z_2
z_3
z_4
?
To
date
the
cleverest
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
q_1
q_1
q_1
q_1
k_1
k_2
k_3
k_4
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
,
,
,
\Bigg[
\begin{array}{l}
\end{array}
\Bigg.
\Bigg]
\begin{array}{l}
\end{array}
\Bigg.
/\sqrt{d_k}
a_{11}
a_{14}
a_{12}
a_{13}
v_4
v_2
v_3
v_1
a_{11}
a_{14}
a_{12}
a_{13}
=
x_1
x_2
x_3
x_4
z_1
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
dot-product magnitude tend to grow with \(d_k\); without the scaling they saturate the softmax, and gradients get tiny
e.g., softmax\((1, 2, 3, 4) \approx (0.03, 0.09, 0.24, 0.64)\), but softmax\((10, 20, 30, 40) \approx (0, 0, 0, 1)\)
z_1
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
q_1
q_1
q_1
q_1
k_1
k_2
k_3
k_4
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
,
,
,
\Bigg[
\begin{array}{l}
\end{array}
\Bigg.
\Bigg]
\begin{array}{l}
\end{array}
\Bigg.
/\sqrt{d_k}
a_{11}
a_{14}
a_{12}
a_{13}
v_4
v_2
v_3
v_1
a_{11}
a_{14}
a_{12}
a_{13}
=
x_1
x_2
x_3
x_4
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
softmax'd attention scores: mixing percentages that sum to 1 (100%)
\(a_{12}\): how much \(x_1\) attends to \(x_2\) when updating \(x_1\)'s representation
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
q_1
q_1
q_1
q_1
k_1
k_2
k_3
k_4
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
,
,
,
\Bigg[
\begin{array}{l}
\end{array}
\Bigg.
\Bigg]
\begin{array}{l}
\end{array}
\Bigg.
/\sqrt{d_k}
=
v_4
v_2
v_3
v_1
+
+
+
=
a_{11}
a_{14}
a_{12}
a_{13}
a_{11}
a_{14}
a_{12}
a_{13}
x_1
x_2
x_3
x_4
z_1
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
\(\in \mathbb{R}^{d_k}\)
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
q_2
q_2
q_2
q_2
k_1
k_2
k_3
k_4
attention mechanism
x_1
x_2
x_3
x_4
z_1
z_2
z_3
z_4
?
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
=
a_{21}
a_{24}
a_{22}
a_{23}
q_2
q_2
q_2
q_2
k_1
k_2
k_3
k_4
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
,
,
,
\Bigg[
\begin{array}{l}
\end{array}
\Bigg.
\Bigg]
\begin{array}{l}
\end{array}
\Bigg.
/\sqrt{d_k}
v_4
v_2
v_3
v_1
a_{21}
a_{24}
a_{22}
a_{23}
x_1
x_2
x_3
x_4
z_2
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
=
=
a_{21}
a_{24}
a_{22}
a_{23}
q_2
q_2
q_2
q_2
k_1
k_2
k_3
k_4
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
,
,
,
\Bigg[
\begin{array}{l}
\end{array}
\Bigg.
\Bigg]
\begin{array}{l}
\end{array}
\Bigg.
/\sqrt{d_k}
v_4
v_2
v_3
v_1
+
+
+
a_{21}
a_{24}
a_{22}
a_{23}
x_1
x_2
x_3
x_4
z_2
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
q_4
q_4
q_4
q_4
k_1
k_2
k_3
k_4
attention mechanism
x_1
x_2
x_3
x_4
z_1
z_2
z_4
z_3
?
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
=
a_{41}
a_{44}
a_{42}
a_{43}
q_4
q_4
q_4
q_4
k_1
k_2
k_3
k_4
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
,
,
,
\Bigg[
\begin{array}{l}
\end{array}
\Bigg.
\Bigg]
\begin{array}{l}
\end{array}
\Bigg.
/\sqrt{d_k}
v_1
v_2
v_3
v_4
a_{41}
a_{42}
a_{43}
a_{44}
x_1
x_2
x_3
x_4
z_4
parallel and structurally identical processing
can calculate \(z_4\) without \(z_3\)
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
=
a_{41}
a_{44}
a_{42}
a_{43}
q_4
q_4
q_4
q_4
k_1
k_2
k_3
k_4
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
,
,
,
\Bigg[
\begin{array}{l}
\end{array}
\Bigg.
\Bigg]
\begin{array}{l}
\end{array}
\Bigg.
/\sqrt{d_k}
=
v_2
v_3
v_1
v_4
+
+
+
a_{41}
a_{44}
a_{42}
a_{43}
x_1
x_2
x_3
x_4
z_4
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
attention mechanism
x_1
x_2
x_3
x_4
z_1
z_2
z_3
z_4
a head maps the sequence of \(x\)'s to a sequence of context-aware \(z\)'s, in two steps:
1. learned (query, key, value) projection
2. attention mechanism
W_k
W_v
W_q
W_k
W_v
W_q
W_k
W_v
W_q
W_k
W_v
W_q
Attention head
v_1^T
v_2^T
v_3^T
v_4^T
\in \mathbb{R}^{d \times d_k}
W_q
W_k
W_v
X
Q = XW_q
K = XW_k
V = XW_v
1. (query, key, value) projection
q_4^T
q_1^T
q_2^T
q_3^T
k_2^T
k_1^T
k_3^T
k_4^T
\in \mathbb{R}^{d \times d_k}
\in \mathbb{R}^{d \times d_k}
stack each token \(x_i\) as a row
\overbrace{\hspace{3.5cm}}^{d}
\left.\begin{array}{l} \\ \\ \\ \\ \end{array}\right\}n
\underbrace{\hspace{1.43cm}}_{d_k}
\left.\begin{array}{l} \\ \\ \\ \\ \end{array}\right\}n
x_1^T
To
x_2^T
date
x_3^T
the
x_4^T
cleverest
q_4^T
q_1^T
q_2^T
q_3^T
Q
K
v_1^T
v_2^T
v_3^T
v_4^T
V
W_q
W_k
W_v
X
2a. dot-product similarity
q_1^{T}k_3
k_2^T
k_4^T
q_4^T
q_2^T
q_3^T
k_4^T
k_2^T
k_1^T
k_1^T
k_3^T
compare \(q_i\) and \(k_j\): entry \((i, j)\) is \(q_i^Tk_j\)
W_q
W_k
W_v
X
k_2^T
k_1^T
k_3^T
k_4^T
q_4^T
q_1^T
q_2^T
q_3^T
v_1^T
v_2^T
v_3^T
v_4^T
QK^T
\in \mathbb{R}^{n \times n}
2a. dot-product similarity
Q
K
V
assemble all \(n \times n\) similarities: row \(i\) is query \(i\)
A =
\left]
\begin{array}{l}
\\
\\
\\
\\
\\
\end{array}
\right.
a_{41}
a_{42}
a_{43}
a_{44}
=
a_{31}
a_{34}
a_{32}
a_{3 3}
a_{21}
a_{24}
a_{22}
a_{23}
a_{11}
a_{14}
a_{12}
a_{13}
\in \mathbb{R}^{n \times n}
each row sums up to 1
(
)
softmax
/\sqrt{d_k}
(
)
softmax
/\sqrt{d_k}
(
)
softmax
/\sqrt{d_k}
(
)
softmax
/\sqrt{d_k}
W_q
W_k
W_v
X
k_2^T
k_1^T
k_3^T
k_4^T
q_4^T
q_1^T
q_2^T
q_3^T
v_1^T
v_2^T
v_3^T
v_4^T
\left[
\begin{array}{l}
\\
\\
\\
\\
\\
\end{array}
\right.
=
QK^T
/\sqrt{d_k}
softmax
(
)
row
\(A\)
2b. attention matrix
Q
K
V
a_{41}
a_{42}
a_{43}
a_{44}
a_{31}
a_{34}
a_{32}
a_{3 3}
a_{21}
a_{24}
a_{22}
a_{23}
a_{11}
a_{14}
a_{12}
a_{13}
W_q
W_k
W_v
X
k_2^T
k_1^T
k_3^T
k_4^T
q_4^T
q_1^T
q_2^T
q_3^T
v_1^T
v_2^T
v_3^T
v_4^T
Z = AV
2c. attention-weighted values \(Z\)
Q
K
V
A
QK^T
/\sqrt{d_k}
softmax
(
)
row
=
q_1
a_{41}
a_{42}
a_{43}
a_{44}
a_{31}
a_{34}
a_{32}
a_{3 3}
a_{21}
a_{24}
a_{22}
a_{23}
a_{11}
a_{14}
a_{12}
a_{13}
q_1
q_2
q_3
q_4
k_4
k_1
k_2
k_3
attention mechanism
x_1
x_2
x_3
x_4
attention scores depend on the (query, key) only
recap
a_{41}
a_{42}
a_{43}
a_{44}
a_{31}
a_{34}
a_{32}
a_{3 3}
a_{21}
a_{24}
a_{22}
a_{23}
a_{11}
a_{14}
a_{12}
a_{13}
q_1
q_1
q_2
q_3
q_4
k_4
k_1
k_2
k_3
attention mechanism
\in \mathbb{R}^{d_k}
x_1
x_2
x_3
x_4
z_3
+
+
+
a_{31}
a_{34}
a_{32}
a_{33}
=
v_4
v_2
v_3
v_1
v_1
v_2
v_3
v_4
v_4
v_4
v_4
\(z_3\) is meant to help predict \(x_4\)
but \(z_3\) already uses \(v_4\): it peeks at the answer
recap
v_1
k_1
q_1
v_2
k_2
q_2
v_3
k_3
q_3
v_4
q_4
k_4
v_4
v_4
v_4
q_1
=
=
a_{31}
a_{34}
a_{32}
a_{33}
q_3
q_3
q_3
q_3
k_1
k_2
k_3
k_4
softmax
\Bigg(
\begin{array}{l}
\end{array}
\Bigg.
\Bigg)
\begin{array}{l}
\end{array}
\Bigg.
,
,
,
\Bigg[
\begin{array}{l}
\end{array}
\Bigg.
\Bigg]
\begin{array}{l}
\end{array}
\Bigg.
/\sqrt{d_k}
v_4
v_2
v_3
v_1
+
+
+
a_{31}
a_{34}
a_{32}
a_{33}
x_1
x_2
x_3
x_4
z_3
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
d_k
To
date
the
cleverest
masking
mask the future:
\(q_3^Tk_4 \to -\infty\), so \(a_{34} = 0\)
\phantom{a_{34}}
\phantom{a_{34}}
0
0
-\infty
masking: no token can peek ahead
dot-product similarity scores
\tfrac{1}{\sqrt{d_k}}
\textcolor{#b8860b}{q_1^T}\textcolor{#ac4477}{k_1}
\textcolor{#b8860b}{q_1^T}\textcolor{#ac4477}{k_2}
\textcolor{#b8860b}{q_1^T}\textcolor{#ac4477}{k_3}
\textcolor{#b8860b}{q_1^T}\textcolor{#ac4477}{k_4}
\textcolor{#b8860b}{q_2^T}\textcolor{#ac4477}{k_1}
\textcolor{#b8860b}{q_2^T}\textcolor{#ac4477}{k_2}
\textcolor{#b8860b}{q_2^T}\textcolor{#ac4477}{k_3}
\textcolor{#b8860b}{q_2^T}\textcolor{#ac4477}{k_4}
\textcolor{#b8860b}{q_3^T}\textcolor{#ac4477}{k_1}
\textcolor{#b8860b}{q_3^T}\textcolor{#ac4477}{k_2}
\textcolor{#b8860b}{q_3^T}\textcolor{#ac4477}{k_3}
\textcolor{#b8860b}{q_3^T}\textcolor{#ac4477}{k_4}
\textcolor{#b8860b}{q_4^T}\textcolor{#ac4477}{k_1}
\textcolor{#b8860b}{q_4^T}\textcolor{#ac4477}{k_2}
\textcolor{#b8860b}{q_4^T}\textcolor{#ac4477}{k_3}
\textcolor{#b8860b}{q_4^T}\textcolor{#ac4477}{k_4}
+
causal mask \(M\)
0
-\infty
-\infty
-\infty
0
0
-\infty
-\infty
0
0
0
-\infty
0
0
0
0
=
\tfrac{1}{\sqrt{d_k}}
masked scores
\textcolor{#b8860b}{q_1^T}\textcolor{#ac4477}{k_1}
-\infty
-\infty
-\infty
\textcolor{#b8860b}{q_2^T}\textcolor{#ac4477}{k_1}
\textcolor{#b8860b}{q_2^T}\textcolor{#ac4477}{k_2}
-\infty
-\infty
\textcolor{#b8860b}{q_3^T}\textcolor{#ac4477}{k_1}
\textcolor{#b8860b}{q_3^T}\textcolor{#ac4477}{k_2}
\textcolor{#b8860b}{q_3^T}\textcolor{#ac4477}{k_3}
-\infty
\textcolor{#b8860b}{q_4^T}\textcolor{#ac4477}{k_1}
\textcolor{#b8860b}{q_4^T}\textcolor{#ac4477}{k_2}
\textcolor{#b8860b}{q_4^T}\textcolor{#ac4477}{k_3}
\textcolor{#b8860b}{q_4^T}\textcolor{#ac4477}{k_4}
\xrightarrow{\text{softmax}_\text{row}}
causal attention
matrix \(A\)
a_{11}
0
0
0
a_{21}
a_{22}
0
0
a_{31}
a_{32}
a_{33}
0
a_{41}
a_{42}
a_{43}
a_{44}
\(Z = AV\), same as before
each \(z_i\) now mixes only \(v_1, \dots, v_i\)
Outline
- Intuition and overview
- Attention mechanism
- Multi-head attention
- Attention beyond language
z_1^1
z_2^1
z_3^1
z_4^1
W^1_k
W^1_v
W^1_q
v^1_1
k^1_1
q^1_1
W^1_k
W^1_v
W^1_q
v^1_2
k^1_2
q^1_2
W^1_k
W^1_v
W^1_q
W^1_k
W^1_v
W^1_q
v^1_3
k^1_3
q^1_3
q^1_4
k^1_4
v^1_4
attention mechanism
x_1
x_2
x_3
x_4
One attention head
z_1^1
z_2^1
z_3^1
z_4^1
W^1_k
W^1_v
W^1_q
v^1_1
k^1_1
q^1_1
W^1_k
W^1_v
W^1_q
v^1_2
k^1_2
q^1_2
W^1_k
W^1_v
W^1_q
W^1_k
W^1_v
W^1_q
v^1_3
k^1_3
q^1_3
q^1_4
k^1_4
v^1_4
z_1^2
z_2^2
z_3^2
z_4^2
W^2_k
W^2_v
W^2_q
v^2_1
k^2_1
q^2_1
W^2_k
W^2_v
W^2_q
v^2_2
k^2_2
q^2_2
W^2_k
W^2_v
W^2_q
W^2_k
W^2_v
W^2_q
v^2_3
k^2_3
q^2_3
q^2_4
k^2_4
v^2_4
attention mechanism
x_1
x_2
x_3
x_4
Two attention heads
head 2 learns its own \(\textcolor{#6aa84f}{W_q^2, W_k^2, W_v^2}\), so it can attend to different things
z_1^1
z_2^1
z_3^1
z_4^1
W^1_k
W^1_v
W^1_q
v^1_1
k^1_1
q^1_1
W^1_k
W^1_v
W^1_q
v^1_2
k^1_2
q^1_2
W^1_k
W^1_v
W^1_q
W^1_k
W^1_v
W^1_q
v^1_3
k^1_3
q^1_3
q^1_4
k^1_4
v^1_4
z_1^2
z_2^2
z_3^2
z_4^2
W^2_k
W^2_v
W^2_q
v^2_1
k^2_1
q^2_1
W^2_k
W^2_v
W^2_q
v^2_2
k^2_2
q^2_2
W^2_k
W^2_v
W^2_q
W^2_k
W^2_v
W^2_q
v^2_3
k^2_3
q^2_3
q^2_4
k^2_4
v^2_4
z_1^3
z_2^3
z_3^3
z_4^3
W^3_k
W^3_v
W^3_q
v^3_1
k^3_1
q^3_1
W^3_k
W^3_v
W^3_q
v^3_2
k^3_2
q^3_2
W^3_k
W^3_v
W^3_q
W^3_k
W^3_v
W^3_q
v^3_3
k^3_3
q^3_3
q^3_4
k^3_4
v^3_4
z_1^H
z_2^H
z_3^H
z_4^H
W^H_k
W^H_v
W^H_q
v^H_1
k^H_1
q^H_1
W^H_k
W^H_v
W^H_q
v^H_2
k^H_2
q^H_2
W^H_k
W^H_v
W^H_q
W^H_k
W^H_v
W^H_q
v^H_3
k^H_3
q^H_3
q^H_4
k^H_4
v^H_4
attention mechanism
\dots
x_1
x_2
x_3
x_4
H attention heads
In particular, each head:
- learns its own \(W_q, W_k, W_v\)
- creates its own sequence of \(q, k, v\)
- computes its own sequence of \(z\)
Parallel, and structurally identical processing across all heads and tokens.
Multi-head Attention
index along heads
index along sequence
\dots
x_1
x_2
x_3
x_4
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
concatenate each token's \(H\) head outputs \(z_i^h \in \mathbb{R}^{d_k}\): \(z_i^{\text{concat}} \in \mathbb{R}^{H d_k}\)
W^1_k
W^1_v
W^1_q
v^1_1
k^1_1
q^1_1
W^1_k
W^1_v
W^1_q
v^1_2
k^1_2
q^1_2
W^1_k
W^1_v
W^1_q
W^1_k
W^1_v
W^1_q
v^1_3
k^1_3
q^1_3
q^1_4
k^1_4
v^1_4
attention mechanism
z_1^1
z_2^1
z_3^1
z_4^1
W^2_k
W^2_v
W^2_q
v^2_1
k^2_1
q^2_1
W^2_k
W^2_v
W^2_q
v^2_2
k^2_2
q^2_2
W^2_k
W^2_v
W^2_q
W^2_k
W^2_v
W^2_q
v^2_3
k^2_3
q^2_3
q^2_4
k^2_4
v^2_4
attention mechanism
z_1^2
z_2^2
z_3^2
z_4^2
W^3_k
W^3_v
W^3_q
v^3_1
k^3_1
q^3_1
W^3_k
W^3_v
W^3_q
v^3_2
k^3_2
q^3_2
W^3_k
W^3_v
W^3_q
W^3_k
W^3_v
W^3_q
v^3_3
k^3_3
q^3_3
q^3_4
k^3_4
v^3_4
attention mechanism
z_1^3
z_2^3
z_3^3
z_4^3
W^H_k
W^H_v
W^H_q
v^H_1
k^H_1
q^H_1
W^H_k
W^H_v
W^H_q
v^H_2
k^H_2
q^H_2
W^H_k
W^H_v
W^H_q
W^H_k
W^H_v
W^H_q
v^H_3
k^H_3
q^H_3
q^H_4
k^H_4
v^H_4
attention mechanism
z_1^H
z_2^H
z_3^H
z_4^H
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
z_1^{\text{concat}}
z_2^{\text{concat}}
z_3^{\text{concat}}
z_4^{\text{concat}}
z_1^1
z_2^1
z_3^1
z_4^1
W^1_k
W^1_v
W^1_q
v^1_1
k^1_1
q^1_1
W^1_k
W^1_v
W^1_q
v^1_2
k^1_2
q^1_2
W^1_k
W^1_v
W^1_q
W^1_k
W^1_v
W^1_q
v^1_3
k^1_3
q^1_3
q^1_4
k^1_4
v^1_4
attention mechanism
z_1^2
z_2^2
z_3^2
z_4^2
W^2_k
W^2_v
W^2_q
v^2_1
k^2_1
q^2_1
W^2_k
W^2_v
W^2_q
v^2_2
k^2_2
q^2_2
W^2_k
W^2_v
W^2_q
W^2_k
W^2_v
W^2_q
v^2_3
k^2_3
q^2_3
q^2_4
k^2_4
v^2_4
attention mechanism
z_1^3
z_2^3
z_3^3
z_4^3
W^3_k
W^3_v
W^3_q
v^3_1
k^3_1
q^3_1
W^3_k
W^3_v
W^3_q
v^3_2
k^3_2
q^3_2
W^3_k
W^3_v
W^3_q
W^3_k
W^3_v
W^3_q
v^3_3
k^3_3
q^3_3
q^3_4
k^3_4
v^3_4
attention mechanism
z_1^H
z_2^H
z_3^H
z_4^H
W^H_k
W^H_v
W^H_q
v^H_1
k^H_1
q^H_1
W^H_k
W^H_v
W^H_q
v^H_2
k^H_2
q^H_2
W^H_k
W^H_v
W^H_q
W^H_k
W^H_v
W^H_q
v^H_3
k^H_3
q^H_3
q^H_4
k^H_4
v^H_4
attention mechanism
\dots
x_1
x_2
x_3
x_4
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
W_o
attention output projection
o_1
o_2
o_3
o_4
\(\in \mathbb{R}^{d}\)
\(\textcolor{#6aa84f}{W_o}\): the learned mixer, blending what the heads found back into the word space \(\mathbb{R}^d\)
\underbrace{Z^{\text{concat}}}_{n \times H d_k} \; \underbrace{\textcolor{#6aa84f}{W_o}}_{H d_k \times d} = \underbrace{O}_{n \times d}
typically \(H \cdot d_k = d\)
(e.g. GPT-2: 12 heads, \(12 \times 64 = 768,\) all hyperparameters)
z_1^{\text{concat}}
z_2^{\text{concat}}
z_3^{\text{concat}}
z_4^{\text{concat}}
\in \mathbb{R}^{Hd_k\times d}
z_1^1
z_2^1
z_3^1
z_4^1
W^1_k
W^1_v
W^1_q
v^1_1
k^1_1
q^1_1
W^1_k
W^1_v
W^1_q
v^1_2
k^1_2
q^1_2
W^1_k
W^1_v
W^1_q
W^1_k
W^1_v
W^1_q
v^1_3
k^1_3
q^1_3
q^1_4
k^1_4
v^1_4
attention mechanism
z_1^2
z_2^2
z_3^2
z_4^2
W^2_k
W^2_v
W^2_q
v^2_1
k^2_1
q^2_1
W^2_k
W^2_v
W^2_q
v^2_2
k^2_2
q^2_2
W^2_k
W^2_v
W^2_q
W^2_k
W^2_v
W^2_q
v^2_3
k^2_3
q^2_3
q^2_4
k^2_4
v^2_4
attention mechanism
z_1^3
z_2^3
z_3^3
z_4^3
W^3_k
W^3_v
W^3_q
v^3_1
k^3_1
q^3_1
W^3_k
W^3_v
W^3_q
v^3_2
k^3_2
q^3_2
W^3_k
W^3_v
W^3_q
W^3_k
W^3_v
W^3_q
v^3_3
k^3_3
q^3_3
q^3_4
k^3_4
v^3_4
attention mechanism
z_1^H
z_2^H
z_3^H
z_4^H
W^H_k
W^H_v
W^H_q
v^H_1
k^H_1
q^H_1
W^H_k
W^H_v
W^H_q
v^H_2
k^H_2
q^H_2
W^H_k
W^H_v
W^H_q
W^H_k
W^H_v
W^H_q
v^H_3
k^H_3
q^H_3
q^H_4
k^H_4
v^H_4
attention mechanism
\dots
x_1
x_2
x_3
x_4
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
\left\{
\begin{array}{l}
\\
\\
\end{array}
\right.
W_o
attention output projection
\in \mathbb{R}^{Hd_k\times d}
o_1
o_2
o_3
o_4
all in word space \(\mathbb{R}^{d}\)
multi-head attention
attention layer
x_1
x_2
x_3
x_4
W_o
attention output projection
W_k^h
W_v^h
W_q^h
\((qkv)\) projection, \(H\) heads
attention mechanism, \(H\) heads
word space
\(\mathbb{R}^{d}\)
word space
\(\mathbb{R}^{d}\)
each head's working space
\(\mathbb{R}^{d_k}\)
concat heads
\(z_i^{\text{concat}} \in \mathbb{R}^{Hd_k}\)
1.
2.
o_1
o_2
o_3
o_4
multi-head attention layer, zoomed out
single-head version: slide 17
zoomed out further, inside a transformer block: slide 16
Shape Example:
\(W\)s are the learned weights
\(n\)
num tokens
5
\(d\)
word-space dim
6
\(X\)
input
\(n \times d\)
\(5 \times 6\)
\(W_q^h\)
query proj
\(d \times d_k\)
\(6 \times 3\)
\(W_k^h\)
key proj
\(d \times d_k\)
\(6 \times 3\)
\(W_v^h\)
value proj
\(d \times d_k\)
\(6 \times 3\)
\(Q^h\)
query
\(n \times d_k\)
\(5 \times 3\)
\(K^h\)
key
\(n \times d_k\)
\(5 \times 3\)
\(V^h\)
value
\(n \times d_k\)
\(5 \times 3\)
\(A^h\)
attn matrix
\(n \times n\)
\(5 \times 5\)
\(Z^h\)
attn head out
\(n \times d_k\)
\(5 \times 3\)
for a single
attention head
\left\{
\begin{array}{l}
\\
\\
\\
\\
\\
\\
\end{array}
\right.
\(H\)
num heads
2
\(Z^{\text{concat}}\)
concatenated heads
\(n \times Hd_k\)
\(5 \times 6\)
\(W_o\)
output proj
\(Hd_k \times d\)
\(6 \times 6\)
\(O\)
attn layer out
\(n \times d\)
\(5 \times 6\)
\(d_k\)
working-space dim
3
Practical techniques
Layer normalization
Residual connection
Positional encoding
Explore interactively: Transformer Explainer
attention alone ignores token order; positional encoding adds each token's position into \(x_i\)
we'll work through these in the lab
Outline
- Intuition and overview
- Attention mechanism
- Multi-head attention
- Attention beyond language
image credit: Nicholas Pfaff
Generative Boba by Boyuan Chen (now at OpenAI)
😉
😉
We can tokenize anything.
General strategy: chop the input up into chunks, project each chunk to an embedding
[images credit: visionbook.mit.edu]
Multi-modality (image q&a)
- (query, key, value) come from different input modality
- cross-attention
[images credit: visionbook.mit.edu]
Biology: attention learns protein structure
- large \(a_{ij}\) ≈ amino acids \(i\) and \(j\) touch once the protein folds; no one taught it folding
- why it matters: a protein's 3D shape decides what it does (e.g., which drugs bind); sequences are cheap to read, shapes take months to measure in a lab
- ESMFold and AlphaFold predict the full 3D shape from sequence
Attention map highlights objects (DINO)
Attention maps can have artifacts
input
ViT (DINOv2)
same ViT + 4 registers
a few background patches get hijacked as scratch memory (bright dots, middle column)
blank "register" tokens give the model scratch space instead, and the maps clean up (right column)
Attention sinks: a big attention score isn't always meaningful
- average attention logits in Llama-2 (grey: masked future)
- each row must hand out 100%; every row can see the first token, so heads park leftover attention there
- so a big \(a_{i1}\) is a parking spot, not importance (like the registers)
- emerged on its own; later models add a dedicated sink token on purpose
Does a head matter? Disable it and check
task: retrieve one fact hidden in a long document
Removing retrieval heads hurts more than removing random heads.
x: document length (5K–100K tokens) · y: where the fact sits (0.1 = start, 0.9 = end) · green = retrieved correctly
TabPFN: no training on your table; it attends to it at prediction time
trained once on millions of synthetic tables; on a new table: no training, just one forward pass
test row = query, training rows = keys, their labels = values: the soft dictionary lookup from slide 3
Summary
Attention creates context-aware representations: learned \(\textcolor{#6aa84f}{W_q, W_k, W_v}\) project each token into a query, key, and value; queries and keys set the mixing percentages, values supply what gets combined.
Multiple heads each learn their own \(\textcolor{#6aa84f}{W_q, W_k, W_v}\) and working space \(\mathbb{R}^{d_k}\); \(\textcolor{#6aa84f}{W_o}\) maps the concatenated outputs back to word space \(\mathbb{R}^d\).
Training processes tokens in parallel: causal masking prevents access to future tokens.
Attention extends beyond language: the same mechanism supports vision, proteins, and tabular data.
Human-like concepts emerge in LLMs
- from 4.7M odd-one-out judgments, each model gets 66 interpretable concept dimensions: animal, food, vehicle, ...
- these line up with human judgments, and with brain activity in category-selective regions
parting thoughts
6.390 IntroML (Fall26) - Lecture 9 Transformers
By Shen Shen
