Exact tiled softmax attention with block-recomputed vector-Jacobian products

Inputs: Positive integers $n,d,v$, finite real matrices $Q,K\in {\mathbb{R}}^{n\times d}$ and $V\in {\mathbb{R}}^{n\times v}$, a real scale $c$, a positive integer key-block size $b$, and fixed nonempty allowed-key sets $A_1,\dots ,A_n\subseteq \{1,\dots ,n\}$. The algorithm reads membership in each $A_i$ on demand; any caller-owned representation of the mask is input storage. All operations below are in exact real arithmetic. Set $J_t=\{(t-1)b+1,\dots ,\min (tb,n)\}$ for $t=1,\dots ,\lceil n\mathrm{/}b\rceil$.

Forward procedure:

1. For each query row $i$, set the current summary to $\mathrm{\perp }$.
2. For $t=1,\dots ,\lceil n\mathrm{/}b\rceil$, form $I=J_t\cap A_i$. If $I=\mathrm{\varnothing }$, skip this block and leave the summary unchanged. Otherwise compute ${\mathrm{\Sigma }}_i(I)$. If the current summary is $\mathrm{\perp }$, copy ${\mathrm{\Sigma }}_i(I)$; otherwise merge the current summary with ${\mathrm{\Sigma }}_i(I)$ using $\star$.
3. After the final block, write $Y_i=z_i\mathrm{/}{\mathrm{\ell }}_i$ and retain the final $m_i$ and ${\mathrm{\ell }}_i$ for a later reverse pass.
4. Return $Y\in {\mathbb{R}}^{n\times v}$ and the retained row statistics.

Reverse procedure: Given a forward run with the same $Q,K,V,c,A_1,\dots ,A_n$ and an upstream matrix $G\in {\mathbb{R}}^{n\times v}$, initialize $dQ,dK,dV$ to zero. For each query row $i$, compute ${\delta }_i={\sum }_{a=1}^v{G}_{ia}{Y}_{ia}$. For every block $J_t$ and each $j\in J_t\cap A_i$, recompute $s_{ij}$ from $Q,K$ and set

$$p_{ij}=\frac{e^{s_{ij}-m_i}}{{\mathrm{\ell }}_i},{\gamma }_{ij}=p_{ij}(\sum _{a=1}^v{G}_{ia}{V}_{ja}-{\delta }_i)\mathrm{.}$$

Then, for $r=1,\dots ,d$, add to the query and key gradient arrays, and for $a=1,\dots ,v$, add to the value gradient array:

$$(dQ{)}_{ir}+=c{\gamma }_{ij}{K}_{jr},(dK{)}_{jr}+=c{\gamma }_{ij}{Q}_{ir},(dV{)}_{ja}+=p_{ij}{G}_{ia}\mathrm{.}$$

Return $dQ,dK,dV$. Do not store a full score or probability matrix: the reverse procedure recomputes one block's scores from the same inputs and fixed allowed-key sets. The prescribed reverse procedure treats the forward maxima and normalizers as frozen scaling statistics; it does not differentiate a maximum operation. The base algorithm has no dropout. Any dropout extension must replay the same realized dropout mask in its reverse pass.

Summary state and merge convention

Masked scaled dot-product attention

Addition of contributions from repeated uses of shared inputs

Exact reverse-mode differentiation of a finite smooth real-arithmetic graph

Dao et al., §§2–3, Appendix B.2–B.4, and Appendix C: tiled forward, recomputed backward, and the exactness proof

Milakov and Gimelshein, §§2–3: stable online softmax normalizer

MIT OCW, Chapter 8 §§8.3–8.4: reverse-mode chain rule and accumulation at shared inputs