Multi-Head Attention with Einsum - Einstein Summation (Einsum)

Multi-Head Attention with Einsum

What you'll build

A function multi_head_attention_einsum(Q, K, V, num_heads) that splits d_model into heads, runs scaled dot-product attention in every head at once with einsum, and merges the heads back.

Background

Multi-head attention lets each head attend to different positions using its own slice of the features. The usual implementation reshapes to (batch, seq, heads, d_k) and then transposes to (batch, heads, seq, d_k) before two batched matmuls. With einsum you can skip the transposes: put the head letter h where it already is and let the subscripts reorder axes for you. This problem covers only the attention core; the learned projections W_Q, W_K, W_V, W_O of a full layer are left out (see Multi-Head Attention in the Transformers section).

Algorithm

Let H = num_heads and d_k = d_model / H.

1. Split heads: reshape Q to (batch, seq_q, H, d_k) and K, V to (batch, seq_k, H, d_k). Head h owns feature channels h*d_k to (h+1)*d_k - 1.
2. Scores per head:

   scores[b, h, q, k] = sum_d Q[b, q, h, d] * K[b, k, h, d] / sqrt(d_k)     einsum 'bqhd,bkhd->bhqk'
   

3. Softmax over the key axis (last axis) of scores, giving A of shape (batch, H, seq_q, seq_k).
4. Weighted values:

   out[b, q, h, d] = sum_k A[b, h, q, k] * V[b, k, h, d]                   einsum 'bhqk,bkhd->bqhd'
   

5. Merge heads: reshape out to (batch, seq_q, d_model).

Inputs

• Q: queries, float array of shape (batch, seq_q, d_model)
• K: keys, float array of shape (batch, seq_k, d_model)
• V: values, float array of shape (batch, seq_k, d_model)
• num_heads: positive integer that divides d_model

Expected Return Format

A float array of shape (batch, seq_q, d_model). Output channels h*d_k .. (h+1)*d_k - 1 come from head h.

Edge cases and conventions

• Scale by sqrt(d_k) (the per-head size), not sqrt(d_model).
• Use a numerically stable softmax (subtract the row max before exp).
• With num_heads = 1 the result is plain single-head attention.

Constraints

Use reshape and np.einsum; no explicit transpose, no loops over heads.

Follow-up

• Count the FLOPs: does splitting into H heads change the total cost compared with one head of size d_model?
• How would you add a causal mask, and in which step?

Input: X = np.array([[[1., 0.], [0., 1.]]])    # (batch=1, seq=2, d_model=2)
np.round(multi_head_attention_einsum(X, X, X, 2), 4)
Output: array([[[0.7311, 0.5   ],
        [0.5   , 0.7311]]])

Input: np.round(multi_head_attention_einsum(X, X, X, 1), 4)
Output: array([[[0.6698, 0.3302],
        [0.3302, 0.6698]]])

import numpy as np

def multi_head_attention_einsum(Q: np.ndarray, K: np.ndarray, V: np.ndarray,
                                 num_heads: int) -> np.ndarray:
    """
    Multi-head scaled dot-product attention using einsum (no projections).

    Args:
        Q: Queries (batch, seq_q, d_model)
        K: Keys (batch, seq_k, d_model)
        V: Values (batch, seq_k, d_model)
        num_heads: Number of heads; divides d_model

    Returns:
        output: (batch, seq_q, d_model); channels h*d_k:(h+1)*d_k come from head h
    """
    batch, seq_q, d_model = Q.shape
    seq_k = K.shape[1]
    d_k = d_model // num_heads

    # Your code here
    pass