Batch Operations with Einsum - Einstein Summation (Einsum)

Batch Operations with Einsum

What you'll build

A function einsum_batch_ops(Q, K, V) that computes batched scaled dot-product attention with two einsums: raw scores, softmax weights and the weighted sum of values.

Background

Batched operations add a letter (here b) that appears in every operand and in the output, so it is carried through untouched: every batch element gets its own independent product. Attention is the canonical example. Each query is compared with every key, the scores are turned into weights with a softmax, and the output is the weighted average of the values.

Formula

With d the feature dimension of Q and K (d = Q.shape[-1]):

1. Scores (unscaled):

   scores[b, q, k] = sum_d Q[b, q, d] * K[b, k, d]          einsum 'bqd,bkd->bqk'
   

2. Weights: softmax over the keys (last axis) of the scaled scores:

   weights[b, q, :] = softmax(scores[b, q, :] / sqrt(d))
   softmax(x)_k = exp(x_k - max(x)) / sum_j exp(x_j - max(x))
   

3. Output:

   output[b, q, v] = sum_k weights[b, q, k] * V[b, k, v]    einsum 'bqk,bkv->bqv'
   

Inputs

• Q: queries, float array of shape (batch, seq_q, d)
• K: keys, float array of shape (batch, seq_k, d)
• V: values, float array of shape (batch, seq_k, d_v)

Expected Return Format

Return a dictionary with these keys:

• 'scores': raw (unscaled) scores Q @ K^T per batch, shape (batch, seq_q, seq_k)
• 'attention_weights': softmax(scores / sqrt(d)) over the last axis, shape (batch, seq_q, seq_k); each row sums to 1
• 'output': weighted sum of values, shape (batch, seq_q, d_v)

Edge cases and conventions

• Subtract the row maximum before exp so large scores don't overflow; this doesn't change the softmax.
• 'scores' is before scaling; only the weights use / sqrt(d).

Constraints

Use np.einsum for both products (no explicit transposes or @) and NumPy for the softmax.

Follow-up

• What is the time and memory cost in terms of batch, seq_q, seq_k, d? Which term dominates for long sequences?
• Why divide by sqrt(d) at all?

Input: Q = np.array([[[1., 0.], [0., 1.]]])            # (1, 2, 2)
K = np.array([[[1., 1.], [2., 0.], [0., 3.]]])  # (1, 3, 2)
V = np.array([[[1., 0.], [0., 1.], [1., 1.]]])  # (1, 3, 2)
r = einsum_batch_ops(Q, K, V)
r["scores"]
Output: array([[[1., 2., 0.],
        [1., 0., 3.]]])

Input: np.round(r["attention_weights"], 4)
Output: array([[[0.284 , 0.576 , 0.14  ],
        [0.1784, 0.0879, 0.7337]]])

Input: np.round(r["output"], 4)
Output: array([[[0.424 , 0.716 ],
        [0.9121, 0.8216]]])

import numpy as np

def einsum_batch_ops(Q: np.ndarray, K: np.ndarray, V: np.ndarray) -> dict:
    """
    Batched scaled dot-product attention with np.einsum.

    Args:
        Q: Queries of shape (batch, seq_q, d)
        K: Keys of shape (batch, seq_k, d)
        V: Values of shape (batch, seq_k, d_v)

    Returns:
        dict with keys:
            'scores':            Q @ K^T per batch (unscaled), shape (batch, seq_q, seq_k)
            'attention_weights': softmax(scores / sqrt(d)) over keys, shape (batch, seq_q, seq_k)
            'output':            weights @ V per batch, shape (batch, seq_q, d_v)
    """
    # Your code here
    pass