Attention / Transformer Module with Backward Pass

ML coding round on autoregressive architectures: implement Attention and Transformer components, then complete the backward pass around attention when softmax backward is provided. Requirements Implement the core Attention module used in an autoregressive architecture. Implement or sketch the surrounding Transformer module. Complete the backward logic for the attention computation. Softmax backward may be given; the remaining gradient flow must be filled in. Be ready to explain training behavior for autoregressive models, not only the forward pass. Notes The prompt tests both tensor-shape discipline and gradient reasoning. Write down the shapes for Q, K, V, scores, probabilities, and output before coding. Scaled dot-product attention has the canonical form softmax(QK^T / sqrt(d_k) + mask) V; for autoregressive self-attention, the mask must block future keys before the softmax. If softmax backward is provided, the remaining work is chaining gradients through O = P V, score scaling, and S = QK^T: dV = P^T dO, dP = dO V^T, dQ = dS K / sqrt(d_k), and dK = dS^T Q / sqrt(d_k). Expect conceptual follow-ups around Transformer structure, autoregressive training, and why masking is required. Preparation On paper, derive the backward pass for a single-head attention block with batch omitted, then re-add batch and head dimensions. Implement a NumPy attention forward/backward on a toy tensor and finite-difference check Q, K, and V gradients. Rehearse the causal-mask explanation: logits for future positions become effectively impossible before softmax, so training can parallelize without leaking target tokens.