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: correct the supplied diagram

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 image lists \((0.26,0.41,0.33,0)\), but that vector is not the softmax of the displayed scores. This course uses the calculated values above.

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). \]

Move the query and watch the distribution

\(i=\)3
KeyRaw scoreAfter maskSoftmax weightValue

Weighted output:

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?

Retrieval check

What must the weights in one attention row sum to?

Practice before moving on

  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.