You are reading immutable version 40. 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 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)