Tiled attention and its recomputed reverse pass are exact

Fix a query row $i$ and let $S$ be the union of the allowed keys in the blocks processed so far. Initially $S=\mathrm{\varnothing }$ and the algorithm's state is $\mathrm{\perp }$. An empty allowed block leaves both $S$ and the state unchanged. If $S=\mathrm{\varnothing }$ and the next allowed block $I$ is nonempty, copying ${\mathrm{\Sigma }}_i(I)$ gives the defining summary of $S\cup I$. If $S$ and $I$ are both nonempty, the merged maximum is $\max (m_i(S),m_i(I))=m_i(S\cup I)$, and $e^{m_i(S)-m_i(S\cup I)}{\mathrm{\ell }}_i(S)={\sum }_{j\in S}{e}^{s_{ij}-m_i(S\cup I)}$ while $e^{m_i(I)-m_i(S\cup I)}{\mathrm{\ell }}_i(I)={\sum }_{j\in I}{e}^{s_{ij}-m_i(S\cup I)}$. The same two identities with $V_j$ inside each sum give the merged $z$ coordinate. Disjoint blocks therefore preserve the summary invariant.

The key blocks partition $\{1,\dots ,n\}$, so after the last block $S=A_i\mathrm{\ne }\mathrm{\varnothing }$. The maximum $m_i$ is finite and at least one term in ${\mathrm{\ell }}_i={\sum }_{j\in A_i}{e}^{s_{ij}-m_i}$ equals $1$, hence ${\mathrm{\ell }}_i>0$. Dividing $z_i$ by ${\mathrm{\ell }}_i$ and multiplying numerator and denominator by $e^{m_i}$ gives the displayed dense masked attention row. This also proves every retained $m_i,{\mathrm{\ell }}_i$ is defined.

With the allowed sets fixed, every score is a polynomial in the entries of $Q,K$; finite exponentials, sums, and division by the strictly positive dense denominator make the dense row map smooth on all of ${\mathbb{R}}^{n(2d+v)}$. Its scalarization $\mathrm{\Phi }$ is therefore a finite $C^1$ real-arithmetic graph. The reverse-mode theorem and chain rule apply to this dense graph; the maximum used only to scale the tiled forward computation is not a primitive of this graph.

For $j,k\in A_i$, write $p_{ij}=e^{s_{ij}}\mathrm{/}{\sum }_{\mathrm{\ell }\in A_i}{e}^{s_{i\mathrm{\ell }}}$, as justified by step 2.1. Direct quotient differentiation gives $\frac{\mathrm{\partial }{p}_{ik}}{\mathrm{\partial }{s}_{ij}}=p_{ik}({\mathbf{1}}_{\{k=j\}}-p_{ij})$. Thus for ${\mathrm{\Phi }}_i={\sum }_a{G}_{ia}{Y}_{ia}$ the score derivative is $\frac{\mathrm{\partial }{\mathrm{\Phi }}_i}{\mathrm{\partial }{s}_{ij}}={\sum }_{k\in A_i}({\sum }_a{G}_{ia}{V}_{ka})\frac{\mathrm{\partial }{p}_{ik}}{\mathrm{\partial }{s}_{ij}}=p_{ij}({\sum }_a{G}_{ia}{V}_{ja}-{\sum }_a{G}_{ia}{Y}_{ia})={\gamma }_{ij}$. The direct value derivative is ${\mathrm{\nabla }}_{V_j}{\mathrm{\Phi }}_i=p_{ij}{G}_i$.

Since $s_{ij}=c{\sum }_r{Q}_{ir}{K}_{jr}$, its coordinate derivatives are $\mathrm{\partial }{s}_{ij}\mathrm{/}\mathrm{\partial }{Q}_{ir}=c{K}_{jr}$ and $\mathrm{\partial }{s}_{ij}\mathrm{/}\mathrm{\partial }{K}_{jr}=c{Q}_{ir}$. The chain rule therefore gives the row-$i$ contributions $c{\sum }_{j\in A_i}{\gamma }_{ij}{K}_j$ to $Q_i$, $c{\gamma }_{ij}{Q}_i$ to $K_j$, and $p_{ij}{G}_i$ to $V_j$. Entries of other query rows receive no contribution from this row.

Set $\mathrm{\Theta }={\mathbb{R}}^{2nd+nv}$, $H=\mathbb{R}$, $h_0(\theta )=0$, $F_i(\theta ,h)=h+{\sum }_a{G}_{ia}{Y}_{ia}(\theta )$ for $i=1,\dots ,n$, and terminal loss $\lambda (\theta ,h)=h$. Step 3.1 makes every $F_i$ continuously differentiable on the open domain $\mathrm{\Theta }\times H$. By induction on $i$, the recurrence state after $i$ transitions is $h_i={\sum }_{r=1}^i{\sum }_a{G}_{ra}{Y}_{ra}(\theta )$: the formula holds at $i=0$ because $h_0=0$, and adding the $i$th transition term proves it at $i$ from its value at $i-1$. Thus its scalar output is $h_n=\mathrm{\Phi }(\theta )$. The same parameter $\theta =(Q,K,V)$ is supplied at every transition, and the shared-parameter gradient formula sums each transition's parameter pullback. It therefore adds each row's pullback at the common $K,V$ inputs, while the row-specific $Q_i$ pullbacks occupy their corresponding rows. There is no direct loss or initial-state parameter term because $\lambda$ is $h$ and $h_0$ is constant. Thus the algorithm's reverse updates give ${\mathrm{\nabla }}_Q\mathrm{\Phi },{\mathrm{\nabla }}_K\mathrm{\Phi },{\mathrm{\nabla }}_V\mathrm{\Phi }$.

At the forward values, the summary invariant in steps 1.1–2.1 gives ${\mathrm{\ell }}_i={\sum }_{j\in A_i}{e}^{s_{ij}-m_i}$. Therefore, for each $j\in A_i$, $p_{ij}=\frac{e^{s_{ij}-m_i}}{{\mathrm{\ell }}_i}=\frac{e^{s_{ij}}}{{\sum }_{\mathrm{\ell }\in A_i}{e}^{s_{i\mathrm{\ell }}}}\mathrm{.}$ This is the dense softmax probability evaluated at the current scores. Steps 4.1–6.1 differentiate the smooth dense attention map and show that substituting these probabilities into the reverse updates gives its exact gradient at these inputs. The stored $m_i$ and ${\mathrm{\ell }}_i$ are used as normalization shifts in those updates; no derivative of either statistic is needed. The argument applies whether or not the maximum score is tied.