Part 9

Attention

Mehmet Kerem Turkcan
Associate Research Scientist
Center for Smart Streetscapes, Columbia University
New York, USA
keremturkcan.com; mkt2126@columbia.edu

Move through the deck with (i) the arrow keys, (ii) a presentation clicker, or (iii) the buttons at the bottom left. The gear at the bottom right opens the slide settings.

Part 8 Overview

  1. Convolution combines nearby pixel values with one filter.
  2. In their 2020 manuscript, Alexey Dosovitskiy and eleven coauthors divided an image into patches and mapped each patch to a token vector.
  3. A Transformer pipeline can process image token vectors and word token vectors.
  4. Attention computes a weighted sum of value vectors for each token position.

How does a model calculate which permitted token representations contribute and the contribution from each one?

A relationship inside a sentence

Lina moved the glass because it was fragile.
Lina moved the _____ because it was fragile.
What does it refer to?
query from “it”What information fits? key from “Lina” key from “glass” candidate score larger score

Answer: it refers to glass.

Attention can incorporate information from the value vector at glass into the representation at it.

Recurrent hidden states and direct paths

An RNN updates its hidden state once per token. Information from token 1 reaches token 6 through five consecutive state updates.

sequence computationpath from token 1 to token 6training schedule within one layer
RNN1\to2\to3\to4\to5\to6step 6 waits for steps 1 through 5
self attentiontoken 6 scores token 1 directlyall query rows use matrix multiplication together
\text{one query row}\;\times\;\text{all key rows}\;\longrightarrow\;\text{one row of scores}

Each Transformer layer updates the token representations, so the stacked layers revise those representations repeatedly. Within one self attention layer, every permitted pair of positions has a direct scoring path.

Tutor prompt: ask students to count arrows from token 1 to token 6 before introducing any matrix notation.

Operations in one attention head

One attention head, a single copy of the calculation, contains five operations.

token vectorsX 1. projectQ, K, Vthree matrices 2. scoreQKᵀ 3. maskpermittedpositions 4. softmaxweights A 5. weightedsum AV (i) project, (ii) score, (iii) mask, (iv) apply softmax, and (v) compute the weighted sumThe five stages use (i) matrix multiplication, (ii) addition, (iii) exponentiation, and (iv) rowwise division.

Query, key, and value

In a search system, a request is compared with index labels, and a successful match retrieves content. Attention uses the standard terms (i) query, (ii) key, and (iii) value for analogous roles.

standard termquestion for one token positionmathematical job
query, q_iWhat information would help position i?compared with every permitted key
key, k_jWhat kind of information does position j offer?produces a compatibility score with a query
value, v_jWhat content can position j contribute?enters the weighted output sum
x_iW_Q=q_i,\qquad x_iW_K=k_i,\qquad x_iW_V=v_i

One hidden row x_i produces all three vectors through three learned projection matrices.

Illustrative head dimensions

The illustrative token window contains (i) red, (ii) ball, and (iii) rolls. Each input row has three entries. The head produces (i) a query with one entry, (ii) a key with one entry, and (iii) a value with two entries.

symbolworked valuemeaning
n3three token positions
d_{model}3three entries in each input row
d_k1one entry in each query and key
d_v2two entries in each value
W_Q:3\times1,\qquad W_K:3\times1,\qquad W_V:3\times2

With these dimensions, every multiplication can be written explicitly.

Three input rows

For the illustrative window red ball rolls, use a representation with three slots: each token receives a row containing (i) one entry equal to 1 and (ii) two entries equal to 0.

positiontokeninput row x_i
1red(1,0,0)
2ball(0,1,0)
3rolls(0,0,1)
X=\begin{pmatrix}1&0&0\\0&1&0\\0&0&1\end{pmatrix}=I_3

In a trained model, hidden rows are dense vectors computed from tokens and earlier layers. Multiplication by the identity matrix I_3 leaves a matrix unchanged, which exposes the projection arithmetic.

Three projections

