Reductions and axes — Linear algebra for LLMs

Summing over an axis removes it

sum(dim=0)

sum(dim=1)

dim=0 collapses the rows: one sum per column. dim=1 collapses the columns: one sum per row. The axis you name is the one that disappears.

Write

Reducing over dim=k removes axis

k

(3, 4) \to (4,)

for dim 0,

(3,)

for dim 1.

Mean, max, and `dim=-1`

On $X$ of shape $(3, 4)$ :

X.mean(dim=-1) → (3,): each row’s mean.
X.max(dim=-1) → two (3,) tensors: each row’s max, and where it is.
X.argmax(dim=-1) → (3,): the index of each row’s max.
X.sum() → (): everything added up.

dim=-1 always means the last axis, whatever the rank, which is why transformer code uses it everywhere: the last axis is the features (or the keys, for attention scores).

keepdim keeps a size-1 axis

X.mean(dim=-1) X.mean(dim=-1, keepdim=True)

With keepdim=True the reduced axis stays as size 1, so the result broadcasts straight back against $X$ : X - X.mean(dim=-1, keepdim=True) centers every row. Without it you get the $(T,)$ bug from the broadcasting deck.

Write

keepdim=True:

(3, 4) \to (3, 1)

, ready to broadcast back.

Center each row

−

Row means 0.5, 1 and 1, as a $(3, 1)$ column, subtracted from every entry of their row. Each row of the result sums to 0. This is the first half of LayerNorm.

Which axis, in a transformer

In a transformer

Softmax of attention scores $(B, h, T, T)$ : over the keys, dim=-1.
LayerNorm or RMSNorm of $(B, T, d)$ : over one token’s features, dim=-1.
The loss over tokens $(B, T)$ : the mean over all of them.

LayerNorm normalizes each token on its own, over its $d$ features. It never looks across the batch; that’s BatchNorm, which transformers don’t use.

✏️ Your turn

Reduce these

x

Write the shapes of x.mean(-1), x.mean(-1, keepdim=True) and x.sum(dim=(0, 1)).

Work it out on paper, then go to the next card to check.

Answer

Expression	Shape
`x.mean(-1)`	`(2, 3)`
`x.mean(-1, keepdim=True)`	`(2, 3, 1)`
`x.sum(dim=(0, 1))`	`(4,)`

Reducing over two axes removes both: one total per feature, summed over every token of every sequence.

📌 Recap

Reductions and axes in one card

Reducing over an axis removes it; keepdim=True keeps it as size 1.
dim=-1 is the last axis: features for norms, keys for softmax.
Center rows: x - x.mean(-1, keepdim=True).
LayerNorm works per token over $d$ , never across the batch.

Write

The four lines above.