Einstein Summation (Einsum)

Einstein Summation (Einsum)

Why it matters

Einsum describes a tensor operation by naming its axes: 'bqd,bkd->bqk' says "for every batch b, query q and key k, sum over d", which is exactly attention's score computation. Research code for transformers, mixture-of-experts layers and bilinear models uses einsum because the subscripts document the shapes. In interviews it's a fast way to show you understand which axes are kept, broadcast or contracted, and the same strings work in np.einsum and torch.einsum.

Core ideas

The rules

In np.einsum('<inputs>-><output>', *operands), each operand gets one letter per axis.

1. A letter that appears in the output is kept, in the order written after ->.
2. A letter that appears in the inputs but not in the output is summed over (contracted).
3. A letter shared by several operands pairs those axes up; their sizes must match.
4. A letter repeated within one operand ('ii') walks its diagonal.

Mentally, einsum is a nested loop over every distinct letter that multiplies the matching entries and accumulates them into the output slot:

# C = np.einsum('ik,kj->ij', A, B)
for i in range(m):
    for j in range(n):
        for k in range(K):             # k is missing from the output: summed
            C[i, j] += A[i, k] * B[k, j]

The cost of a two-operand einsum is the product of the sizes of all distinct letters: m * K * n here.

Common patterns

Subscripts	Meaning	Equivalent
'ij->'	sum of all entries	A.sum()
'ij->i' / 'ij->j'	row sums / column sums	A.sum(axis=1) / A.sum(axis=0)
'ij->ji'	transpose	A.T
'ii->i' / 'ii->'	diagonal / trace	np.diag(A) / np.trace(A)
'i,i->'	dot product	u @ v
'i,j->ij'	outer product	np.outer(u, v)
'ij,j->i'	matrix-vector	A @ v
'ik,kj->ij'	matrix multiply	A @ B
'ij,ij->ij' / 'ij,ij->'	Hadamard product / Frobenius inner product	A * B / (A * B).sum()
'bij,bjk->bik'	batched matmul	A @ B on 3D arrays
'bqd,bkd->bqk'	attention scores, (batch, seq_q, d) x (batch, seq_k, d)	Q @ K.transpose(0, 2, 1)
'bhqk,bhkd->bhqd'	multi-head weighted sum of values	A @ V on 4D arrays

Worked example

A = np.array([[1, 2],
              [3, 4]])
B = np.array([[5, 6],
              [7, 8]])
np.einsum('ik,kj->ij', A, B)   # [[19, 22], [43, 50]]   C[0, 0] = 1*5 + 2*7 = 19
np.einsum('ii->', A)           # 5                      1 + 4
np.einsum('ij,ij->', A, A)     # 30                     1 + 4 + 9 + 16

Q = np.array([[[1., 0.], [0., 1.]]])               # (batch=1, seq_q=2, d=2)
K = np.array([[[1., 1.], [2., 0.], [0., 3.]]])     # (batch=1, seq_k=3, d=2)
np.einsum('bqd,bkd->bqk', Q, K)                    # (1, 2, 3)
# [[[1., 2., 0.],      query 0 picks feature 0 of each key
#   [1., 0., 3.]]]     query 1 picks feature 1 of each key

Divide those scores by sqrt(d) = 1.414 and softmax each row, and query 0's attention weights over the three keys are [0.284, 0.576, 0.140].

Common pitfalls

• A dropped letter is summed, silently: 'bij,bjk->ik' sums over the batch as well as j. Check that every axis you want to keep appears after ->.
• Output order matters: 'ij->ji' transposes; 'bqk' versus 'bkq' gives scores of a different layout.
• Implicit mode surprises: without ->, the output letters are sorted alphabetically, so np.einsum('ji', A) returns A.T. Always write the arrow.
• Repeated letters need equal sizes: 'ii->i' on a (2, 3) array raises an error; check for squareness first.
• Integer inputs stay integer: einsum of int arrays returns ints. Convert to float before dividing by sqrt(d) or taking a softmax.
• Many operands can be slow: np.einsum doesn't optimise the contraction order by default; pass optimize=True for three or more operands.

Interview angle

• "Write attention with einsum" is common: name the axes, write 'bqd,bkd->bqk' and 'bqk,bkd->bqd', and state the cost, O(batch * seq_q * seq_k * d).
• Expect to translate between einsum and @/transpose/sum in both directions, and to explain which letters are contracted.
• Follow-ups: multi-head attention without explicit transposes ('bqhd,bkhd->bhqk'), a bilinear form x^T W y as 'bi,ij,bj->b', and when einsum is slower than a hand-written @.

Problems in this section

1. Einsum Fundamentals: sums, transpose and diagonal of one matrix.
2. Matrix Operations with Einsum: matmul, matrix-vector, outer, Hadamard, Frobenius and trace.
3. Batch Operations with Einsum: batched attention scores, softmax weights and the weighted sum of values.
4. Multi-Head Attention with Einsum: split heads, attend per head and merge, without transposes.
5. Einsum vs Traditional Operations: check each einsum against its NumPy equivalent.

The attention problems lead straight into Scaled Dot-Product Attention and Multi-Head Attention in the Transformers section.