Parameter values chosen by hand
\begin{aligned}W_Q&=\begin{pmatrix}0\\0\\1\end{pmatrix},&W_K&=\begin{pmatrix}0\\\ln 2\\0\end{pmatrix},&W_V&=\begin{pmatrix}1&0\\0&2\\1&1\end{pmatrix}.\\[-2pt]Q&=XW_Q=W_Q,&K&=XW_K=W_K,&V&=XW_V=W_V.\end{aligned}
token rowquery q_ikey k_ivalue v_i
red00(1,0)
ball0\ln2(0,2)
rolls10(1,1)

The number \ln2\approx0.693 satisfies e^{\ln2}=2.

How does rolls score each key?

Because each query and key has one entry, every compatibility score is one multiplication: s_{ij}=q_i k_j.

The query from rolls is 1. Compare it with the three keys, (0,\ln2,0), one at a time.

key from reds_{31}=1(0)=0
key from balls_{32}=1(\ln2)=\ln2
key from rollss_{33}=1(0)=0
\text{score row from rolls}=(0,\ln2,0)

The head has d_k=1, so division by \sqrt{d_k}=1 leaves these scores unchanged.

Score matrix for all queries

Matrix multiplication computes every query row at once. Matrix rows represent queries, and matrix columns represent keys.

red and ball0(0,\ln2,0)=(0,0,0) for both query rows
query from rolls1(0,\ln2,0)=(0,\ln2,0)
S=QK^{\mathsf T}=\begin{pmatrix}0&0&0\\0&0&0\\0&\ln2&0\end{pmatrix}

The entry in row 3, column 2 records the score from the rolls query to the ball key.

Which positions may each query use?

A model that predicts the next token may use its current position and earlier positions. The causal mask assigns -\infty to every future column before softmax.

M=\begin{pmatrix}0&-\infty&-\infty\\0&0&-\infty\\0&0&0\end{pmatrix}

The symbol -\infty represents an extremely negative score. Because e^{-\infty}=0, softmax assigns zero weight to every masked entry.

queryredballrolls
redallowedfuturefuture
ballallowedallowedfuture
rollsallowedallowedallowed
S+M=\begin{pmatrix}0&-\infty&-\infty\\0&0&-\infty\\0&\ln2&0\end{pmatrix}

Decoders that predict the next token use this triangular mask. In their 2020 Vision Transformer manuscript, Dosovitskiy et al. allowed every patch token to attend to every patch token.

Softmax for the rolls row

The last query, from rolls, may use all three columns. Applying softmax to these scores produces positive weights whose sum is 1.

scores(0,\ln2,0)
exponentiate(e^0,e^{\ln2},e^0)=(1,2,1)
add1+2+1=4
divide by 4(1/4,2/4,1/4)=(1/4,1/2,1/4)
a_{\text{rolls}}=\operatorname{softmax}(0,\ln2,0)=\left(\frac14,\frac12,\frac14\right)

The coefficient for ball is one half; the coefficients for red and rolls are each one quarter.

All attention weights

query rowpermitted exponentiated scoresdivide by their totalweight row
red(1)(1)/1(1,0,0)
ball(1,1)(1,1)/2(1/2,1/2,0)
rolls(1,2,1)(1,2,1)/4(1/4,1/2,1/4)
A=\operatorname{softmax}_{\mathrm{row}}(S+M)=\begin{pmatrix}1&0&0\\\frac12&\frac12&0\\\frac14&\frac12&\frac14\end{pmatrix}
Two invariants hold: (i) every row totals 1, and (ii) every masked future entry has weight 0.

Weighted sum of the value rows

The output for rolls is the weighted sum of three value rows, each with two entries.

source tokenweightvalueweighted contribution
red1/4(1,0)(1/4,0)
ball1/2(0,2)(0,1)
rolls1/4(1,1)(1/4,1/4)
\begin{aligned}o_{\text{rolls}}&=\frac14(1,0)+\frac12(0,2)+\frac14(1,1)\\&=(1/4,0)+(0,1)+(1/4,1/4)=(1/2,5/4).\end{aligned}

The head's output representation for rolls is (0.5,1.25). A second head can produce another output row for the same token.

Output matrix for all queries

