Self-Attention in Matrix Form

Matrix notation computes all query positions together, with query rows and readable key-value columns kept explicit. Apply the causal mask before stable row-wise softmax, then distinguish scores, weights, values, and vector contributions.

Matrix notation computes every query position at once. It is compact, but only if we keep the two position axes distinct: rows are queries, and columns are keys whose values may be read.

For one sequence,

\begin{aligned} Q&=XW_Q, & Q&:T\times d_k,\\ K&=XW_K, & K&:T\times d_k,\\ V&=XW_V, & V&:T\times d_v. \end{aligned}

The score matrix is

S=\frac{QK^\top}{\sqrt{d_k}},\qquad S:T\times T.

Entry $S_{ij}$ compares query $i$ with key $j$ . The transpose is what turns $K$ from $(T,d_k)$ into $(d_k,T)$ , so the inner $d_k$ axes contract.

Mask Before Normalizing

For causal self-attention, define

M_{ij}=\begin{cases} 0 & j\le i,\\ -\infty & j>i. \end{cases}

Then normalize across the final, key-position axis:

A=\operatorname{softmax}_{last}(S+M).

Every row of $A$ sums to 1 over its allowed columns. Future columns receive exactly zero weight. Masking after softmax and simply setting entries to zero would reduce the row sum; the remaining entries would no longer be the same normalized distribution.

In code, a finite minimum value is often used instead of literal $-\infty$ . The required behavior is the same: forbidden entries must contribute zero after softmax, and every allowed row must remain finite and normalized.

Extend the Four-Token Trace

The Language Modeling trace used

X=\begin{bmatrix} 1&0\\ 0&1\\ 1&1\\ -1&1 \end{bmatrix},\qquad W_Q=W_K=I,

and

W_V=\begin{bmatrix}1&1\\1&-1\end{bmatrix}.

For the final position, the scaled scores were approximately

[-0.707,\;0.707,\;0,\;1.414],

the weights were

[0.065,\;0.266,\;0.131,\;0.539],

and the weighted value reading was

h_4\approx[0.592,-1.279].

This chapter treats that calculation as verified prior work. The matrix form performs the same steps for all four rows and produces $H=AV\in\mathbb R^{4\times2}$ .

Keep Four Quantities Separate

For one pair of positions $(i,j)$ :

$S_{ij}$ is an unnormalized compatibility score.
$A_{ij}$ is the row-normalized attention weight.
$v_j$ is the value vector at source position $j$ .
$A_{ij}v_j$ is that source position's vector contribution to $h_i$ .

A large score need not imply a weight near 1 if other scores are similar. A large weight need not imply a large contribution if the value is small. A large contribution in $H$ may later be changed or cancelled by the output projection and other parts of the network.

Q1. Identify the softmax axis

A score tensor for one batch has shape $(5,5)$ . Row 3 belongs to one query, and its five columns refer to key positions. Along which axis must softmax operate?

Choose one

Across the five columns within each query row

Down the five query rows within each key column

Across all 25 scores at once

Select one choice, then check.

Hint

Each query row must sum to 1 over the positions it can read.

Solution

Softmax operates across columns, the final axis. This normalizes each query's scores into a distribution over key-value positions.

Audit the Equations in Dependency Order

Check shapes and values in the same order as the computation: projections, scores, mask, softmax row sums, forbidden zeroes, then weighted values. Looking only at the final matrix makes an earlier axis or mask error much harder to locate.