Einsum and Tensor Contractions

Read enough transformer implementations and eventually you hit a line like torch.einsum('bhqd,bhkd->bhqk', q, k) and have to stop and think. Einstein summation notation (einsum) is not a new mathematical idea, it is a compact way of writing exactly which axes of two tensors get multiplied elementwise, which get summed away, and which survive into the output. Every matmul, batched matmul, and the core of attention itself is a single einsum call, and learning to read the notation pays off precisely because it removes the ambiguity that .reshape() and .transpose() chains otherwise hide.

The notation, by example

Ordinary matrix multiplication C = A @ B, with A shape (m, k) and B shape (k, n), is written einsum('mk,kn->mn', A, B). Read it left to right: A has axes labelled m and k, B has axes labelled k and n, the arrow shows the output has axes m and n. The rule that makes this work is simple and mechanical:

• A label that appears in both inputs but not in the output is summed over (contracted).
• A label that appears in only one input, or in the output, is kept as-is.

k appears in both inputs and not in the output, so it is summed over, exactly reproducing sum_k A[m,k] * B[k,n], the ordinary matmul formula (see matmul-the-core-operation).

Batching and attention as one pattern

Batched matmul is the same rule with a shared batch label that is not contracted because it also appears in the output: einsum('bmk,bkn->bmn', A, B). The b axis is preserved (not summed) because it appears on both sides of the arrow; only k disappears.

Attention's score computation is exactly this pattern extended with a head axis:

scores = einsum('bhqd,bhkd->bhqk', Q, K)

Q has shape (batch, heads, seq_q, d_head), K has shape (batch, heads, seq_k, d_head). The label d (the head dimension) appears in both inputs and not the output, so it is contracted, exactly the dot product between each query and each key. b and h are batch-like labels preserved on both sides. q and k each appear in only one input, so both survive into the output shape (batch, heads, seq_q, seq_k), the full attention score matrix. Applying the weighted sum over values is another einsum, einsum('bhqk,bhkd->bhqd', attn_weights, V), contracting the key-position axis k this time instead of d.

Why this notation, and not just .matmul()

torch.matmul and @ only handle the "batch of matrix multiplications" case cleanly; anything with more axes, or with axes that need to be summed in an order that does not match the standard matmul shape, usually requires a chain of transpose, reshape, and matmul calls whose intent is not visible from reading the code. Einsum expresses the same computation declaratively: name the axes, name what survives, and let the implementation figure out the most efficient contraction order. This also makes it easier to reason about cost: multiplying out the label counts on each side gives the FLOP count directly, and multi-tensor einsums (contracting three or more tensors at once) can sometimes be computed in a different order that is asymptotically cheaper, a genuinely non-trivial optimisation problem in its own right.