red output1(1,0)=(1,0)
ball output\frac12(1,0)+\frac12(0,2)=(1/2,1)
rolls output\frac14(1,0)+\frac12(0,2)+\frac14(1,1)=(1/2,5/4)
O=AV=\begin{pmatrix}1&0&0\\\frac12&\frac12&0\\\frac14&\frac12&\frac14\end{pmatrix}\begin{pmatrix}1&0\\0&2\\1&1\end{pmatrix}=\begin{pmatrix}1&0\\\frac12&1\\\frac12&\frac54\end{pmatrix}.

Matrix multiplication performs the three weighted sums together. Row i of A supplies the coefficients for the permitted value rows in output row i of O.

Tutor move: after covering the matrix equation, compute one output row with students before showing how repeating that weighted sum fills the entire matrix.

Does attention choose one token?

Keeping only the largest weight

The largest score in the rolls row belongs to ball. Keeping that value alone would produce v_{\text{ball}}=(0,2).

Scaled dot product attention applies softmax and retains every permitted value with its calculated weight.

(1/4,1/2,1/4)V=(1/2,5/4)

The standard result is the weighted sum (0.5,1.25).

Question: when would the shortcut become a close numerical approximation?

When one softmax weight lies very near 1 and the other weights lie very near 0.

The complete formula

  1. Project: compute (i) Q=XW_Q, (ii) K=XW_K, and (iii) V=XW_V.
  2. Score: S=QK^{\mathsf T}/\sqrt{d_k}.
  3. Mask: add M, which records the permitted key positions for every query.
  4. Softmax: compute every row of A=\operatorname{softmax}_{\mathrm{row}}(S+M).
  5. Weighted sum: compute O=AV.
\operatorname{Attention}(Q,K,V)=\operatorname{softmax}\!\left(\frac{QK^{\mathsf T}}{\sqrt{d_k}}+M\right)V
arrayshapemeaning
Xn\times d_{model}n token rows
Q,Kn\times d_kquery and key rows
Vn\times d_vvalue rows
S,M,An\times none entry per token pair
On\times d_vone output row per token

Why divide by the square root?

Illustrative comparison with d_k=4
calculationscore for key 1score for key 2softmax weights
raw dot products42(0.881,0.119)
divide by \sqrt4=221(0.731,0.269)

Each score adds d_k products of corresponding query and key entries. As d_k grows, raw score gaps often widen. Scaling keeps several weights able to contribute.

\frac{q\cdot k}{\sqrt{d_k}}

Division by \sqrt{d_k} reduces the effect of head width on the typical score spread.

Because this head has d_k=1, division by \sqrt{d_k}=1 leaves every score unchanged.

A second attention head

Head 1 produced o^{(1)}_{\text{rolls}}=(1/2,5/4). Suppose that a second illustrative head assigns the permitted weight row a^{(2)}_{\text{rolls}}=(1/2,1/4,1/4).

source tokenhead 2 weightvalueweighted contribution
red1/2(1,0)(1/2,0)
ball1/4(0,2)(0,1/2)
rolls1/4(1,1)(1/4,1/4)
o^{(2)}_{\text{rolls}}=(1/2,0)+(0,1/2)+(1/4,1/4)=(3/4,3/4)

The two heads now provide four numbers for the rolls row: (1/2,5/4) and (3/4,3/4).

Parallel attention heads

Each head has separate projection matrices, so multiplying the same input rows by those matrices produces several learned coordinate representations.

input X attention head 1 attention head 2 attention head H concatenateplace rows side by side outputprojection WO
\operatorname{MultiHead}(X)=\operatorname{Concat}(\operatorname{head}_1,\ldots,\operatorname{head}_H)W_O

A learned head may become useful for a recurring pattern. A simple human label summarizes observed examples. The same head can exhibit another pattern for another input.

Concatenation with the output projection

Concatenation places the two head outputs beside each other.

c_{\text{rolls}}=(1/2,5/4,3/4,3/4)

Choose this illustrative output projection, whose shape is 4\times3:

W_O=\begin{pmatrix}1&0&0\\0&1&0\\0&0&1\\0&0&1\end{pmatrix}

