Recompute probabilities and differentiate the normalized row

For one row, the derivative of the output with respect to a probability is the upstream output gradient dotted with that key’s value. Softmax then subtracts a probability-weighted row mean from each of these derivatives. This subtraction captures the denominator’s dependence on every score. Ignoring it treats probabilities as independent exponentials and produces gradients whose sum generally fails the score-shift invariance check.

The weighted mean can also be computed by dotting the upstream gradient with the forward output, because the output is the probability-weighted sum of values. This identity avoids another materialized row reduction in suitable tiled schedules. The formulas here omit dropout and use one row with finite allowed logits. Masked scores contribute zero probability and gradient under a fixed mask; an entirely masked row follows the separately declared extension rather than an undefined normalized distribution.

Saving LSE equal to the logarithm of the row’s exponential sum lets backward reconstruct each probability as exp(score−LSE). The score tile is recomputed from Q and K using the forward scaling and bias rules. This is not another independently normalized local softmax: its denominator belongs to the entire original row. After obtaining score gradients, the usual matrix products form gradients with respect to Q, K, and V, including the score scale in both Q and K derivatives.

Tiled accumulation must include contributions from every tile that shares an input. One query receives gradient contributions from many key tiles, and a key/value can receive contributions from many query tiles. GQA adds another sharing axis because several query heads use the same KV head, requiring their gradient contributions to sum. Race-free accumulation, numerical precision, and deterministic reduction choices are implementation concerns that can change performance without changing the required mathematical derivative.

Attention dropout modifies probabilities after softmax in common training formulations. Recomputing backward needs the same logical dropout mask and scaling as forward, often regenerated from a saved random state and a stable mapping from logical indices to random values. A different tile schedule must not silently assign different draws to the same logical operation if exact replay is required. Disable dropout for the first deterministic correctness tests, then test the stochastic contract separately.

An exact algorithm can still have finite-precision differences because the reduction order changes between the materialized and tiled versions. Compare gradients against a higher-precision reference and use small finite-difference tests on selected smooth inputs. Include large score ranges, unusual lengths, shared heads, and masked rows. Extra recomputation is worthwhile when saved traffic and memory dominate its cost, but neither the memory reduction nor the algebra proves that every shape runs faster on every accelerator.