Attention Backward Pass (dQ, dK, dV) - Transformer Architecture

Attention Backward Pass

What you'll build

Single-head scaled dot-product attention and its backward pass by hand: given the upstream gradient dOut = dL/dOut, return dQ, dK and dV. This is what autograd (and FlashAttention's backward kernel) computes for every attention layer.

Background

Interviewers ask for this to check that you can backpropagate through a matrix product and a softmax without a framework. The only non-obvious step is the softmax: its Jacobian is diag(p) - p p^T, but you never build it; you apply it to a vector in O(m) per row.

Forward

Shapes: Q (n, d), K (m, d), V (m, d_v).

S   = Q @ K.T / sqrt(d)                   # (n, m), d = Q.shape[-1]
S   = S with -inf where j > i             # only if causal=True (then n == m)
P   = softmax(S, axis=-1)                 # each row sums to 1
Out = P @ V                               # (n, d_v)

Backward

Given dOut (n, d_v), work back through each line:

dV = P.T @ dOut                           # (m, d_v)
dP = dOut @ V.T                           # (n, m)
dS = P * (dP - rowsum(dP * P))            # softmax Jacobian-vector product, row by row
dQ = dS @ K / sqrt(d)                     # (n, d)
dK = dS.T @ Q / sqrt(d)                   # (m, d)

Masked positions have P = 0, so they get dS = 0 automatically: the backward pass needs no special case for the mask.

Inputs

• Q, K, V as above; dOut: (n, d_v)
• causal: bool; if True, query i may only attend to keys j <= i (requires n == m)

Expected Return Format

• attention_forward(Q, K, V, causal=False): Out, array (n, d_v)
• attention_backward(Q, K, V, dOut, causal=False): tuple (dQ, dK, dV) with the shapes of Q, K, V

Edge cases and conventions

• Use a numerically stable softmax (subtract the row max); the tests use scores of size about 100.
• Recompute P inside the backward pass (FlashAttention does the same instead of storing it).
• The hidden tests compare your gradients with finite differences of L = sum(Out * dOut).

Constraints

NumPy only, vectorised; no autograd.

Follow-up

• FlashAttention's backward pass computes rowsum(dP * P) as rowsum(dOut * Out). Why are they equal, and why is the second one cheaper?
• What is the memory cost of storing P for a 32k-token context with 32 heads in fp16?

Input: attention_forward(np.zeros((2, 4)), np.zeros((3, 4)), np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]))
Output: [[3.0, 4.0], [3.0, 4.0]]

Input: attention_backward(np.zeros((2, 4)), np.zeros((3, 4)), np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]), np.ones((2, 2)))[2]
Output: [[0.6667, 0.6667], [0.6667, 0.6667], [0.6667, 0.6667]]

import numpy as np

def attention_forward(Q, K, V, causal=False):
    """
    Scaled dot-product attention.

    Args:
        Q: Queries (n, d)
        K: Keys (m, d)
        V: Values (m, d_v)
        causal: If True, query i may only attend to keys j <= i (n == m)

    Returns:
        Out: (n, d_v)
    """
    # Your code here
    pass

def attention_backward(Q, K, V, dOut, causal=False):
    """
    Gradients of attention with respect to Q, K and V.

    Args:
        Q, K, V: As in attention_forward
        dOut: Upstream gradient dL/dOut, shape (n, d_v)
        causal: Same flag as in the forward pass

    Returns:
        (dQ, dK, dV) with the shapes of Q, K, V
    """
    # Your code here
    pass