Multiplication by W_O maps four concatenated entries back to d_{model}=3 entries.

m_{\text{rolls}}=c_{\text{rolls}}W_O=(1/2,5/4,3/2)

A residual connection adds the input row x_{\text{rolls}}=(0,0,1).

y_{\text{rolls}}=x_{\text{rolls}}+m_{\text{rolls}}=(1/2,5/4,5/2)

The output projection restores the model width, which makes the residual addition possible.

One Transformer layer

In 2017, Vaswani et al. arranged attention inside each Transformer layer with (i) residual connections, (ii) layer normalization, and (iii) a positionwise feedforward network.

stepcalculation for the token rowspurpose
1U=\operatorname{MultiHead}(X)combine information from permitted positions
2R=\operatorname{LayerNorm}(X+U)add the residual input and normalize each row
3F=\operatorname{FFN}(R)apply the same feedforward network, which contains two linear layers, to every token row
4Y=\operatorname{LayerNorm}(R+F)add the second residual input and normalize

The positionwise feedforward network contains two linear layers with an activation function between them. It updates each token row inside every Transformer layer.

During training, a final linear layer maps every final token row to logits. During generation, the model uses the logits from the newest row.

Initialization of attention parameters

Illustrative initialization scaled by width
import numpy as np
rng = np.random.default_rng(9)
d_model, d_k, d_v, H = 3, 1, 2, 2
After importing NumPy and initializing its generator with seed 9, set the dimensions for the two worked heads.
def scaled_normal(rows, cols):
  sd = np.sqrt(2.0 / (rows + cols))
  return rng.normal(0.0, sd,
      size=(rows, cols))
After computing a standard deviation from the input and output widths, draw every entry.
Wq = [scaled_normal(d_model, d_k)
      for _ in range(H)]
Wk = [scaled_normal(d_model, d_k)
      for _ in range(H)]
Wv = [scaled_normal(d_model, d_v)
      for _ in range(H)]
Initialize (i) a query projection, (ii) a key projection, and (iii) a value projection for each head.
Wo = scaled_normal(H * d_v, d_model)Initialize the output projection with shape 4\times3.

The scale \sqrt{2/(\mathrm{rows}+\mathrm{cols})} decreases as the matrix dimensions grow.

Training the attention parameters

The matrices in the red ball rolls example use entries chosen by hand so that every result remains exact. A training program can create initial parameter arrays with a pseudorandom generator.

  1. Forward pass: the model uses W_Q,W_K,W_V,W_O to compute a prediction and a loss L.
  2. Backward pass: backpropagation computes the gradient of L with respect to every matrix entry.
  3. Update: gradient descent subtracts the learning rate times the gradient from each entry.
W\leftarrow W-\eta\,\frac{\partial L}{\partial W}

Seed 9 reproduces the same draws when (i) the NumPy version, (ii) the generator type, (iii) the array shapes, and (iv) the generator call order also match.

Why does token order need position information?

dog bites person
person bites dog

Because both sentences contain the same token identities, their positions distinguish the grammatical roles. The additive method in the 2017 Transformer included position information before attention.

x_i=e_i+p_i
symbolmeaningsource of its entries
e_itoken embedding at position ithe model's learned token embedding table
p_iadditive position encodinga fixed sine and cosine encoding or a learned position embedding

Vaswani et al. used sinusoidal position encodings in 2017 and reported similar results with learned position embeddings. In 2018, Peter Shaw, Jakob Uszkoreit, and Ashish Vaswani placed relative position representations inside the attention calculation.

NumPy projections and scores

NumPy codeshape and worked value
import numpy as npImport the NumPy library.
X = np.eye(3)3\times3 identity matrix for the three token rows.
Wq = np.array([[0.], [0.], [1.]])3\times1 query matrix chosen by hand.
Wk = np.array([[0.], [np.log(2.)], [0.]])3\times1 key matrix containing \ln2.
Wv = np.array([[1., 0.],
               [0., 2.],
               [1., 1.]])
