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.

When it falls down

• Readable does not mean fast by default. A naive einsum implementation may materialise large intermediate tensors or choose a suboptimal contraction order for multi-tensor expressions; production code often falls back to hand-fused kernels (like FlashAttention; see flash-attention) specifically because a generic einsum evaluator will not discover that fusion on its own.
• It hides memory cost as easily as it hides compute. The attention score einsum above materialises a full (batch, heads, seq_q, seq_k) tensor, quadratic in sequence length, just as plainly as writing it out by hand would; the compact notation does not make that cost disappear, only easier to overlook.
• Label collisions are a common bug source. Reusing the same letter for two axes that are not actually the same size, or forgetting to include an axis in the output that should have been preserved, produces either a shape error or, worse, a silently wrong contraction.
• Not every tensor operation is expressible as a pure contraction. Non-linearities, normalisation, indexing, and gather operations are not einsums; einsum covers exactly the multiply-and-sum family of operations, which happens to be most but not all of a transformer's compute.

Further reading

• NumPy documentation, einsum - the reference specification of the notation, shared essentially unchanged across NumPy, PyTorch, JAX, and TensorFlow.
• Vaswani et al., 2017, Attention Is All You Need, arXiv:1706.03762 - the QK^T and weighted-value-sum operations that the einsum examples above implement directly.