Pipeline-generated

How statement and proof provenance work

The first chip identifies the source of the statement or construction; the second identifies the source of its local proof or verification.

Literature-sourced: the exact statement appears in a cited source; only wording and notation differ.
AI-adapted: a semantically identical restatement of literature-sourced material, modulo indexing, notation, and boundary cases adopted by the library.
AI-generated: a genuinely novel statement formulated by AI, with no source for the claim itself.

These labels describe origin, not correctness: citations and verification chips remain separate evidence.

✓ 2 results · all verified · 2 also independently AI-judged

Every result on this page is machine-checked by a proof checker and read in full and owner-audited; the judge is an additional, independent cross-model AI review of the proofs; all 2 also cleared it.

Exact tiled attention and memory-aware differentiation

1 · Prerequisites

2 · Logical flowchart

3 · Definitions, theorems and proofs

DefinitionDefinition: Literature-sourcedProof: Not applicableOpen item page →

Stable summary of an attention key block

Definition

Let $n, d, v$ be positive integers, let $Q, K \in R^{n \times d}$ and $V \in R^{n \times v}$ , let $c \in R$ , and for each query row $i \in {1, \dots, n}$ let $A_{i} \subseteq {1, \dots, n}$ be a nonempty set of allowed keys. Write $q_{i}$ for row $i$ of $Q$ , $k_{j}$ for row $j$ of $K$ , $v_{j}$ for row $j$ of $V$ , and $s_{i j} = c \sum_{r = 1}^{d} Q_{i r} K_{j r}$ for $j \in A_{i}$ . For a nonempty finite set $S \subseteq A_{i}$ , its attention summary is

Σ_{i} (S) = (m_{i} (S), ℓ_{i} (S), z_{i} (S)), m_{i} (S) = \max_{j \in S} s_{i j}, ℓ_{i} (S) = \sum_{j \in S} e^{s_{i j} - m_{i} (S)}, z_{i} (S) = \sum_{j \in S} e^{s_{i j} - m_{i} (S)} v_{j} .

The empty set has the separate state $⊥$ and has no maximum or numeric summary coordinates. For disjoint nonempty $S, T \subseteq A_{i}$ , define their merge by