3\times2 value matrix.
Q, K, V = X @ Wq, X @ Wk, X @ WvThe @ operator means matrix multiplication. The output shapes are (i) 3\times1 for Q, (ii) 3\times1 for K, and (iii) 3\times2 for V.
scores = Q @ K.T / np.sqrt(Q.shape[1])After .T transposes K, matrix multiplication produces the 3\times3 score matrix.

The code expresses the same row calculations as arithmetic by hand. NumPy performs all query rows together.

NumPy mask, softmax, and output

NumPy codemathematical action
allowed = np.tril(np.ones_like(
    scores, dtype=bool))
Create a lower triangular Boolean mask in which current and earlier columns are true.
scores = np.where(allowed,
    scores, -np.inf)
Keep permitted scores and assign negative infinity to future columns.
shifted = scores - np.max(
    scores, axis=-1, keepdims=True)
Subtracting each row maximum leaves the softmax weights unchanged and makes the largest shifted score 0, which prevents unnecessarily large exponential values.
A = np.exp(shifted)
A = A / A.sum(axis=-1,
    keepdims=True)
Exponentiate each score and divide every row by its own total.
O = A @ VMultiplying A by V computes weighted sums of the value rows and produces an output with shape 3\times2.
A=\begin{pmatrix}1&0&0\\0.5&0.5&0\\0.25&0.5&0.25\end{pmatrix},\qquad O=\begin{pmatrix}1&0\\0.5&1\\0.5&1.25\end{pmatrix}

Training, generation, and cached keys and values

settingavailable token rowsattention computationnext action
trainingthe full training sequence is storedthe causal mask permits all query rows to run togethercompute losses for the predicted positions
generationthe prompt and previously chosen tokens are storedthe program compares the newest query with the permitted keys and uses the corresponding valuesafter the program chooses one token, it appends that token and repeats the calculation

A generation implementation can keep earlier key and value rows in a cache of keys and values. After the program generates a new token, it appends (i) one new key row and (ii) one new value row.

K_{1:t}=\begin{bmatrix}K_{1:t-1}\\k_t\end{bmatrix},\qquad V_{1:t}=\begin{bmatrix}V_{1:t-1}\\v_t\end{bmatrix}

Training exposes parallel query rows. Autoregressive generation remains sequential because the token chosen at step t becomes part of the input at step t+1.

Quadratic growth of the score matrix

With n token positions, each of the n queries scores n keys. One full attention head therefore forms n^2 score entries.

n\text{ rows}\times n\text{ columns}=n^2\text{ scores}

Doubling the sequence length multiplies the score count by four.

Illustrative sequence lengths and exact score counts
tokens nscores per head n^2
416
864
12816,384
1,0241,048,576

These counts describe score entries for one head and one example. Full implementations also store or recompute other intermediate arrays.

What does one attention weight establish?

An entry A_{ij} is the coefficient for value row v_j in the weighted sum that produces output row o_i for (i) one head, (ii) one layer, and (iii) one forward pass.

o_i=\sum_j A_{ij}v_j
  1. Within this weighted sum, a larger coefficient gives v_j a larger direct multiplier.
  2. The final prediction also depends on (i) the value vectors, (ii) later layers, (iii) residual connections, and (iv) the output readout.
  3. After changing a token or an internal vector, recompute the model and measure the resulting change in its prediction.

In 2019, Jain and Wallace produced similar predictions from attention distributions that differed greatly across several language tasks. Wiegreffe and Pinter argued later that year that conclusions depend on the definition of explanation and the test being used.

The path from soft alignment to the Transformer

2014 manuscriptDzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio let a translation decoder softly search source sentence positions, thereby easing the bottleneck created when the entire source was compressed into one encoder vector of fixed length.
2015Minh Thang Luong, Hieu Pham, and Christopher Manning compared global attention over every source word with local attention over a selected region.
2017Ashish Vaswani and seven coauthors introduced the Transformer, whose layers used (i) scaled attention based on dot products, (ii) multiple heads, (iii) position encodings, and (iv) positionwise feedforward networks.
\text{soft alignment}\longrightarrow\text{scaled self attention}\longrightarrow\text{Transformer layers}

