You are reading immutable version 39. The current guide may be newer.

Study the paper

Attention Is All You Need

Lessons, visuals, quizzes, flashcards, and resources—organized in teaching order.

All activities

Scaled Dot-Product Attention

Scaled Dot-Product Attention

Scaled Dot-Product Attention

Source equation

Attention(Q,K,V)=softmax(QKTdk)V\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V

The Scaled Dot-Product Attention mechanism computes attention weights over a set of values based on queries and keys. The scaling factor dk\sqrt{d_k} is introduced to prevent the dot products from growing extremely large in magnitude for high-dimensional keys, which would otherwise push the softmax function into regions with extremely small gradients.

Sources

S3.E1

Attention​(Q,K,V)=softmax​(Q​KTdk)​V\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V (1)
\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V

Attention(Q,K,V)=softmax(QKTdk)V\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V

Sources

S3.E1

Attention​(Q,K,V)=softmax​(Q​KTdk)​V\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V (1)
\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V
Implementation detail

Illustrative Toy Calculation

Let us compute the attention output for a single query vector with dk=4d_k = 4:

  • Query vector Q=[1.0,0.0,1.0,0.0]Q = [1.0, 0.0, 1.0, 0.0] (shape: 1×41 \times 4)
  • Key matrix K=[1.00.01.00.00.01.00.01.0]K = \begin{bmatrix} 1.0 & 0.0 & 1.0 & 0.0 \\ 0.0 & 1.0 & 0.0 & 1.0 \end{bmatrix} (shape: 2×42 \times 4)
  • Value matrix V=[10.020.0]V = \begin{bmatrix} 10.0 \\ 20.0 \end{bmatrix} (shape: 2×12 \times 1)

Step 1: Compute the dot product QKTQ K^T QKT=[1.0⋅1.0+0.0⋅0.0+1.0⋅1.0+0.0⋅0.0,  1.0⋅0.0+0.0⋅1.0+1.0⋅0.0+0.0⋅1.0]=[2.0,0.0]Q K^T = [1.0 \cdot 1.0 + 0.0 \cdot 0.0 + 1.0 \cdot 1.0 + 0.0 \cdot 0.0,\; 1.0 \cdot 0.0 + 0.0 \cdot 1.0 + 1.0 \cdot 0.0 + 0.0 \cdot 1.0] = [2.0, 0.0]

Step 2: Scale by dk=4=2.0\sqrt{d_k} = \sqrt{4} = 2.0 QKTdk=[2.02.0,0.02.0]=[1.0,0.0]\frac{Q K^T}{\sqrt{d_k}} = \left[ \frac{2.0}{2.0}, \frac{0.0}{2.0} \right] = [1.0, 0.0]

Step 3: Apply Softmax softmax([1.0,0.0])=[e1e1+e0,e0e1+e0]≈[0.731,0.269]\mathrm{softmax}([1.0, 0.0]) = \left[ \frac{e^1}{e^1 + e^0}, \frac{e^0}{e^1 + e^0} \right] \approx [0.731, 0.269]

Step 4: Multiply by VV Attention(Q,K,V)=0.731⋅[10.0]+0.269⋅[20.0]=[7.31+5.38]=[12.69]\mathrm{Attention}(Q,K,V) = 0.731 \cdot [10.0] + 0.269 \cdot [20.0] = [7.31 + 5.38] = [12.69]

Sources

S3.E1

Attention​(Q,K,V)=softmax​(Q​KTdk)​V\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V (1)
\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{QK^{T}}{\sqrt{d_{k}}})V
QQ
Query matrix of shape (batch_size, sequence_length_q, d_k) · [B, S_Q, d_k]
KK
Key matrix of shape (batch_size, sequence_length_k, d_k) · [B, S_K, d_k]
VV
Value matrix of shape (batch_size, sequence_length_k, d_v) · [B, S_K, d_v]
dkd_k
Dimensionality of the keys and queries · Scalar
Attention(Q,K,V)Attention(Q,K,V)
The resulting attention representation of shape (batch_size, sequence_length_q, d_v) · [B, S_Q, d_v]