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 maps a set of query vectors, key vectors, and value vectors to an output. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.

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
Deep dive

Why scale by 1/dk1/\sqrt{d_k}?

For large values of dkd_k, the dot products grow large in magnitude, pushing the softmax function into regions with extremely small gradients. To counteract this effect, the dot products are scaled by 1dk\frac{1}{\sqrt{d_k}}.

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

Let's walk through a deterministic toy example with a single query, two keys, and two values.

Let:

  • Q=(1.02.0)Q = \begin{pmatrix} 1.0 & 2.0 \end{pmatrix} (shape: 1×21 \times 2, so dk=2d_k = 2)
  • K=(1.00.00.01.0)K = \begin{pmatrix} 1.0 & 0.0 \\ 0.0 & 1.0 \end{pmatrix} (shape: 2×22 \times 2)
  • V=(10.020.0)V = \begin{pmatrix} 10.0 \\ 20.0 \end{pmatrix} (shape: 2×12 \times 1)

Step 1: Compute the dot products QKTQ K^T: QKT=(1.02.0)(1.00.00.01.0)=(1.02.0)Q K^T = \begin{pmatrix} 1.0 & 2.0 \end{pmatrix} \begin{pmatrix} 1.0 & 0.0 \\ 0.0 & 1.0 \end{pmatrix} = \begin{pmatrix} 1.0 & 2.0 \end{pmatrix}

Step 2: Scale by dk=2≈1.4142\sqrt{d_k} = \sqrt{2} \approx 1.4142: QKTdk=(1.01.41422.01.4142)≈(0.70711.4142)\frac{Q K^T}{\sqrt{d_k}} = \begin{pmatrix} \frac{1.0}{1.4142} & \frac{2.0}{1.4142} \end{pmatrix} \approx \begin{pmatrix} 0.7071 & 1.4142 \end{pmatrix}

Step 3: Apply softmax: softmax((0.70711.4142))≈(0.33030.6697)\text{softmax}(\begin{pmatrix} 0.7071 & 1.4142 \end{pmatrix}) \approx \begin{pmatrix} 0.3303 & 0.6697 \end{pmatrix}

Step 4: Multiply by VV: Attention(Q,K,V)≈(0.33030.6697)(10.020.0)=0.3303×10.0+0.6697×20.0=16.697\text{Attention}(Q,K,V) \approx \begin{pmatrix} 0.3303 & 0.6697 \end{pmatrix} \begin{pmatrix} 10.0 \\ 20.0 \end{pmatrix} = 0.3303 \times 10.0 + 0.6697 \times 20.0 = 16.697

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 containing packed query vectors · matrix of shape (d_seq_q, d_k)
KK
Key matrix containing packed key vectors · matrix of shape (d_seq_k, d_k)
VV
Value matrix containing packed value vectors · matrix of shape (d_seq_k, d_v)
dkd_k
Dimensionality of the key vectors · scalar
Attention(Q,K,V)Attention(Q,K,V)
The resulting attention representation · matrix of shape (d_seq_q, d_v)