Study the paper

Attention Is All You Need

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

Scaled Dot-Product Attention

Source equation

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

Scaled Dot-Product Attention computes attention weights by taking the dot product of queries with keys, scaling them by the square root of the key dimension, applying a softmax function, and using the resulting weights to compute a weighted sum of values.

Sources

S3.E1

Attention(Q,K,V)=softmax(QKTdk)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

equation

Attention(Q,K,V)=softmax(QKTdk)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

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

Sources

S3.E1

Attention(Q,K,V)=softmax(QKTdk)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

equation

Attention(Q,K,V)=softmax(QKTdk)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

The scaling factor $\frac{1}{\sqrt{d_k}}$ is crucial. For large values of $d_k$ , the dot products grow large in magnitude, pushing the softmax function into regions with extremely small gradients. Scaling by $\sqrt{d_k}$ counteracts this effect.

Sources

S3.E1

Attention(Q,K,V)=softmax(QKTdk)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

equation

Attention(Q,K,V)=softmax(QKTdk)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 perform a deterministic toy calculation. Suppose we have:

$Q = \begin{pmatrix} 1.0 & 0.0 \end{pmatrix}$ (shape: $1 \times 2$ )
$K = \begin{pmatrix} 1.0 & 0.0 \\ 0.0 & 1.0 \end{pmatrix}$ (shape: $2 \times 2$ )
$V = \begin{pmatrix} 2.0 & 3.0 \\ 4.0 & 5.0 \end{pmatrix}$ (shape: $2 \times 2$ )
$d_k = 2$ , so $\sqrt{d_k} = \sqrt{2} \approx 1.41421356$

Compute $Q K^T$ : $Q K^T = \begin{pmatrix} 1.0 & 0.0 \end{pmatrix} \begin{pmatrix} 1.0 & 0.0 \\ 0.0 & 1.0 \end{pmatrix} = \begin{pmatrix} 1.0 & 0.0 \end{pmatrix}$
Scale by $\sqrt{d_k}$ : $\frac{Q K^T}{\sqrt{d_k}} = \begin{pmatrix} \frac{1.0}{\sqrt{2}} & \frac{0.0}{\sqrt{2}} \end{pmatrix} \approx \begin{pmatrix} 0.70710678 & 0.0 \end{pmatrix}$
Apply Softmax: $\text{softmax}(\begin{pmatrix} 0.70710678 & 0.0 \end{pmatrix}) = \begin{pmatrix} \frac{e^{0.70710678}}{e^{0.70710678} + e^0} & \frac{e^0}{e^{0.70710678} + e^0} \end{pmatrix} \approx \begin{pmatrix} 0.66974149 & 0.33025851 \end{pmatrix}$
Multiply by $V$ : $\text{Attention}(Q,K,V) = \begin{pmatrix} 0.66974149 & 0.33025851 \end{pmatrix} \begin{pmatrix} 2.0 & 3.0 \\ 4.0 & 5.0 \end{pmatrix} \approx \begin{pmatrix} 2.66051702 & 3.66051702 \end{pmatrix}$

Sources

S3.E1

Attention(Q,K,V)=softmax(QKTdk)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

equation

Attention(Q,K,V)=softmax(QKTdk)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

$Q$: Query matrix · (n_q, d_k)
$K$: Key matrix · (n_k, d_k)
$V$: Value matrix · (n_k, d_v)
$d_k$: Dimensionality of keys · scalar
$Attention$: Attention output matrix · (n_q, d_v)