Historical precision: the 2017 paper defined (i) encoder self attention, (ii) causal self attention in the decoder, and (iii) attention from the decoder to the encoder. Language models that contain only a decoder use causal self attention to predict the next token.

Tutor practice

1. In a causal window with three tokens, which key columns may the query at position 2 use?

Columns 1 and 2. Column 3 is a future position.

2. Compute the softmax weights for the score row (0,\ln2,0).

Exponentiation gives (1,2,1). Because the total is 4, the softmax weights are (1/4,1/2,1/4).

3. Use those weights with values (i) (1,0), (ii) (0,2), and (iii) (1,1).

(1/4,0)+(0,1)+(1/4,1/4)=(1/2,5/4).

4. If (i) n=5, (ii) d_k=2, and (iii) d_v=3, what are the shapes of the score matrix and the head output?

The score matrix is 5\times5; the output is 5\times3.

For the classroom: require students to show (i) the exponentiation, (ii) the total, and (iii) the weighted sum of the values before revealing each numerical answer.

Part 9 Overview

  1. Three projections: each hidden row produces (i) a query, (ii) a key, and (iii) a value through learned matrices.
  2. Scores and mask: scaled dot products between queries and keys measure compatibility. A causal mask adds -\infty to future scores, so softmax gives those positions zero weight.
  3. Weights and content: rowwise softmax produces weights whose sum is one, and matrix multiplication computes weighted sums of value rows.
  4. Multihead output: separate heads concatenate their outputs before W_O projects each row back to the model width.
  5. Layer structure: a Transformer layer places residual connections and layer normalization around (i) multihead attention and (ii) the positionwise feedforward network.
  6. Score matrix size: direct pairwise scores allow parallel computation of all query rows during training and require an n\times n matrix.

A diffusion model learns to reverse gradual corruption.

Sources and further reading

  • D. Bahdanau, K. Cho, Y. Bengio, 2014 manuscript, “Neural Machine Translation by Jointly Learning to Align and Translate,” arXiv:1409.0473; published at ICLR 2015. The paper identifies using a single encoder vector for the entire source as a bottleneck and introduces learned soft alignment over source positions.
  • M. T. Luong, H. Pham, C. D. Manning, 2015, “Effective Approaches to Attention-based Neural Machine Translation,” EMNLP, pages 1412 to 1421, DOI 10.18653/v1/D15-1166.
  • A. Vaswani et al., 2017, “Attention Is All You Need,” Advances in Neural Information Processing Systems 30. Sections 3 and 4 define (i) scaled dot product attention, (ii) multihead attention, (iii) causal masking, (iv) position encodings, and (v) the computational comparison with recurrence.
  • A. Dosovitskiy et al., 2020 manuscript, “An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale,” arXiv:2010.11929; published at ICLR 2021. The paper maps image patches to token vectors and processes their sequence with a Transformer encoder.
  • P. Shaw, J. Uszkoreit, A. Vaswani, 2018, “Self-Attention with Relative Position Representations,” NAACL HLT, pages 464 to 468, DOI 10.18653/v1/N18-2074.

Interpretation studies and provenance

  • S. Jain, B. C. Wallace, 2019, “Attention is not Explanation,” NAACL HLT, pages 3543 to 3556, DOI 10.18653/v1/N19-1357.
  • S. Wiegreffe, Y. Pinter, 2019, “Attention is not not Explanation,” EMNLP IJCNLP, pages 11 to 20, DOI 10.18653/v1/D19-1002. The paper proposes tests that depend on a stated definition of explanation.
  • The following materials are illustrative: (i) sentence examples, (ii) initialization pattern, (iii) matrices for three tokens, (iv) scaling comparison, (v) code, and (vi) score count tables. Every displayed numerical result follows from the shown values.
  • Owner: Mehmet Kerem Turkcan; Associate Research Scientist; Center for Smart Streetscapes, Columbia University; New York, USA; keremturkcan.com; mkt2126@columbia.edu.
Intro to AI, Part 9: AttentionM. K. Turkcan, Columbia University