m^{'} = \max (m_{i} (S), m_{i} (T)), ℓ^{'} = e^{m_{i} (S) - m^{'}} ℓ_{i} (S) + e^{m_{i} (T) - m^{'}} ℓ_{i} (T), z^{'} = e^{m_{i} (S) - m^{'}} z_{i} (S) + e^{m_{i} (T) - m^{'}} z_{i} (T) .

and set $Σ_{i} (S) ⋆ Σ_{i} (T) = (m^{'}, ℓ^{'}, z^{'})$ . This summary stores the maximum score, a positive shifted normalizer, and the correspondingly shifted weighted value sum; dividing its third coordinate by its second gives the attention output on $S$ .

Definition references

Masked row-softmax attention and its nonempty-row convention

Dao et al., §3.1 and Appendix C: online maximum, normalizer, and weighted-value summaries

Milakov and Gimelshein, §3: online maximum and softmax normalizer

AlgorithmAlgorithm: Literature-sourcedProof: Not applicablejudge pass (gpt-6-luna)Open item page →

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

Definition

Inputs: Positive integers $n, d, v$ , finite real matrices $Q, K \in R^{n \times d}$ and $V \in 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 (t b, n)}$ for $t = 1, \dots, ⌈ n / b ⌉$ .

Forward procedure:

For each query row $i$ , set the current summary to $⊥$ .
For $t = 1, \dots, ⌈ n / b ⌉$ , form $I = J_{t} \cap A_{i}$ . If $I = \emptyset$ , skip this block and leave the summary unchanged. Otherwise compute $Σ_{i} (I)$ . If the current summary is $⊥$ , copy $Σ_{i} (I)$ ; otherwise merge the current summary with $Σ_{i} (I)$ using $⋆$ .
After the final block, write $Y_{i} = z_{i} / ℓ_{i}$ and retain the final $m_{i}$ and $ℓ_{i}$ for a later reverse pass.
Return $Y \in 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 R^{n \times v}$ , initialize $d Q, d K, d V$ to zero. For each query row $i$ , compute $δ_{i} = \sum_{a = 1}^{v} G_{i a} Y_{i a}$ . For every block $J_{t}$ and each $j \in J_{t} \cap A_{i}$ , recompute $s_{i j}$ from $Q, K$ and set

p_{i j} = \frac{e^{s_{i j} - m_{i}}}{ℓ_{i}}, γ_{i j} = p_{i j} (\sum_{a = 1}^{v} G_{i a} V_{j a} - δ_{i}) .

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:

(d Q)_{i r} + = c γ_{i j} K_{j r}, (d K)_{j r} + = c γ_{i j} Q_{i r}, (d V)_{j a} + = p_{i j} G_{i a} .

Return $d Q, d K, d V$ . 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.

Definition references

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

TheoremStatement: Literature-sourcedProof: AI-generatedprecheck passjudge pass (gpt-6-luna)Open item page →

Tiled attention and its recomputed reverse pass are exact

Statement

Let $n, d, v$ be positive integers, $Q, K \in R^{n \times d}$ and $V \in R^{n \times v}$ be finite real matrices, $c \in R$ , $b$ a positive integer, and $A_{i} \subseteq {1, \dots, n}$ a fixed nonempty set of allowed keys for every query row $i$ . Write $Q_{i}$ for row $i$ of $Q$ , and $K_{j}, V_{j}$ for row $j$ of $K, V$ . Let $Y, m_{i}, ℓ_{i}$ be the forward outputs and row statistics of the exact tiled softmax attention algorithm, and fix an upstream matrix $G \in R^{n \times v}$ . Then each row of $Y$ equals the dense masked attention row

Y_{i} = \frac{\sum_{j \in A_{i}} e^{c Q_{i} K_{j}^{T}} V_{j}}{\sum_{j \in A_{i}} e^{c Q_{i} K_{j}^{T}}},

and the algorithm's recomputed reverse procedure returns the Euclidean gradients of the scalar function $Φ (Q, K, V) = \sum_{i = 1}^{n} \sum_{a = 1}^{v} G_{i a} Y_{i a}$ with respect to $Q, K, V$ , while holding $c$ , the sets $A_{i}$ , and $G$ fixed. Its reverse rule treats the stored maxima as fixed normalization shifts, so the conclusion also holds when scores tie at a row maximum. Every row is defined because $A_{i}$ is nonempty; no convention for an all-masked row is asserted.

Statement references

Dao et al., §3.1 and Appendix C: exactness of the online tiled attention recurrence

Dao et al., Appendix B.2–B.4: reverse-mode derivatives and block recomputation

Milakov and Gimelshein, §3: online normalizer recurrence and its induction proof

MIT OCW, Chapter 8 §§8.3–8.4: finite-graph reverse-mode chain rule

Facts & Assumptions

Given: The finite real inputs, the fixed nonempty allowed-key sets, the positive block size, the forward and reverse procedures, and the upstream matrix $G$ .

[F1]

Scaled dot-product attention: masked attention uses row softmax of the scaled query-key scores and is defined only when each row has an allowed key.

[F2]

Stable summary of an attention key block: for nonempty $S$ , the summary coordinates are $\max_{j \in S} s_{i j}$ , $\sum_{j \in S} e^{s_{i j} - m_{i} (S)}$ , and $\sum_{j \in S} e^{s_{i j} - m_{i} (S)} V_{j}$ ; the empty set is the separate symbol $⊥$ .

[F3]

Exact tiled softmax attention with block-recomputed vector-Jacobian products: the forward pass skips empty allowed blocks, copies the first nonempty block summary, merges later summaries, and the reverse pass recomputes $p_{i j}$ and accumulates the stated $d Q, d K, d V$ entries.

[F4]

Finite straight-line computational programs and real-arithmetic tapes: a finite topologically ordered real-arithmetic graph has differentiable primitive maps on open domains and repeated input uses are separate parent slots.

[F5]

Reverse mode computes the exact gradient of the real-arithmetic graph: for a finite $C^{1}$ program with one scalar output on a fixed strict branch, reverse mode returns the exact derivative of that output.

[F6]

The chain rule for total derivatives: $D (g \circ f) (a) = D g (f (a)) \circ D f (a)$ : for totally differentiable finite-dimensional maps, the derivative of a composition is the ordered composition of the derivatives.

[F7]

Backpropagation through a finite shared-parameter recurrence: the same parameter vector is supplied to every transition.

[F8]

Reverse accumulation sums gradients over shared parameter uses: the gradient formula includes a pullback term for each transition using the shared parameter.

Proof

technique · identify the exact summary invariant, then differentiate the equivalent dense masked attention map and group repeated input contributions

1.1F2F3givenconstruct

Fix a query row $i$ and let $S$ be the union of the allowed keys in the blocks processed so far. Initially $S = \emptyset$ and the algorithm's state is $⊥$ . An empty allowed block leaves both $S$ and the state unchanged. If $S = \emptyset$ and the next allowed block $I$ is nonempty, copying $Σ_{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)} ℓ_{i} (S) = \sum_{j \in S} e^{s_{i j} - m_{i} (S \cup I)}$ while $e^{m_{i} (I) - m_{i} (S \cup I)} ℓ_{i} (I) = \sum_{j \in I} e^{s_{i j} - 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.

2.1F1F2F3step 1.1algebra

The key blocks partition ${1, \dots, n}$ , so after the last block $S = A_{i} \neq \emptyset$ . The maximum $m_{i}$ is finite and at least one term in $ℓ_{i} = \sum_{j \in A_{i}} e^{s_{i j} - m_{i}}$ equals $1$ , hence $ℓ_{i} > 0$ . Dividing $z_{i}$ by $ℓ_{i}$ and multiplying numerator and denominator by $e^{m_{i}}$ gives the displayed dense masked attention row. This also proves every retained $m_{i}, ℓ_{i}$ is defined.

3.1F4F5F6step 2.1given

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 $R^{n (2 d + v)}$ . Its scalarization $Φ$ 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.

4.1F5F6step 2.1step 3.1algebra

For $j, k \in A_{i}$ , write $p_{i j} = e^{s_{i j}} / \sum_{ℓ \in A_{i}} e^{s_{i ℓ}}$ , as justified by step 2.1. Direct quotient differentiation gives $\frac{\partial p_{i k}}{\partial s_{i j}} = p_{i k} (1_{{k = j}} - p_{i j})$ . Thus for $Φ_{i} = \sum_{a} G_{i a} Y_{i a}$ the score derivative is $\frac{\partial Φ_{i}}{\partial s_{i j}} = \sum_{k \in A_{i}} (\sum_{a} G_{i a} V_{k a}) \frac{\partial p_{i k}}{\partial s_{i j}} = p_{i j} (\sum_{a} G_{i a} V_{j a} - \sum_{a} G_{i a} Y_{i a}) = γ_{i j}$ . The direct value derivative is $\nabla_{V_{j}} Φ_{i} = p_{i j} G_{i}$ .

5.1F6step 4.1algebra

Since $s_{i j} = c \sum_{r} Q_{i r} K_{j r}$ , its coordinate derivatives are $\partial s_{i j} / \partial Q_{i r} = c K_{j r}$ and $\partial s_{i j} / \partial K_{j r} = c Q_{i r}$ . The chain rule therefore gives the row- $i$ contributions $c \sum_{j \in A_{i}} γ_{i j} K_{j}$ to $Q_{i}$ , $c γ_{i j} Q_{i}$ to $K_{j}$ , and $p_{i j} G_{i}$ to $V_{j}$ . Entries of other query rows receive no contribution from this row.

6.1F7F8step 3.1step 5.1algebra

Set $Θ = R^{2 n d + n v}$ , $H = R$ , $h_{0} (θ) = 0$ , $F_{i} (θ, h) = h + \sum_{a} G_{i a} Y_{i a} (θ)$ for $i = 1, \dots, n$ , and terminal loss $λ (θ, h) = h$ . Step 3.1 makes every $F_{i}$ continuously differentiable on the open domain $Θ \times H$ . By induction on $i$ , the recurrence state after $i$ transitions is $h_{i} = \sum_{r = 1}^{i} \sum_{a} G_{r a} Y_{r a} (θ)$ : 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} = Φ (θ)$ . The same parameter $θ = (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 $λ$ is $h$ and $h_{0}$ is constant. Thus the algorithm's reverse updates give $\nabla_{Q} Φ, \nabla_{K} Φ, \nabla_{V} Φ$ .

7.1F2F3step 2.1step 3.1step 4.1step 5.1step 6.1algebra∎

At the forward values, the summary invariant in steps 1.1–2.1 gives $ℓ_{i} = \sum_{j \in A_{i}} e^{s_{i j} - m_{i}}$ . Therefore, for each $j \in A_{i}$ , $p_{i j} = \frac{e^{s_{i j} - m_{i}}}{ℓ_{i}} = \frac{e^{s_{i j}}}{\sum_{ℓ \in A_{i}} e^{s_{i ℓ}}} .$ 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 $ℓ_{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.

TheoremStatement: Literature-sourcedProof: AI-generatedprecheck passjudge pass (gpt-6-luna)Open item page →

Tiled attention uses linear retained memory and quadratic arithmetic

Statement

In the unit-cost real-arithmetic model where scalar arithmetic, exponential evaluation, comparisons, and one allowed-key membership query each cost $O (1)$ , the forward and reverse procedures of the exact tiled softmax attention algorithm use $O (n^{2} (d + v))$ arithmetic operations. The forward procedure retains the $Q, K, V$ inputs, the $n \times v$ output, and two scalar statistics per query, requiring $O (n (d + v))$ stored input and output scalars, with at most $O (b + v)$ additional live scalars for one block's scores and weighted-value summary. The reverse procedure additionally takes $G$ and returns $d Q, d K, d V$ ; counting these as input and output arrays, it also uses $O (n (d + v))$ stored input and output scalars and at most $O (b + v)$ additional live scalars. Neither procedure materializes an $n \times n$ score or probability array. Storage for a caller-owned mask representation is input storage and is not part of the additional workspace; if it is a dense array, its own $O (n^{2})$ entries must be counted separately. The exactness of the forward and reverse outputs is the assertion of the tiled-attention exactness theorem. No hardware I/O or floating-point operation count is asserted.

Statement references

Dao et al., §3.1 and Appendix C: forward arithmetic and extra-memory count

Dao et al., Appendix B.2–B.4: recomputed reverse pass and its linear extra-memory claim

Exact forward output and reverse-mode derivatives

Facts & Assumptions

Given: The positive integers $n, d, v, b$ , fixed allowed-key sets, and the two procedures of the exact tiled softmax attention algorithm.

[F1]

Exact tiled softmax attention with block-recomputed vector-Jacobian products: the forward procedure processes one query row and one key block at a time, while the reverse procedure recomputes each permitted score and accumulates gradients into output arrays.

[F2]

Stable summary of an attention key block: a nonempty block summary stores one scalar maximum, one scalar normalizer, and one vector in $R^{v}$ ; an empty block stores no numeric state.

[F3]

Tiled attention and its recomputed reverse pass are exact: the two procedures return the dense masked attention and its exact vector-Jacobian product.

[F4]

Exact tiled softmax attention with block-recomputed vector-Jacobian products: the reverse procedure recomputes each block's scores and does not store a full score or probability matrix.

[A1]

The mask membership accessor is read-only, has $O (1)$ query cost, and any storage for its representation is caller-owned input storage rather than algorithmic scratch.

[A2]

The operation count is scalar arithmetic, exponential evaluation, comparisons, and mask membership queries; all are unit cost in this model.

Proof

technique · count the live arrays, block workspace, and operations in the explicit loops

1.1F1F2A1givenalgebra

The forward procedure retains $Q, K, V, Y$ and $(m_{i}, ℓ_{i})_{i = 1}^{n}$ . Their scalar counts are $n d + n d + n v + n v + 2 n = O (n (d + v))$ because $d, v \geq 1$ . At any query and key block it stores at most $b$ scores and the current and block weighted-value vectors of length $v$ , plus a constant number of scalars. This is $O (b + v)$ extra live storage, and it discards the block workspace before moving to the next block.

1.2F1F2A2givenalgebra

The forward loops examine at most $n^{2}$ query-key pairs. For each pair, the score dot product uses $O (d)$ operations; its normalized exponential and weighted-value contribution use $O (1 + v)$ operations. Block maxima, membership queries, and scalar loop overhead cost $O (n^{2})$ in total. There are at most $n ⌈ n / b ⌉ \leq n^{2}$ block visits, each summary copy or merge costs $O (v)$ , and final row divisions cost $O (n v)$ . Thus forward arithmetic is $O (n^{2} (d + v))$ . Only one block's scores and weighted-value summaries are live, so no $n \times n$ score or probability array is allocated.

1.3F4A1givenalgebra

The reverse procedure's retained arrays are $Q, K, V, Y, G, d Q, d K, d V$ and the $2 n$ row statistics. Their scalar counts are all $O (n (d + v))$ . For one row and one key it recomputes one score and one probability, forms scalar inner products, and updates entries in the already allocated gradient arrays; it needs only a constant number of additional scalars, which is within $O (b + v)$ . It does not store the recomputed probabilities across pairs.

2.1F3F4A2givenalgebra∎

Each query-key pair in the reverse pass costs $O (d + v)$ operations: one score, one inner product with $V_{j}$ , and vector updates with lengths $d$ and $v$ . There are at most $n^{2}$ pairs, and the finite loop overhead is at most $O (n^{2})$ , so reverse arithmetic is $O (n^{2} (d + v))$ . Since scores and probabilities are recomputed one at a time, this pass also creates no $n \times n$ array. The separate exactness claim is supplied by [F3].

RemarkRemark: Literature-sourcedProof: Not applicablejudge pass (gpt-6-luna)Open item page →

The scalar workspace result is not a hardware I/O bound

Remarks

The workspace theorem counts scalar storage and operations for the explicitly defined real-arithmetic procedure. It does not count reads and writes between GPU high-bandwidth memory and on-chip SRAM, and it gives no hardware speedup claim. The FlashAttention paper's $Θ (n^{2} d^{2} / M)$ HBM-access result assumes its stated memory model, block schedule, and range $d \leq M \leq n d$ ; its endpoint input-size argument is not a lower bound for every fixed memory size. Neither claim follows from the scalar workspace theorem here.

The maximum shift makes every exact-real exponent argument $s_{i j} - m_{i}$ nonpositive, but this does not specify rounded arithmetic or rule out underflow of very small finite-precision terms. Different floating-point accumulation orders can return different bit patterns. The base procedure includes no dropout; if an extension uses dropout, its fixed realized mask must be reused in the reverse pass because a different mask changes the differentiated function.

Remarks references

The scalar operation and workspace theorem proved for this algorithm

Dao et al., §§2.1–3.2 and Appendix B: the paper's hardware-memory model and recomputed backward pass

4 · Examples, counterexamples and false statements

None yet.

Exact tiled attention and memory-aware differentiation

2 · Logical flowchart

3 · Definitions, theorems and proofs

Stable summary of an attention key block

Definition

Definition references

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

Definition

Definition references

Tiled attention and its recomputed reverse pass are exact

Statement

Statement references

Facts & Assumptions

Proof

Tiled attention uses linear retained memory and quadratic arithmetic

Statement

Statement references

Facts & Assumptions

Proof

The scalar workspace result is not a hardware I/O bound

Remarks

Remarks references

4 · Examples, counterexamples and false statements

Sources