A Short Course for Modern Large Language ModelsGlossary

Lesson 19 of 23 · Transformer block

Turn comparisons into a blend.

Your win: calculate softmax probabilities from masked scores and use them to form a weighted sum of value vectors.

15 minutesNeeds: exponentialsOutcome: calculate one attention output
Mission link: \(Q'\) and \(K'\) decide the weights; the weights select and blend information carried by \(V\).

Softmax normalizes an entire score row

\[ P_{ij}=\operatorname{softmax}(A_{i,:})_j =\frac{e^{A_{ij}}}{\sum_{r=1}^{T}e^{A_{ir}}}, \qquad \sum_jP_{ij}=1. \]

Numbers substituted: a rounded softmax check

For masked scores \((0.7,1.1,0.2,-\infty)\):

\[ (e^{0.7},e^{1.1},e^{0.2},e^{-\infty}) \approx(2.014,3.004,1.221,0), \] \[ Z=2.014+3.004+1.221=6.239, \] \[ \mathbf p\approx \left(\frac{2.014}{6.239},\frac{3.004}{6.239},\frac{1.221}{6.239},0\right) =(0.323,0.482,0.196,0). \]

The calculated values above are the reference distribution for this lesson. The displayed weights are rounded to three decimals, so use the unrounded exponentials when checking that they sum to exactly one.

The probabilities blend value vectors

\[O=PV,\qquad \mathbf o_i=\sum_jP_{ij}\mathbf v_j.\]

Let \(\mathbf v_1=(1,0)\), \(\mathbf v_2=(0,2)\), \(\mathbf v_3=(1,1)\), and \(\mathbf v_4=(-1,1)\). Then for query 3:

\[ \mathbf o_3 =0.323(1,0)+0.482(0,2)+0.196(1,1)+0(-1,1) \approx(0.519,1.159). \]
Keep the roles separate: query-key scores answer “how much?”; values answer “what content is carried into the output?” RoPE changes the former by rotating Q and K, while V stays unrotated.

Code checkpoint · probabilities blend values

Softmax acts across the key axis. The following matrix multiplication replaces that key axis with the value feature axis.

Show softmax and weighted values
weights = F.softmax(scores.float(), dim=-1).to(dtype=q.dtype)
# weights: [B, H, T, T], each row sums to 1

context = weights @ v
# v:       [B, H, T, Dh]
# context: [B, H, T, Dh]

Trace it: Why does the final dimension of context equal Dh rather than T?

Engineering bridge · dtype-safe attention output

MiniMind keeps the probability calculation in float32 for numerical stability, casts the weights back to the query dtype, optionally drops attention links during training, then projects and drops the residual update.

Show the explicit attention path
scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
weights = self.attn_dropout(
    F.softmax(scores.float(), dim=-1).type_as(xq)
)
context = weights @ xv                       # [B, H, Tq, Dh]
context = context.transpose(1, 2).reshape(bsz, seq_len, -1)
output = self.resid_dropout(self.o_proj(context))  # [B, Tq, D]

Runtime check: dropout probabilities are zero in evaluation mode, so the same shape contract holds for training and inference.

Retrieval check

Answer once before reopening the code; use the feedback to correct your explanation.

What must the weights in one attention row sum to?

STAGE 4 · PROVE IT

Practice studio · recall → build → verify

  1. Compute \(\operatorname{softmax}(0,0)\).
  2. Compute \(\operatorname{softmax}(\ln2,0)\).
  3. Why does adding the same constant to every finite score leave softmax unchanged?
  4. Use weights \((0.25,0.75)\) and values \((2,0)\), \((0,4)\) to calculate the output.
  5. Using \((0.323,0.482,0.196,0)\), verify the weights sum to approximately one. Explain the rounding discrepancy.
  6. If all attention weight is on \(\mathbf v_2\), what is the output?
Check solutions
  1. \((0.5,0.5)\).
  2. \((2/3,1/3)\), because \(e^{\ln2}=2\).
  3. The common factor \(e^c\) cancels between numerator and denominator.
  4. \(0.25(2,0)+0.75(0,4)=(0.5,3)\).
  5. The displayed sum is \(1.001\); the exact unrounded values sum to \(1\).
  6. Exactly \(\mathbf v_2\).

Primary source: Vaswani et al., Section 3.2.1.