Multi-Layer Transformers Gradient Can be Approximated in Almost Linear Time

ArXiv: 2408.13233

🎯 Pitch

The backpropagation gradient for a multi-layer transformer, long thought to require quadratic Ω(n²) time in the input length, can actually be computed in almost linear n^{1+o(1)} time with only a polynomial small error. This robs the attention mechanism of its most infamous scaling bottleneck for training, holding even with residual connections, causal masks, and multi-head attention—and thus unlocks sub-quadratic training for long-context models.

---

1. Executive Summary

This paper proves that the gradient computation in multi-layer transformer models can be approximated in almost linear time $n^{1+o(1)}$ — where $n$ is the input sequence length — while maintaining a polynomially small approximation error of $1/\text{poly}(n)$ across the entire model. The theoretical results hold for general loss functions and practical transformer sub-modules, including residual connections, causal masks, and multi-head attention, with the key mechanism being a low-rank approximation of the attention matrix via polynomial kernel methods (enabling decomposition of the $n \times n$ attention matrix into $U_1 V_1^\top$ where $U_1, V_1 \in \mathbb{R}^{n \times k_1}$ with $k_1 = n^{o(1)}$) that is then propagated through all gradient components — the gradient on intermediate variables $T_i(X)$, on the key-query weight product $W_i = W_{Q_i} W_{K_i}^\top$, and on the value weight $W_{V_i}$. The algorithm breaks the quadratic $\Omega(n^2)$ bottleneck in standard backpropagation, establishing that training can be accelerated from quadratic to sub-quadratic time complexity with bounded error propagation across $m$ layers — a guarantee that holds under the mild assumptions that the hidden dimension $d = O(\log n)$ and each matrix entry uses $O(\log n)$ bits.

2. Context and Motivation

The Core Problem: The Quadratic Bottleneck in Transformer Training

The fundamental problem this paper addresses is stark and well-known to anyone who has trained large language models: the self-attention mechanism at the heart of transformer architectures has quadratic time complexity $\Omega(n^2 d)$ with respect to the input sequence length $n$. This is not merely a theoretical concern — it is the dominant computational cost in both training and inference, and it scales disastrously as context lengths grow.

To understand why this is so serious, consider the numbers. A modern LLM like LLaMA 3.1 405B supports a context length of $n = 128$k tokens with hidden dimension $d = 4096$. Each self-attention block must compute attention scores between every pair of tokens — that is $n^2$ pairwise interactions — and then multiply by the value matrix. The paper notes that training such a model requires 30.84M GPU hours. Even if only, say, 20% of that compute goes to attention, that is still millions of GPU hours spent on a quadratic operation. If the sequence length doubles, the attention cost quadruples. This is unsustainable for the trajectory the field is on, where context lengths are rapidly growing (from 2k to 8k to 32k to 128k and beyond) to support tasks like long-document summarization, multi-turn conversation, retrieval-augmented generation, and in-context learning with many examples.

The problem is both practical and fundamental. Practically, it drives up training costs, energy consumption, and carbon emissions. The paper explicitly cites this:

"this quadratic time complexity results in critical challenges: (i) a marked decrease in training efficiency... and (ii) significant energy usage, which in turn contributes to higher carbon dioxide emissions" (Section 1)

Theoretically, it raises a natural question: is the quadratic dependence on $n$ inherent to the attention mechanism, or is it an artifact of the naive algorithm? The input to a self-attention layer is $X \in \mathbb{R}^{n \times d}$, which contains only $n \cdot d = n^{1+o(1)}$ numbers (when $d = O(\log n)$). Yet the intermediate attention matrix $f(X) \in \mathbb{R}^{n \times n}$ has $n^2$ entries. This discrepancy — an $n \times n$ object constructed from an $n \times d$ input — suggests that the attention matrix may be intrinsically low-rank, and that the quadratic cost of working with it explicitly might be avoidable. The paper's entire approach is built on exploiting this suspected low-rank structure.

Forward vs. Backward: Where Prior Work Stopped Short

A critical distinction the paper makes — and which motivates its existence — is between the forward pass (inference) and the backward pass (training). The forward pass computes the attention output:

$\text{Attn}(X) = \text{Softmax}(X W_Q W_K^\top X^\top / d) \cdot X W_V$

The backward pass computes the gradients of the loss with respect to all parameters $(W_Q, W_K, W_V)$ and intermediate variables $(T_i(X))$ via backpropagation. The backward pass is significantly more complex than the forward pass, requiring gradients of the softmax normalization, the matrix multiplications, and the exponential function — all composed through the chain rule across multiple layers.

A seminal prior work by Alman and Song (2023) — which this paper builds on directly — proved that the forward pass of self-attention can be approximated in almost linear time $n^{1+o(1)}$ using polynomial kernel methods to obtain a low-rank representation of the attention matrix. This was a breakthrough: it showed that inference with long contexts need not be quadratic. But their result only covered inference, not training. Training involves computing gradients, which the Alman and Song (2023) work did not address.

Another prior work, Alman and Song (2024a), took the first step toward addressing the backward pass. They proved that the gradient computation for a single self-attention layer could also be approximated in almost linear time. However, this work had several critical limitations that prevented it from being applicable to real transformer training (Section 4.3):

1. Single-layer only: They computed gradients for one attention layer in isolation. There was no mechanism for propagating gradients through multiple layers — the chain rule across layers introduces error accumulation that their analysis did not handle.

2. Specific loss function: Their analysis assumed a $\ell_2$ loss, which is far from the cross-entropy loss universally used in language model training. Extending to general loss functions is non-trivial because the gradient expressions become more complex and involve different matrix structures.

3. Incomplete gradient computation: They only computed gradients for $W_{Q_i}$ and $W_{K_i}$ (actually their product $W_i = W_{Q_i} W_{K_i}^\top$). They did not compute:

   • The gradient on $W_{V_i}$ (the value weight matrix)
   • The gradient on the intermediate variables $T_i(X)$, which is essential for backpropagation through multiple layers — without it you cannot chain layers together
   • The gradient on the input $X$ itself, which matters for applications like prompt tuning

These omissions meant that Alman and Song (2024a) provided a theoretical result for a toy setting — a single attention layer, trained with $\ell_2$ loss, without computing all gradients — rather than a result applicable to the multi-layer transformers used in practice.

The Gap: No Sub-Quadratic Training Algorithm for Real Multi-Layer Transformers

Despite these prior results on accelerating individual attention operations, no existing work had provided a sub-quadratic algorithm for training a complete multi-layer transformer with general loss functions. This is the gap the paper fills. The challenge is not just accelerating individual gradient components, but doing so in a way that:

1. Handles error propagation across layers: In an $m$-layer transformer, gradients flow backward from layer $m$ to layer 1 through the chain rule. If each layer introduces an approximation error of $\epsilon$, a naive analysis would suggest the final error grows as $m \cdot \epsilon$ or worse. For the result to be meaningful, the paper must prove that the error remains bounded by $1/\text{poly}(n)$ across the entire model, not just per-layer.

2. Works for general loss functions: Real training uses cross-entropy loss (or variants). The gradient expressions differ from $\ell_2$ loss — different terms appear, different matrix structures emerge. Prior work's $\ell_2$-specific derivations do not carry over.

3. Handles all gradient components required for backpropagation: Computing $\frac{dL}{dW_{Q_i}}$ and $\frac{dL}{dW_{K_i}}$ is not enough. You need $\frac{dL}{dW_{V_i}}$ to update the value weights, and critically, you need $\frac{dL}{dT_{i-1}(X)}$ to propagate the gradient to the previous layer. Without $\frac{dL}{dT_{i-1}(X)}$, there is no way to chain layers together in backpropagation.

4. Accommodates practical transformer sub-modules: Real transformers have residual connections (which change the gradient flow), causal attention masks (which break the low-rank structure), and multi-head attention (which splits the attention into multiple parallel heads). Any practically relevant algorithm must handle these.

How This Paper Positions Itself

The paper positions itself not as proposing a fundamentally new attention mechanism, but as providing the theoretical foundation for efficient transformer training by extending the low-rank approximation approach from inference to the full backward pass of multi-layer models. The intellectual architecture is:

• Base technique: Polynomial kernel approximation of the attention matrix (from Alman and Song, 2023), which gives $f(X) \approx U_1 V_1^\top$ where $U_1, V_1 \in \mathbb{R}^{n \times k_1}$ with $k_1 = n^{o(1)}$. The key insight is that multiplying matrices in the right order — e.g., $U_1 (V_1^\top V)$ instead of $(U_1 V_1^\top) V$ — avoids ever forming the full $n \times n$ matrix.

• Extension to gradients: The paper shows that all gradient components that appear in backpropagation can be expressed in forms amenable to this low-rank multiplication trick — either as $f(X) \cdot H$ for some matrix $H$, or as $(f(X) \odot (U V^\top)) \cdot H$ (Hadamard product with another low-rank matrix). By categorizing gradient terms into these two patterns and providing fast algorithms for each, the paper covers every gradient computation needed.

• Extension to multiple layers: The paper proves error propagation bounds showing that composing $m$ approximate gradient computations does not blow up the error beyond $1/\text{poly}(n)$. This is the crucial step that makes the result applicable to real deep transformers rather than just single-layer curiosities.

• Handling practical sub-modules: The paper shows that residual connections (which add identity mappings to the gradient flow) introduce only $O(nd)$ overhead, causal masks can be handled by a specialized $O(nk)$ algorithm that exploits the lower-triangular structure, and multi-head attention is just a parallel application of the single-head algorithm with reduced per-head dimension.

The paper explicitly positions itself as enabling faster training, not just faster inference. This is a crucial distinction because training involves many forward and backward passes over enormous datasets — the quadratic attention cost is paid billions of times. Inference is often run once per user query. Training is where the quadratic bottleneck most severely limits what is practical.

The paper also positions its algorithm as complementary to system-level optimizations like FlashAttention (Dao et al., 2022; Dao, 2023), which reduce memory I/O but do not change the $O(n^2)$ algorithmic complexity. The paper suggests that combining its theoretical $n^{1+o(1)}$ algorithm with FlashAttention-style I/O optimization could yield further practical speedups, though implementing this on GPUs is left as future work due to coding challenges (Section B.4).

3. Technical Approach

3.1 Reader Orientation

The paper builds a gradient computation system that takes a multi-layer transformer model and a training loss, and computes approximate gradients for all weight matrices and intermediate variables in almost linear time rather than the standard quadratic time. It solves the backpropagation bottleneck by recognizing that the $n \times n$ attention matrix at each layer can be well-approximated by a low-rank factorization, and that all gradient components appearing in the chain rule can be expressed in forms amenable to this factorization — either as a straightforward matrix-vector product with the low-rank attention matrix, or as a Hadamard product with another low-rank matrix — enabling a cascade of carefully-ordered matrix multiplications that never explicitly construct the full $n \times n$ object.

3.2 Big-Picture Architecture

The system has five major components, organized around the backpropagation computation:

1. Base Transformer Model: An $m$-layer transformer defined as $F_m(X) = g_m \circ \text{Attn}_m \circ g_{m-1} \circ \text{Attn}_{m-1} \circ \dots \circ g_1 \circ \text{Attn}_1 \circ g_0(X)$. Each layer $i$ has: a self-attention module $\text{Attn}_i$ (with weights $W_{Q_i}, W_{K_i}, W_{V_i}$ that produce the attention matrix $f(X)$ and output $s(X)$), and a non-attention component $g_i$ (MLP, layer norm, activation function, residual connection) that the paper models as $g_i(Z) = \varphi(Z \cdot W_g)$ where $\varphi$ is any element-wise activation.

2. Low-Rank Attention Approximator: Uses polynomial kernel methods (from Alman and Aggarwal, 2022) to decompose the $n \times n$ attention matrix $f(X) = D^{-1} \exp(XWX^\top/d)$ into two tall matrices $U_1, V_1 \in \mathbb{R}^{n \times k_1}$ where $k_1 = n^{o(1)}$, such that $\|U_1 V_1^\top - f(X)\|_\infty \leq \epsilon / \text{poly}(n)$. These matrices are constructed in $n^{1+o(1)}$ time and are the foundation for all subsequent acceleration.

3. Gradient Component Decomposer: Takes the closed-form expressions of the gradient on intermediate variables $T_i(X)$ (derived from the chain rule and prior work by Deng et al., 2023) and breaks them into five computable terms — $D_6, D_7, D_8, D_2, D_4$ — each of which can be rewritten in one of two canonical forms amenable to low-rank acceleration: either $z(X) \cdot X \cdot W^\top$ (dot-product form) or $(f(X) \odot (UV^\top)) \cdot H$ (Hadamard-product form).

4. Fast Gradient Computer (SingleGrad): For each layer $i$, given the upstream gradient $\frac{dL}{dT_i(X)}$, this module computes three things: (a) the gradient on the layer input $\frac{dL}{dT_{i-1}(X)}$ (needed to continue backpropagation), (b) the gradient on $W_i = W_{Q_i}W_{K_i}^\top$, and (c) the gradient on $W_{V_i}$. Each computation exploits the low-rank representation of $f(X)$ and the canonical gradient forms, using careful multiplication ordering to keep the cost at $n^{1+o(1)}$.

5. Multi-Layer Backpropagation Loop (MultiGrad): Starting from the loss layer and the gradient $\frac{dL}{dT_m(X)}$ (computed in $O(nd)$ time from the loss definition), this loop iterates backward from layer $m$ down to layer 1. At each step, it calls SingleGrad to get approximate gradients for the weights and the layer input, then passes that input gradient to the previous layer. The algorithm tracks error accumulation and proves that the total approximation error across all $m$ layers is bounded by $1/\text{poly}(n)$.

Information Flow: The loss function produces $\frac{dL}{dT_m(X)}$ → this gradient enters layer $m$'s SingleGrad, which uses the $g_m$ function to compute $G_m = \frac{dL}{d\text{Attn}_m}$ via the chain rule → SingleGrad decomposes the gradient into the five $D_k$ terms, approximates each using low-rank multiplication, and sums them to produce $\frac{dL}{dT_{m-1}(X)}$, $\frac{dL}{dW_m}$, and $\frac{dL}{dW_{V_m}}$ → the weight gradients go to the optimizer, $\frac{dL}{dT_{m-1}(X)}$ goes to layer $m-1$ → repeats until layer 1 → layer 1 additionally produces $\frac{dL}{dX}$ if needed (for prompt tuning or input gradients).

3.3 Roadmap for the Deep Dive

• First, the mathematical foundation: how self-attention is defined, how the chain rule yields closed-form gradient expressions for multi-layer transformers, and how the gradient on intermediate variables decomposes into five terms — because this decomposition drives the entire algorithmic structure.

• Second, the low-rank approximation technique: how the $n \times n$ attention matrix $f(X)$ is approximated by $U_1 V_1^\top$ in $n^{1+o(1)}$ time, why the existence of such an approximation is plausible given that $f(X)$ is constructed from only $nd$ input entries, and the critical "multiplication ordering trick" that avoids ever forming the full $n \times n$ matrix.

• Third, the gradient computation for $T_i(X)$ (the intermediate variables): this is the hardest part because it involves decomposing the raw gradient from Lemma D.1 into the five terms $B_6, B_7, B_8, B_2, B_4$, reducing the double summation over $i_0, j_0$ into matrix products involving auxiliary matrices $z_6, z_7, z_8, z_2, z_4$, and then showing that each $z$-matrix can itself be approximated in low-rank form — enabling fast computation of each $D_k$ term.

• Fourth, the gradient computation for the weight matrices $W_i$ and $W_{V_i}$: these are relatively simpler but require extending prior work from $\ell_2$ loss to general loss functions, and showing that the resulting expressions ($X^\top p(X) X$ for $W_i$ and $X^\top f(X) G_i$ for $W_{V_i}$) are amenable to the low-rank multiplication trick.

• Fifth, the gradient computation through the non-attention components $g_i$: deriving the closed form $G_i = (g_{T_i} \odot \varphi'(\text{Attn}_i W_g)) W_g^\top$ and showing it costs only $O(nd^2)$, which is linear in $n$.

• Sixth, the error propagation analysis across multiple layers: using mathematical induction to prove that if each layer's approximation error is at most $\epsilon$, the total error after $m$ layers remains bounded by $m \cdot \text{poly}(n) \cdot \epsilon$, and choosing $\epsilon = 1/\text{poly}(n)$ to make the final error $1/\text{poly}(n)$.

3.4 Detailed, Sentence-Based Technical Breakdown

This is a theoretical complexity paper whose core contribution is proving that the gradient of a multi-layer transformer can be approximated in almost linear time with bounded error. There is no implementation or empirical evaluation — the "system" is a mathematical proof and an algorithmic template (Algorithm 1). The technical content consists of: (1) deriving closed-form gradient expressions, (2) showing that each expression can be rewritten in a form amenable to low-rank matrix multiplication, (3) providing low-rank approximation lemmas for the attention matrix and its derived quantities, (4) proving running time bounds by counting matrix multiplications with the right ordering, and (5) proving error bounds by tracking the amplification of approximation error through linear algebra operations.

---

Mathematical Foundation: Self-Attention, Loss, and Chain Rule

Self-attention definition. The paper defines the self-attention module $\text{Attn}(X)$ with input $X \in \mathbb{R}^{n \times d}$, where $n$ is the number of tokens and $d$ is the hidden dimension. Given weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d \times d}$, the attention output is:

$\text{Attn}(X) = \text{Softmax}(X W_Q W_K^\top X^\top / d) \cdot X W_V$

For notational convenience, the paper defines the combined key-query weight $W := W_Q W_K^\top \in \mathbb{R}^{d \times d}$. The attention computation then proceeds through three intermediate quantities:

$A := \exp(X W X^\top / d) \in \mathbb{R}^{n \times n}$ $D := \text{diag}(A \mathbf{1}_n) \in \mathbb{R}^{n \times n}$ $f(X) := D^{-1} A \in \mathbb{R}^{n \times n}$

where $A$ is the element-wise exponentiated similarity scores (the "energy" matrix before normalization), $D$ is a diagonal matrix containing the row sums of $A$ (one per token, representing the normalization constant for the softmax), $\mathbf{1}_n$ is the all-ones vector of length $n$, and $f(X)$ is the row-stochastic attention matrix where $f(X)_{i,j}$ is the attention weight that token $i$ assigns to token $j$. The value computation is $h(X) := X W_V \in \mathbb{R}^{n \times d}$, yielding the final output $s(X) := f(X) \cdot h(X) \in \mathbb{R}^{n \times d}$.

What this computes: For each token position $i$, the attention mechanism computes a weighted average of the value vectors of all tokens, where the weights are determined by the softmax-normalized dot-product similarity between the query vector of token $i$ and the key vectors of all tokens. This is the standard scaled dot-product attention from Vaswani et al. (2017), with the scaling factor $d$ in the denominator inside the softmax.

Why this form: The separation into $f(X)$ and $h(X)$ is critical for the subsequent analysis because the quadratic complexity comes entirely from $f(X)$ — an $n \times n$ matrix. The value computation $h(X) = X W_V$ is linear in $n$ (it is just an $n \times d$ matrix multiplied by a $d \times d$ matrix, costing $O(nd^2)$). If $f(X)$ can be approximated by a low-rank factorization, the entire attention computation becomes sub-quadratic.

Multi-layer transformer definition. The paper defines an $m$-layer transformer as a composition of functions:

$F_m(X) := g_m \circ \text{Attn}_m \circ g_{m-1} \circ \text{Attn}_{m-1} \circ \dots \circ g_1 \circ \text{Attn}_1 \circ g_0(X)$

where $g_i$ represents all non-attention components in layer $i$ (layer normalization, MLP, residual connection, dropout, positional encoding, multi-head concatenation). The intermediate variables $T_i(X)$ are defined recursively: $T_0(X) := g_0(X)$, and for $i \in [m]$, $T_i(X) := (g_i \circ \text{Attn}_i)(T_{i-1}(X))$. So $T_i(X) \in \mathbb{R}^{n \times d}$ is the output of the $i$-th transformer layer — the hidden state after processing by layers 0 through $i$.

Loss function definition. The paper defines a general loss function $L(X)$ as the sum of per-element losses:

$L(X) = \sum_{j=1}^{n} \sum_{k=1}^{d} \ell(X)_{j,k}$

where $\ell(X)_{j,k} : \mathbb{R}^{n \times d} \to \mathbb{R}$ is assumed differentiable for each position $(j, k)$. This formulation is general enough to cover cross-entropy loss (as shown in Remark 3.2), which is the standard training objective. For cross-entropy, the output of the final transformer layer passes through an additional linear layer mapping dimension $d$ to vocabulary size $d_{\text{voc}}$, producing $Y_{\text{pred}} \in \mathbb{R}^{n \times d_{\text{voc}}}$, and the loss is the negative log-likelihood summed over all positions and vocabulary items. The key property this definition captures is that $L(X)$ decomposes as a sum over output positions — a property inherited by all gradients via linearity of differentiation.

Chain rule for multi-layer gradients. Lemma 3.4 (formal version C.4) provides the closed-form gradient expressions that are the starting point for the entire paper. Let $G_i \in \mathbb{R}^{n \times d}$ be the "upstream gradient" — the gradient of the loss with respect to the output of the attention module in layer $i$, i.e., $G_i = \frac{dL}{d \text{Attn}_i(T_{i-1}(X))}$. This $G_i$ encapsulates all information from layers $i+1$ through $m$ that is needed to compute gradients in layer $i$.

Then the gradient of the loss with respect to $T_{i-1}(X)$ (the input to layer $i$) is:

$\frac{dL}{dT_{i-1}(X)} = \sum_{i_2=1}^{n} \sum_{j_2=1}^{d} G_i(i_2, j_2) \cdot \frac{d \text{Attn}_i(T_{i-1}(X))_{i_2, j_2}}{d T_{i-1}(X)}$

where $\frac{d \text{Attn}_i(T_{i-1}(X))_{i_2, j_2}}{d T_{i-1}(X)} \in \mathbb{R}^{n \times d}$ is the gradient of the $(i_2, j_2)$-th scalar entry of $\text{Attn}_i$ with respect to the entire input matrix.

Similarly, the gradient with respect to a weight matrix $W_{*i}$ (where $*$ is $Q$, $K$, or $V$) is:

$\frac{dL}{dW_{*i}} = \sum_{i_2=1}^{n} \sum_{j_2=1}^{d} G_i(i_2, j_2) \cdot \frac{d \text{Attn}_i(T_{i-1}(X))_{i_2, j_2}}{d W_{*i}}$

What these expressions compute: These are the standard chain rule for matrix-valued functions composed with a scalar loss. The outer double summation arises because $L$ is a function of all $n \times d$ outputs of $\text{Attn}_i$, and the total derivative is the sum of partial derivatives weighted by the upstream gradient $G_i$. The inner terms $\frac{d \text{Attn}_i(\cdot)_{i_2, j_2}}{d T_{i-1}(X)}$ are the local Jacobians — how each output entry depends on each input entry — evaluated at the current input $T_{i-1}(X)$.

Why this decomposition matters: The paper's task is to compute these expressions for every layer efficiently. The "trick" is that $\frac{d \text{Attn}_i(\cdot)_{i_2, j_2}}{d T_{i-1}(X)}$ is not an arbitrary $n \times d$ matrix — it has specific structure inherited from the attention mechanism (softmax, matrix multiplications). By plugging in the known closed-form gradient of the attention output (from Deng et al., 2023, Lemma D.1), the double sum collapses into a small set of structured matrix operations that can be accelerated.

Role of $G_i$ and $T_{i-1}(X)$. The intermediate variable $T_{i-1}(X)$ serves as the "input" to the $i$-th attention layer when viewed in isolation, and $G_i$ is the "upstream gradient" flowing into that layer from everything above it. The computation $\frac{dL}{dT_{i-1}(X)}$ is the result of "pulling back" the upstream gradient through the attention layer — it becomes the upstream gradient for the next layer down ($i-1$). This recursive structure is why Lemma 3.4 Part 1 is the linchpin of multi-layer backpropagation: it computes the quantity needed by the next iteration of the backward pass.

---

Low-Rank Approximation of the Attention Matrix

The core insight. The attention matrix $f(X) \in \mathbb{R}^{n \times n}$ is computed from $X \in \mathbb{R}^{n \times d}$, which contains only $nd$ entries. When $d = O(\log n)$, this is $n^{1+o(1)}$ — much smaller than the $n^2$ entries in $f(X)$. This information-theoretic argument suggests $f(X)$ should be highly compressible. The paper formalizes this via polynomial kernel approximation.

Polynomial kernel approximation. Lemma C.13 (from Alman and Song, 2023) states: For any accuracy parameter $\epsilon = 1/\text{poly}(n)$, there exists an integer $k_1 = n^{o(1)}$ and two matrices $U_1, V_1 \in \mathbb{R}^{n \times k_1}$ such that:

$\|U_1 V_1^\top - f(X)\|_\infty \leq \epsilon / \text{poly}(n)$

where $\| \cdot \|_\infty$ is the maximum absolute entry-wise error.

What this means operationally: The $n \times n$ attention matrix is approximated as the product of an $n \times k_1$ matrix and a $k_1 \times n$ matrix, where $k_1$ is sub-polynomial in $n$ (formally $n^{o(1)}$, which grows slower than $n^\delta$ for any $\delta > 0$). The construction and the proof that such $U_1, V_1$ exist and can be built in $n^{1+o(1)}$ time comes from the referenced polynomial method — it uses optimal-degree polynomial approximations to the exponential function to construct a low-degree polynomial kernel that approximates the Gaussian kernel, which in turn approximates the softmax attention.

Why the order of multiplication matters. Even with $U_1 V_1^\top$ available, computing the attention output naively as $(U_1 V_1^\top) h(X)$ would cost $O(n^2 k_1)$ because you would first form the $n \times n$ product. The key optimization is to compute $U_1 (V_1^\top h(X))$ instead:

1. Compute $V_1^\top h(X) \in \mathbb{R}^{k_1 \times d}$: cost $O(k_1 n d) = n^{1+o(1)}$
2. Compute $U_1 (V_1^\top h(X)) \in \mathbb{R}^{n \times d}$: cost $O(n k_1 d) = n^{1+o(1)}$

This avoids ever materializing the $n \times n$ matrix. The same trick — associating matrix multiplications so that the small dimension $k_1$ stays on the inside — is applied throughout the paper to every gradient term that involves $f(X)$.

Assumptions enabling the low-rank approximation. The paper assumes $d = O(\log n)$ (so the hidden dimension grows at most logarithmically with sequence length) and that each entry of $X$ and the weight matrices can be represented using $O(\log n)$ bits (so values have bounded precision that grows logarithmically). The first assumption is mild — in practice $d$ is typically 64–128 in older models or 4096–8192 in modern ones, but the theoretical result treats $d$ as asymptotically small relative to $n$. The second assumption ensures that the infinity norms of the matrices are bounded by $\text{poly}(n)$, which is used in the error analysis to bound the amplification of approximation errors through matrix multiplications.

---

Gradient on Intermediate Variables $T_i(X)$: The Core Computation

This is the most complex part of the paper and the main technical contribution beyond prior work.

The starting point: gradient of a single attention output entry. Lemma D.1 (from Deng et al., 2023) gives the partial derivative of one attention output entry $s(X)_{i_0, j_0}$ with respect to one input entry $X_{i_1, j_1}$. There are two cases based on whether the output position $i_0$ equals the input position $i_1$:

For $i_0 = i_1$ (the same token position — this captures the effect of changing token $i_0$'s own embedding on its own attention output):

$\frac{ds(X)_{i_0, j_0}}{dX_{i_0, j_1}} = C_1 + C_2 + C_3 + C_4 + C_5$

For $i_0 \neq i_1$ (different token positions — this captures cross-token effects):

$\frac{ds(X)_{i_0, j_0}}{dX_{i_1, j_1}} = C_6 + C_7 + C_8$

where the $C_k$ terms are defined as:

• $C_1(X) := -s(X)_{i_0, j_0} \cdot f(X)_{i_0, i_0} \cdot \langle W_{j_1,*}, X_{i_0,*}\rangle$
• $C_2(X) := -s(X)_{i_0, j_0} \cdot \langle f(X)_{i_0,*}, X W_{*,j_1}\rangle$
• $C_3(X) := f(X)_{i_0, i_0} \cdot h(X)_{i_0, j_0} \cdot \langle W_{j_1,*}, X_{i_0,*}\rangle$
• $C_4(X) := \langle f(X)_{i_0,*} \odot (X W_{*,j_1}), h(X)_{*,j_0}\rangle$
• $C_5(X) := f(X)_{i_0, i_0} \cdot (W_V)_{j_1, j_0}$
• $C_6(X) := -s(X)_{i_0, j_0} \cdot f(X)_{i_1, i_0} \cdot \langle W_{j_1,*}, X_{i_0,*}\rangle$
• $C_7(X) := f(X)_{i_1, i_0} \cdot h(X)_{i_1, j_0} \cdot \langle W_{j_1,*}, X_{i_0,*}\rangle$
• $C_8(X) := f(X)_{i_1, i_0} \cdot (W_V)_{j_1, j_0}$

What these expressions represent: Each $C_k$ is a scalar capturing one "pathway" through which a change in input $X_{i_1, j_1}$ affects output $s(X)_{i_0, j_0}$. The terms come from differentiating the softmax normalization ($C_1, C_6$ capture the effect through the normalizing denominator $D$; $C_2$ captures the effect through the attention weights between $i_0$ and other tokens), the value computation ($C_3, C_7$ capture the effect through the value vector of the attended token; $C_4$ captures cross-effects through the Hadamard product), and the direct value weight effect ($C_5, C_8$ capture the gradient flowing directly through $W_V$). The asymmetry between $i_0 = i_1$ and $i_0 \neq i_1$ arises because when $i_0 \neq i_1$, token $i_1$ only affects $s(X)_{i_0, j_0}$ through the attention weights $f(X)_{i_0, i_1}$ (and then through the value $h(X)_{i_1, j_0}$), whereas when $i_0 = i_1$, there are additional pathways through the self-attention weight $f(X)_{i_0, i_0}$ and through the normalizer.

From scalar entries to matrix form. The paper's crucial step is converting these per-entry partial derivatives into operations on whole matrices. This is done in Section D, where each $C_k$ term is "assembled" into a matrix $B_k(X) \in \mathbb{R}^{n \times d}$ or $B_k(X) \in \mathbb{R}^d$, and then the double summation over $(i_0, j_0)$ from the chain rule is collapsed.

For example, $C_6(X, i_1, j_1) = -s(X)_{i_0, j_0} \cdot f(X)_{i_1, i_0} \cdot \langle W_{j_1,*}, X_{i_0,*}\rangle$. When we fix $(i_0, j_0)$ and vary $(i_1, j_1)$, this defines a matrix $B_6(X) \in \mathbb{R}^{n \times d}$ where the $(i_1, j_1)$ entry is $C_6(X, i_1, j_1)$. Lemma D.3 shows that this matrix can be written as:

$B_6(X) = -s(X)_{i_0, j_0} \cdot f(X)_{*,i_0} \cdot (W \cdot X_{i_0,*})^\top$

where $f(X)_{*,i_0} \in \mathbb{R}^{n \times 1}$ is the $i_0$-th column of the attention matrix (the attention weights TO token $i_0$), and $(W \cdot X_{i_0,*})^\top \in \mathbb{R}^{1 \times d}$ is the query/key projection of token $i_0$'s embedding. This is an outer product of a column vector and a row vector, producing an $n \times d$ matrix — but crucially, it is rank-1 for a given $(i_0, j_0)$.

Similarly, Lemma D.6 shows that $C_2(X, j_1) = -s(X)_{i_0, j_0} \cdot \langle f(X)_{i_0,*}, X W_{*,j_1}\rangle$ (which depends only on $j_1$, not $i_1$) yields a vector $B_2(X) \in \mathbb{R}^{d}$:

$B_2(X) = -s(X)_{i_0, j_0} \cdot W^\top \cdot X^\top \cdot f(X)_{i_0,*}$

Analogous matrix forms are derived for $B_7, B_8, B_4$ in Lemmas D.4, D.5, D.7.

Reducing the double sum to matrix products. The gradient with respect to the input $X$ involves a double sum over $(i_0, j_0)$ of $G_i(i_0, j_0)$ times the appropriate $B_k$ matrices. This is where the paper's main technical work happens — for each $B_k$ term, the double sum is collapsed into a matrix multiplication involving an auxiliary "z-matrix" that encodes the summation.

Consider the $B_6$ term. Lemma D.10 defines an auxiliary matrix $z_6(X) \in \mathbb{R}^{n \times n}$, where the $i_0$-th column is:

$z_6(X)_{*,i_0} := (G_i(i_0, *) \cdot s(X)_{i_0,*}) \cdot f(X)_{*,i_0}$

Let us parse this: $G_i(i_0, *) \in \mathbb{R}^{d}$ is the $i_0$-th row of the upstream gradient (treated as a column vector), $s(X)_{i_0,*} \in \mathbb{R}^{d}$ is the $i_0$-th row of the attention output, and their dot product $G_i(i_0, *) \cdot s(X)_{i_0,*}$ is a scalar — call it $K_{i_0}$. This scalar is the cumulative sensitivity of the loss to token $i_0$'s entire output vector. Then the $i_0$-th column of $z_6(X)$ is this scalar times the $i_0$-th column of the attention matrix — the attention weights from all tokens TO token $i_0$.

In matrix form, $z_6(X) = f(X) \cdot \text{diag}(K)$, where $K \in \mathbb{R}^n$ is the vector of these scalars $K_{i_0} = G_i(i_0, *) \cdot s(X)_{i_0,*}$. The double-summed $B_6$ term then collapses to:

$D_6 := \sum_{i_0=1}^{n} \sum_{j_0=1}^{d} G_i(i_0, j_0) B_6(X) = -z_6(X) \cdot X \cdot W^\top$

What this achieves: The double sum over $n \times d$ terms, each involving $O(nd)$ operations in a naive implementation, is replaced by: (1) computing the $n$-vector $K$ in $O(nd)$ time, (2) constructing a low-rank approximation of $z_6(X) = f(X) \cdot \text{diag}(K)$ using the already-available low-rank approximation of $f(X)$, and (3) multiplying in the order that keeps the small dimension $k_1$ on the inside. The cost drops from $O(n^2 d)$ to $n^{1+o(1)}$.

The five $D_k$ terms and their z-matrices. The paper performs the same reduction for all five relevant terms. Table 1 of the paper (implicit in Lemma D.17) summarizes:

where $K \in \mathbb{R}^n$ has $K_{i_0} = G_i(i_0, *) \cdot s(X)_{i_0,*}$, and $\odot$ is the element-wise (Hadamard) product.

Distinction between dot-product and Hadamard forms. Terms $D_8, D_2, D_6$ involve straightforward matrix multiplication with $f(X)$ — they are in the "dot-product form" $f(X) \cdot H$ for some $H$. Terms $D_7$ and $D_4$ involve a Hadamard product between $f(X)$ and another low-rank matrix — they are in the "Hadamard form" $(f(X) \odot (UV^\top)) \cdot H$. The acceleration technique differs between these two forms, but both are handled in $n^{1+o(1)}$ time.

Fast computation of each z-matrix. The paper proves that each $z_k(X)$ matrix can itself be approximated by a low-rank factorization $U_k V_k^\top$ with $k = n^{o(1)}$, constructed in $n^{1+o(1)}$ time. The constructions are:

• $z_6(X) = f(X) \cdot \text{diag}(K)$: Take the low-rank approximation $f(X) \approx U_1 V_1^\top$ from Lemma C.13. Set $U_6 = U_1$ and $V_6 = \text{diag}(K) V_1$. Then $U_6 V_6^\top = U_1 V_1^\top \text{diag}(K) \approx f(X) \text{diag}(K) = z_6(X)$. The construction cost is $O(n k_1)$ for the $\text{diag}(K) V_1$ multiplication.

• $z_2(X) = \text{diag}(K) \cdot f(X)$: Take $f(X) \approx U_1 V_1^\top$. Set $U_9 = \text{diag}(K) U_1$ and $V_9 = V_1$. Then $U_9 V_9^\top = \text{diag}(K) U_1 V_1^\top \approx \text{diag}(K) f(X) = z_2(X)$.

• $z_7(X) = f(X) \odot (h(X) G_i^\top)$: This uses the row-wise Kronecker product (Fact C.2). Given $f(X) \approx U_1 V_1^\top$, and letting $h(X)$ and $G_i$ play the role of the second factor, set $U_7 = U_1 \oslash h(X) \in \mathbb{R}^{n \times k_1 d}$ and $V_7 = V_1 \oslash G_i \in \mathbb{R}^{n \times k_1 d}$. Here $U_1 \oslash h(X)$ means: for each row $i$, take the Kronecker product of $U_1(i, *)$ (length $k_1$) and $h(X)(i, *)$ (length $d$), producing a vector of length $k_1 d$. This is constructed in $O(n d k_1) = n^{1+o(1)}$ time.

• $z_4(X) = f(X) \odot (G_i h(X)^\top)$: Similar construction with $U_{10} = U_1 \oslash G_i$ and $V_{10} = V_1 \oslash h(X)$.

Why the Hadamard product requires the row-wise Kronecker product. The identity from Fact C.2 states that $(U_1 V_1^\top) \odot (U_2 V_2^\top) = (U_1 \oslash U_2)(V_1 \oslash V_2)^\top$. This allows the Hadamard product of two low-rank matrices to be expressed as a single low-rank matrix (with rank equal to the product of the individual ranks). Without this identity, the Hadamard product would destroy the low-rank structure. The cost is that the rank increases from $k_1$ to $k_1 d$, but since $d = O(\log n) = n^{o(1)}$, this remains $n^{o(1)}$.

Error bounds for the z-matrix approximations. For each $z_k$, the paper proves:

$\|\tilde{z}_k(X) - z_k(X)\|_\infty \leq \epsilon / \text{poly}(n)$

The proofs use the fact that the infinity norm error of a matrix product is bounded by the product of dimensions times the infinity norms of the factors times the error of the factors. For example, for $z_6$:

$\|U_1 V_1^\top \text{diag}(K) - f(X) \text{diag}(K)\|_\infty \leq n \cdot \|U_1 V_1^\top - f(X)\|_\infty \cdot \|\text{diag}(K)\|_\infty$

Since $\|U_1 V_1^\top - f(X)\|_\infty \leq \epsilon / \text{poly}(n)$ by Lemma C.13, and $\|\text{diag}(K)\|_\infty \leq \text{poly}(n)$ because each entry of $K$ is a bounded dot product, the product error remains $\epsilon / \text{poly}(n)$.

Putting everything together: Lemma E.11 (Fast computation for $\frac{dL}{dT_{i-1}(X)}$). With low-rank approximations of all five $z_k$ matrices (or direct fast computation for $D_8$ which only needs $f(X) \approx U_1 V_1^\top$), the algorithm computes approximate versions $\tilde{D}_6, \tilde{D}_7, \tilde{D}_8, \tilde{D}_2, \tilde{D}_4$ and sums them:

$\frac{dL}{dT_{i-1}(X)} \approx \tilde{gt} := \tilde{D}_6 + \tilde{D}_7 + \tilde{D}_8 + \tilde{D}_2 + \tilde{D}_4$

Running time: Each $\tilde{D}_k$ is computed via the multiplication ordering trick: for $D_6$, compute $V_6^\top X$ (cost $k_6 n d$), then $(V_6^\top X) W^\top$ (cost $k_6 d^2$), then $U_6$ times the result (cost $n k_6 d$). All steps are $n^{1+o(1)}$ because $k_6 = n^{o(1)}$ and $d = O(\log n)$. The summation costs $5 n d = n^{1+o(1)}$.

Error bound: Using triangle inequality:

$\|\tilde{gt} - \frac{dL}{dT_{i-1}(X)}\|_\infty \leq \sum_{k \in \{2,4,6,7,8\}} \|\tilde{D}_k - D_k\|_\infty \leq 5 \cdot \epsilon / \text{poly}(n)$

Choosing $\epsilon = 1/\text{poly}(n)$ makes the final error $1/\text{poly}(n)$.

---

Gradient on Weight Matrices $W_i$ and $W_{V_i}$

Gradient on $W_i = W_{Q_i} W_{K_i}^\top$. The paper extends the result of Alman and Song (2024a) from $\ell_2$ loss to general loss functions. Lemma F.4 shows that the gradient collapses to a simple form:

$\frac{dL}{dW_i} = X^\top \cdot p(X) \cdot X$

where $p(X) \in \mathbb{R}^{n \times n}$ is defined in Definition C.12. For each row $j_0 \in [n]$:

$p(X)_{j_0} := (\text{diag}(f(X)_{j_0}) - f(X)_{j_0} f(X)_{j_0}^\top) \cdot q(X)_{j_0}$

where $q(X) \in \mathbb{R}^{n \times n}$ is defined as $q(X) = G_i \cdot h(X)^\top$ (Definition C.11). This expression comes from differentiating the softmax: the first term $\text{diag}(f(X)_{j_0}) q(X)_{j_0}$ is the "direct" gradient through the softmax probabilities, and the second term $f(X)_{j_0} f(X)_{j_0}^\top q(X)_{j_0}$ is the "indirect" gradient through the normalizing denominator. The result $p(X)$ can be decomposed as $p(X) = p_1(X) - p_2(X)$ where $p_1(X) = f(X) \odot q(X)$ and $p_2(X) = \text{diag}(p_1(X) \cdot \mathbf{1}_n) f(X)$.

Low-rank approximation of $p_1(X)$ and $p_2(X)$. Lemmas C.15 and C.16 (from Alman and Song, 2024a) provide low-rank factorizations $p_1(X) \approx U_3 V_3^\top$ and $p_2(X) \approx U_4 V_4^\top$, constructed in $n^{1+o(1)}$ time.

Fast computation of $X^\top p(X) X$. The computation proceeds as:

1. Compute $X^\top U_3 \in \mathbb{R}^{d \times k_3}$: $O(n d k_3) = n^{1+o(1)}$
2. Compute $(X^\top U_3) V_3^\top \in \mathbb{R}^{d \times n}$: $O(d k_3 n) = n^{1+o(1)}$
3. Compute $(X^\top U_3 V_3^\top) X \in \mathbb{R}^{d \times d}$: $O(d^2 n) = n^{1+o(1)}$

The same three-step multiplication is done for $p_2(X)$, and the results are subtracted. The error bound follows from the same $\epsilon/\text{poly}(n)$ propagation argument.

What this computes: The expression $X^\top p(X) X$ is the gradient of the loss with respect to the combined key-query weight matrix $W = W_Q W_K^\top$. In practice, one would then compute $\frac{dL}{dW_Q} = \frac{dL}{dW} \cdot W_K$ and $\frac{dL}{dW_K} = W_Q^\top \cdot \frac{dL}{dW}$ via the chain rule.

Gradient on $W_{V_i}$. Lemma G.3 shows a much simpler form:

$\frac{dL}{dW_{V_i}} = X^\top \cdot f(X) \cdot G_i$

What this represents: The gradient of the loss with respect to the value weight matrix is the input $X^\top$ times the attention-weighted upstream gradient. Each row of $G_i$ corresponds to one output token's gradient, and $f(X)$ redistributes this gradient back to the input tokens according to the attention weights.

Fast computation: Using $f(X) \approx U_1 V_1^\top$:

1. Compute $X^\top U_1 \in \mathbb{R}^{d \times k_1}$: $O(n d k_1) = n^{1+o(1)}$
2. Compute $(X^\top U_1) V_1^\top \in \mathbb{R}^{d \times n}$: $O(d k_1 n) = n^{1+o(1)}$
3. Compute $(X^\top U_1 V_1^\top) G_i \in \mathbb{R}^{d \times d}$: $O(d^2 n) = n^{1+o(1)}$

The error bound is $\|X^\top \tilde{f}(X) G_i - X^\top f(X) G_i\|_\infty \leq n^2 \|X\|_\infty \| \tilde{f}(X) - f(X) \|_\infty \|G_i\|_\infty \leq \epsilon/\text{poly}(n)$.

---

Gradient Through Non-Attention Components $g_i$

Modeling $g_i$. The paper assumes each $g_i$ takes the form $g_i(Z) = \varphi(Z \cdot W_g)$ where $W_g \in \mathbb{R}^{d \times d}$ is a weight matrix and $\varphi: \mathbb{R} \to \mathbb{R}$ is an element-wise activation function (e.g., ReLU, GELU, SiLU). Let $\varphi'$ denote its derivative. This covers the MLP block (typically two linear layers with an activation) plus any normalization that can be expressed as an element-wise operation.

The gradient of $g_i$. Lemma H.1 computes the Jacobian of $T_i = g_i(\text{Attn}_i(T_{i-1}))$ with respect to its input $\text{Attn}_i$. For a fixed output position $(i_4, j_4)$ and input position $(i_5, j_5)$:

$\frac{d T_i(i_4, j_4)}{d \text{Attn}_i(i_5, j_5)} = \begin{cases} \varphi'(\text{Attn}_i(i_4, *) \cdot W_g(*, j_4)) \cdot W_g(j_5, j_4) & \text{if } i_4 = i_5 \\ 0 & \text{if } i_4 \neq i_5 \end{cases}$

What this means: The non-attention component $g_i$ operates independently on each token position — the output at position $i_4$ depends only on the input at position $i_4$. This is a crucial property: it means the Jacobian is block-diagonal (each $d \times d$ block corresponds to one position), and the full $n d \times n d$ Jacobian can be represented as $n$ independent $d \times d$ operations.

Computing $G_i$ from $\frac{dL}{dT_i(X)}$. Lemma H.2 shows that when backpropagating through $g_i$, the gradient $G_i = \frac{dL}{d \text{Attn}_i}$ can be computed as:

$G_i = (g_{T_i} \odot \varphi'(\text{Attn}_i \cdot W_g)) \cdot W_g^\top$

where $g_{T_i} = \frac{dL}{dT_i(X)}$ is the upstream gradient coming from the layer above. The computation is:

1. Element-wise multiply $g_{T_i}$ by $\varphi'(\text{Attn}_i \cdot W_g)$ (a matrix of activation derivatives): $O(n d)$ time
2. Multiply the result by $W_g^\top \in \mathbb{R}^{d \times d}$: $O(n d^2)$ time

Why this is linear in $n$: The key property is that $g_i$ operates position-wise — there is no interaction between different token positions. This means backpropagation through $g_i$ costs $O(n d^2)$, which is linear in $n$ (since $d = O(\log n)$, $d^2 = (\log n)^2 = n^{o(1)}$). The quadratic bottleneck comes entirely from the attention mechanism.

---

Error Propagation Across Multiple Layers

The single-layer guarantee. Lemma H.3 establishes the base case: for a single-layer transformer $F_1(X) = g_1 \circ \text{Attn}_1 \circ g_0(X)$, the gradient $\frac{dL}{dX}$ can be approximated in $n^{1+o(1)}$ time with error bounded by $1/\text{poly}(n)$. The proof combines all previous lemmas: $G_1$ is computed from $\frac{dL}{dT_1}$ in $n^{1+o(1)}$ time via Lemma H.2, the intermediate gradient $\frac{dL}{dT_0}$ is computed in $n^{1+o(1)}$ time via Lemma E.11, and $\frac{dL}{dX}$ is computed from $\frac{dL}{dT_0}$ via another $g_0$ backpropagation step.

The induction step. Lemma H.4 proves the multi-layer case by mathematical induction. The induction hypothesis: for a $k$-layer transformer, the gradient with respect to the input can be approximated in $n^{1+o(1)}$ time with error $1/\text{poly}(n)$. For a $(k+1)$-layer transformer $F_{k+1}(X) = F_k(T_1)$ where $T_1 = g_1(\text{Attn}_1(g_0(X)))$, the algorithm:

1. Gets $\frac{dL}{dT_1}$ from the induction hypothesis applied to the $k$-layer part (treating $T_1$ as input to $F_k$)
2. Computes $\frac{dL}{dX}$ from $\frac{dL}{dT_1}$ using the single-layer method (Lemmas H.2, E.11, H.2 again)

The running time is $n^{1+o(1)}$ per layer, and there are $k+1 = m$ layers, so total time is $m \cdot n^{1+o(1)} = n^{1+o(1)}$ since $m = n^{o(1)}$.

Error bound propagation. The critical part is proving that errors do not compound catastrophically across layers. The paper shows:

$\|\tilde{g}_{T_1} - g_{T_1}\|_\infty \leq 1/\text{poly}(n) \quad \text{(by induction hypothesis)}$

$\|\tilde{g}_X - g_X\|_\infty \leq n \cdot d \cdot \|\tilde{g}_{T_1} - g_{T_1}\|_\infty \cdot \| \frac{d T_1}{dX} \|_\infty \leq \text{poly}(n) \cdot (1/\text{poly}(n)) \cdot \text{poly}(n) = 1/\text{poly}(n)$

The multiplication by $n \cdot d$ and $\|\frac{d T_1}{dX}\|_\infty$ comes from the matrix product in the chain rule — these factors are bounded by $\text{poly}(n)$ because each matrix entry uses $O(\log n)$ bits. The crucial point is that the error grows by only a $\text{poly}(n)$ factor per layer, not exponentially. After $m$ layers, the error is $m \cdot \text{poly}(n) \cdot (1/\text{poly}(n)) = 1/\text{poly}(n)$, since we can choose the per-layer $\epsilon$ to be $1/m \cdot \text{poly}(n) = 1/\text{poly}(n)$.

Why this works: The key is that the chain rule involves only linear operations (matrix multiplications and additions) between the per-layer gradient computations. Linear operations amplify errors at most multiplicatively by norms of the involved matrices, which are bounded. There is no exponential accumulation because there are no iterative or feedback processes — backpropagation is a single backward pass through a fixed-depth network. With $m = n^{o(1)}$, the depth is asymptotically negligible compared to the polynomial error bound.

---

Handling Practical Sub-Modules

Causal attention mask (Section I). The causal mask $M \in \{0, 1\}^{n \times n}$ (lower triangular, with $M_{i,j} = 1$ if $i \geq j$, else 0) modifies the attention matrix to $\tilde{f}(X) := D^{-1}(M \odot A)$ where $D := \text{diag}((M \odot A) \cdot \mathbf{1}_n)$. The challenge is that $M \odot (U_1 V_1^\top)$ is not itself low-rank — the mask breaks the factorization. The paper's solution uses Algorithm 2 from prior work (Liang et al., 2024d): for any vector $v \in \mathbb{R}^n$, the product $(M \odot (U_1 V_1^\top)) \cdot v$ can be computed in $O(n k_1)$ time by maintaining a cumulative sum. The algorithm is:

1. Compute $b_j = V_1(j, *) \cdot v_j$ for each $j$ (a $k_1$-vector)
2. Maintain $c_j = c_{j-1} + b_j$ (cumulative sum of $b$ vectors)
3. Output $Y_j = \langle U_1(j, *), c_j \rangle$ for each $j$

This works because the lower-triangular mask means position $j$ can only attend to positions $\leq j$, and the cumulative sum $c_j$ exactly captures $\sum_{i=1}^{j} V_1(i, *) \cdot v_i$. Lemma I.4 extends this to matrices $H$ with $k_H$ columns by running the algorithm $k_H$ times, costing $O(n k_1 k_H) = n^{1+o(1)}$.

All gradient terms are then categorized into two patterns: dot-product terms $\tilde{f}(X) H$ (handled by substituting the masked multiplication) and Hadamard-product terms $(\tilde{f}(X) \odot (U V^\top)) H$ (handled via the same algorithm after rewriting using Fact C.2 to combine the mask with the low-rank factors).

Residual connection (Section J). The paper defines two residual connections per layer: $Z_i = T_{i-1} + \text{Attn}_i(T_{i-1})$ (around the attention) and $T_i = Z_i + g_i(Z_i)$ (around the MLP). Lemma J.3 proves that if a module's gradient can be computed in $n^{1+o(1)}$ time without residual connection, then adding the residual connection increases the cost by only $O(nd)$ (for the addition of the identity gradient) and does not affect the error bound beyond a constant factor. The chain rule gives $\frac{dL}{dX_R} = \frac{dL}{dY_R} \cdot (I + \frac{dH}{dX_R}) = \frac{dL}{dY_R} + \frac{dL}{dY_R} \cdot \frac{dH}{dX_R}$. The first term is free (it is just the upstream gradient), and the second term is the gradient through the sub-module, which is already fast by assumption.

Multi-head attention (Section K). With $h$ heads, the hidden dimension $d$ is split into $h$ sub-dimensions $d_h = d/h$. Each head $l$ computes $\text{Attn}_l$ independently on its slice of the input. The gradient with respect to the input is the sum over heads: $\frac{dL}{dX} = \frac{dL}{dY} \cdot \sum_{l=1}^{h} \frac{d s_l}{dX}$. Since $h$ is constant and each head's gradient uses the single-head fast algorithm (with dimension $d_h$ instead of $d$), the total cost is $h \cdot n^{1+o(1)} = n^{1+o(1)}$. The error sums over $h$ heads but remains bounded by $h \cdot 1/\text{poly}(n) = 1/\text{poly}(n)$.

Prompt tuning (Section B.5). Because the algorithm computes gradients with respect to the input $X$ (through the entire chain rule back to $T_0 = g_0(X)$), it naturally supports prompt tuning — the computation of $\frac{dL}{dX}$ for the soft prompt embeddings. The gradient on $X$ is obtained as the final output of the MultiGrad loop (line 32 in Algorithm 1, which implicitly computes $\frac{dL}{dX}$ through the $g_0$ backpropagation if needed).

---

Summary of Design Choices and Their Justifications

• Low-rank approximation of $f(X)$ via polynomial methods over alternatives (e.g., random projections, kernel density estimation): provides provable entry-wise error bounds ($\|\cdot\|_\infty$) of $1/\text{poly}(n)$, which is necessary for the error propagation analysis — other methods typically provide weaker guarantees (Frobenius norm or probabilistic bounds).

• Decomposition of the gradient into five $D_k$ terms over a monolithic computation: isolates the terms into two structural patterns (dot-product vs. Hadamard) that each require different low-rank handling, but collectively cover all pathways in the chain rule. The five-term decomposition comes from differentiating the softmax normalization, the value computation, and the cross-attention effects separately.

• Multiplication ordering trick ($U_1 (V_1^\top H)$ instead of $(U_1 V_1^\top) H$): the single algorithmic device that makes all fast computations work. By keeping the $k_1 \times n$ matrix on the inside of the multiplication, the cost goes from $O(n^2)$ to $O(n k_1 d) = n^{1+o(1)}$.

• Row-wise Kronecker product for Hadamard low-rank products (Fact C.2): converts a Hadamard product of low-rank matrices (which is not obviously low-rank) into a single low-rank matrix with composite rank $k_1 d$. This is essential for the $D_7$ and $D_4$ terms, which involve element-wise products between $f(X)$ and gradient-related matrices.

• Mathematical induction for multi-layer error propagation over a probabilistic union bound or martingale argument: linear structure of the chain rule means errors compound additively (through matrix norms), so $m$ layers give at most $m \cdot \text{poly}(n)$ error amplification — controlled by choosing per-layer $\epsilon$ appropriately.

• Assumption $d = O(\log n)$ for theoretical tractability: ensures that the rank parameters ($k_1 = n^{o(1)}$, $k_1 d = n^{o(1)}$, etc.) remain sub-polynomial, and that $d^2 = (\log n)^2 = n^{o(1)}$ so that the $g_i$ backpropagation cost is sub-quadratic. The paper acknowledges this is mild — in practice $d$ is large but constant relative to $n$, so the $O(n d^2)$ cost is practically linear in $n$ even if not asymptotically $n^{1+o(1)}$ under this specific definition.

• Modeling $g_i$ as $\varphi(Z W_g)$ over arbitrary architectures: captures the essential structure — position-wise operations with a learned linear transformation — while being general enough to cover MLPs, activations, and layer norms (which can be folded into $\varphi$). The position-wise property is what makes the gradient through $g_i$ linear in $n$.

4. Key Insights and Innovations

Innovation 1: A Comprehensive Closed-Form Decomposition of Multi-Layer Transformer Gradients Into Two Structural Patterns

The paper's deepest conceptual contribution is not the acceleration itself — it is the taxonomy of gradient structure that makes acceleration possible. Prior work on attention gradients (Deng et al., 2023; Alman and Song, 2024a) treated each gradient term as an isolated algebraic expression, computing them one by one without recognizing higher-level patterns. This paper performs a systematic structural analysis of the full backpropagation chain and discovers that every gradient term across intermediate variables $T_i(X)$, key-query weights $W_i$, and value weights $W_{V_i}$ falls into exactly one of two canonical forms: either a dot-product form $f(X) \cdot H$ (where $f(X)$ is the attention matrix and $H$ is some matrix derived from upstream gradients and hidden states) or a Hadamard-product form $(f(X) \odot (U V^\top)) \cdot H$.

This is a fundamental shift in how to think about transformer gradients. Before this work, the field viewed the gradient computation as a heterogeneous collection of terms — $C_1$ through $C_8$ in Lemma D.1, each with different algebraic structure — and the dominant assumption was that accelerating one term did not imply anything about the others. Alman and Song (2024a) only handled a subset of terms for a single layer with a specific loss. This paper's taxonomy shows that the apparent diversity is superficial: after collapsing the double summations from the chain rule into matrix products (via the $z$-matrices in Lemma D.17), every term reduces to either multiplying a low-rank-approximable matrix by something, or taking a Hadamard product with another low-rank matrix and then multiplying by something. The five $D_k$ terms — which look like five different problems — turn out to be two.

The significance goes beyond this particular algorithm. The taxonomy provides a diagnostic tool for future work: any proposed attention variant or new gradient-based training method can be checked against these two patterns. If its gradients fall into one of these categories, the low-rank acceleration machinery applies immediately. If not, one has identified a genuinely new computational challenge. This reframes the problem from "can we accelerate this specific expression" to "what structural properties guarantee that gradient computation is sub-quadratic." It is analogous to how the discovery that many graph algorithms reduce to matrix multiplication reframed combinatorial optimization — it is a unifying abstraction that changes what questions researchers ask.

The evidence for this taxonomy is distributed across the paper but anchored in Lemma D.17, which explicitly classifies each $D_k$ term into one of the two forms, and in Sections E–G, which provide separate acceleration lemmas for each pattern. The pattern holds across all three gradient components ($T_i$, $W_i$, $W_{V_i}$) and across all layers, which is why the same low-rank multiplication trick works uniformly.

Innovation 2: The First Proof That Multi-Layer Error Propagation Remains Controlled Under Low-Rank Approximation

Accelerating a single attention layer with low-rank approximations was known (Alman and Song, 2024a). The critical unsolved problem was: does the approximation error compound catastrophically when you chain $m$ such accelerated layers together through backpropagation? This is not an incremental extension — it is the difference between a theoretical curiosity (a single-layer result with no path to practice) and a result that applies to real deep transformers.

The field's prior implicit assumption was pessimistic. In general, composing $m$ approximate functions can cause errors to grow exponentially — each layer's output error becomes the next layer's input error, and if the function is sensitive to its inputs, the error amplifies geometrically. If this happened in transformers, the low-rank approach would be useless for any $m > 1$, because the final gradient would be dominated by noise after just a few layers. The paper proves this does NOT happen. The error propagation analysis (Lemma H.4, via mathematical induction) shows that each layer amplifies the incoming error by at most a $\text{poly}(n)$ factor — specifically, the product of matrix dimensions ($nd$) and the infinity norms of the gradient matrices. Since these norms are bounded because each entry uses $O(\log n)$ bits, the error after $m$ layers is at most $m \cdot \text{poly}(n) \cdot \epsilon$. By choosing the per-layer approximation tolerance $\epsilon$ to be $1/(m \cdot \text{poly}(n))$, the final error is $1/\text{poly}(n)$.

This is a fundamental theoretical advance in understanding how approximation errors flow through deep architectures during training. It provides a template for analyzing other approximate training methods: check whether the chain rule involves only linear operations (matrix multiplications, additions) between layers, and whether the intermediate matrices have bounded norms. If so, error accumulation is additive, not multiplicative. The proof technique itself — bounding the infinity norm of matrix products and using induction over the layer index — is novel in this context and may generalize.

The evidence is Theorem 4.2 and Lemma H.4. Without this error propagation result, the single-layer acceleration would remain a standalone theoretical result, like many in the efficient attention literature. With it, the paper makes a claim about training full transformer models — a qualitatively different kind of statement. The key numbers: the error bound is $1/\text{poly}(n)$ for any $m = n^{o(1)}$, meaning the depth can grow sub-polynomially with sequence length while maintaining accuracy. For practical transformers where $m$ is 12–96, this is more than sufficient.

Innovation 3: A Loss-Function-Agnostic Framework That Reveals Cross-Entropy Training Is As Tractable As $\ell_2$

Alman and Song (2024a) proved fast gradient computation for a single layer with $\ell_2$ loss. At first glance, this seems like a minor restriction — just swap the loss function, right? But the gradient expressions for different loss functions can be structurally different. The $\ell_2$ loss produces a particularly simple upstream gradient $G_i$ (essentially the prediction error), and it was unclear whether general loss functions would introduce terms that break the low-rank structure.

The paper's innovation is a general loss function framework (Definition 3.1) that abstracts away the specific form of the loss while preserving the structural properties needed for acceleration. The loss is defined as $L(X) = \sum_{j=1}^{n} \sum_{k=1}^{d} \ell(X)_{j,k}$, where the per-element loss $\ell(X)_{j,k}$ can be anything differentiable. The key is that the loss decomposes as a sum over output positions — a property that cross-entropy, $\ell_2$, and most practical training objectives share. This sum structure propagates through the chain rule (Lemma 3.4), ensuring that the upstream gradient $G_i$ always enters the gradient computations through the same summation pattern, regardless of how $\ell$ is defined.

This reframes the problem from "can we accelerate training with loss $L$" to "can we accelerate training with any loss that decomposes as a sum over positions." The answer is yes — and Remark 3.2 explicitly verifies that standard cross-entropy loss satisfies this property. This is significant because it means the theoretical guarantee covers the actual training objective used by virtually all language models, not a toy proxy.

The framework also clarifies why prior work's $\ell_2$-specific derivations could not be directly reused: they exploited the specific form of the gradient of $\ell_2$ loss to simplify intermediate expressions. By working at the level of the abstract $G_i$ (the upstream gradient matrix, whatever its origin), the paper separates the acceleration problem into two independent parts: (1) compute $G_i$ from the loss (which costs $O(nd)$ for any sum-decomposable loss), and (2) push $G_i$ through the attention layer using the low-rank machinery (which works for any $G_i$). This separation is clean and general.

The evidence is in Lemma 3.4 (which expresses gradients in terms of $G_i$ without specifying the loss) and Remark 3.2 (which maps cross-entropy into the framework). The framework is validated by the fact that Lemma F.4 successfully extends the gradient of $W_i$ from $\ell_2$ to general loss — the previous result by Alman and Song becomes a special case.

Innovation 4: Identification and Resolution of the "Causal Mask Breaks Low-Rank" Problem

The causal attention mask — a lower-triangular matrix $M \in \{0, 1\}^{n \times n}$ that prevents tokens from attending to future positions — is universal in autoregressive language model training. At first encounter, it appears to be a fatal obstacle to low-rank methods. The mask itself is full-rank (it is triangular with ones on the diagonal, so its determinant is 1), and the masked attention matrix $\tilde{f}(X) = D^{-1} (M \odot A)$ inherits this full-rank property. If $\tilde{f}(X)$ cannot be well-approximated by a low-rank factorization, the entire acceleration approach collapses for decoder-only transformers — which is to say, for GPT-style models that dominate the field.

Prior work on fast attention (Alman and Song, 2023; Alman and Song, 2024a) either ignored the causal mask or treated it as an orthogonal concern. The implicit assumption in the field was that handling causality required either a different algorithmic approach or accepting the quadratic cost during training. The paper directly confronts this assumption and shows it is false.

The innovation is the recognition that the low-rank approximation of the unmasked attention matrix $A \approx U_0 V_0^\top$ remains useful even after the mask is applied, because the operation $(M \odot (U_0 V_0^\top)) \cdot v$ for any vector $v$ can be computed in $O(n k_1)$ time by exploiting the lower-triangular structure of $M$. This is a specific algorithmic insight: the cumulative-sum trick (Algorithm 2) works because the $j$-th output position can only attend to positions $1$ through $j$, and maintaining a running sum of $V_0(i, *) \cdot v_i$ for $i \leq j$ avoids recomputing the full triangular product.

What makes this an innovation rather than just an implementation detail is that it classifies all gradient terms into two categories (dot-product and Hadamard-product) and provides separate near-linear-time algorithms for each category with the causal mask applied (Lemmas I.7 and I.8). This classification is not obvious from the raw gradient expressions — it requires the structural analysis of Section D (which decomposes the gradient into the $D_k$ terms) to reveal that every term falls into one of these two patterns. The paper thus solves the causal mask problem not by modifying the mask or the attention mechanism, but by recognizing that the gradient expressions, when properly organized, interact with the mask in only two ways, both of which are amenable to the cumulative-sum algorithm.

The evidence is in Section I and Lemmas I.7–I.8, which explicitly show that after substituting $\tilde{f}(X)$ for $f(X)$ and using the masked multiplication algorithm, all gradient terms compute in $n^{1+o(1)}$ time. The significance is that the main theorem (Theorem 4.2) applies to causal autoregressive transformers — the architecture of GPT, LLaMA, and most deployed LLMs — not just to bidirectional encoders.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper does not report experiments on any empirical dataset. There are no train/test splits, no benchmark tasks, no input sequences, and no loss values reported. All results are theoretical: running time bounds and approximation error bounds stated as functions of the sequence length $n$ and hidden dimension $d$. The paper is a pure computational complexity theory paper — it proves asymptotic time complexity and error bounds for an algorithm, but does not implement or evaluate that algorithm on any data.

• Base model(s). No specific model is evaluated. The theoretical framework assumes a generic multi-layer transformer with $m$ layers, hidden dimension $d$, and sequence length $n$, with weight matrices $W_{Q_i}, W_{K_i}, W_{V_i} \in \mathbb{R}^{d \times d}$ and non-attention components $g_i$ modeled as $g_i(Z) = \varphi(Z W_g)$ for element-wise activation $\varphi$. The assumptions are: (1) $d = O(\log n)$, (2) each matrix entry can be represented using $O(\log n)$ bits, and (3) $m = n^{o(1)}$. No specific architecture (e.g., GPT, LLaMA, BERT) is instantiated or tested.

• Metrics. The paper reports two types of theoretical metrics:

  • Running time complexity: Measured in asymptotic $O(\cdot)$ notation as a function of $n$ and $d$. The central claim is that the gradient computation runs in $n^{1+o(1)}$ time, compared to the standard $\Omega(n^2)$ lower bound for naive exact computation. The $n^{1+o(1)}$ notation means: for any $\delta > 0$, the running time is $O(n^{1+\delta})$ for sufficiently large $n$. The $n^{o(1)}$ terms represent factors that grow slower than any polynomial (e.g., $(\log n)^c$, $2^{\sqrt{\log n}}$, etc.).
  • Approximation error: Measured as the entry-wise $\ell_\infty$ norm of the difference between the approximate gradient and the exact gradient: $\|\tilde{g} - g\|_\infty \leq 1/\text{poly}(n)$. This is a worst-case bound — every single entry of every gradient matrix (for every layer, every weight, and every intermediate variable) is guaranteed to have error at most $1/\text{poly}(n)$. The notation $1/\text{poly}(n)$ means: for any polynomial $P(n)$, there exists a choice of algorithm parameters such that the error is $\leq 1/P(n)$ for all sufficiently large $n$.

• Baselines. There are no empirical baselines. The theoretical baseline is the standard exact gradient computation via backpropagation, which the paper states has time complexity $\Omega(n^2 d) = \Omega(n^2)$ for the attention mechanism (since computing and storing the $n \times n$ attention matrix and propagating gradients through it requires at least quadratic time). The paper's contribution is relative to this theoretical baseline — it provides an algorithm with provably better asymptotic complexity and bounded error.

• Generation budget / compute accounting. The paper measures compute in abstract terms of matrix multiplication operations. The fundamental unit of accounting is: "how many $n \times n$ matrix multiplications are required, and can they be restructured to avoid forming the full $n \times n$ matrix?" The cost model counts operations at the granularity of basic linear algebra: multiplying an $n \times k$ matrix by a $k \times d$ matrix costs $O(n k d)$, and the key parameter is the rank $k$ of the low-rank approximations (always $k = n^{o(1)}$). The paper does not measure FLOPs, wall-clock time, memory usage, or GPU utilization.

• Cross-validation / statistical protocol. None. There is no data, no randomness in the algorithm (the polynomial kernel approximation is deterministic for a given accuracy parameter $\epsilon$), and no statistical inference. The error bounds are worst-case deterministic guarantees — they hold for every possible input matrix $X$ satisfying the bounded-entry assumptions, not "with high probability" or "on average."

---

Main Quantitative Results

There are no empirical quantitative results in the traditional sense — no accuracy percentages, no FLOP comparisons, no tables with numbers from runs. The "results" are theorem statements that specify asymptotic complexity and error bounds. I organize them by the paper's logical groupings:

Single-Layer Gradient Approximation (Theorem 4.1)

The theorem states: For a single-layer self-attention transformer, the three gradient components — $\frac{dL}{dX}$, $\frac{dL}{dW_Q W_K^\top}$, and $\frac{dL}{dW_V}$ — can be approximated in $n^{1+o(1)}$ time with error bounded by $1/\text{poly}(n)$.

This is proven by combining three component lemmas:

1. Lemma 5.1 (informal version E.11): $\frac{dL}{dT_{i-1}(X)}$ is computable in $n^{1+o(1)}$ time with $1/\text{poly}(n)$ error. The proof decomposes the gradient into five $D_k$ terms ($D_6, D_7, D_8, D_2, D_4$), provides low-rank approximations for each via the $z$-matrices in Lemma D.17, and sums them.
2. Lemma 5.2 (informal version F.5): $\frac{dL}{dW_i}$ for $W_i = W_{Q_i}W_{K_i}^\top$ is computable in $n^{1+o(1)}$ time with $1/\text{poly}(n)$ error. The proof uses the form $X^\top p(X) X$ with low-rank approximations of $p_1(X)$ and $p_2(X)$ from Lemmas C.15 and C.16.
3. Lemma 5.3 (informal version G.4): $\frac{dL}{dW_{V_i}}$ is computable in $n^{1+o(1)}$ time with $1/\text{poly}(n)$ error. The proof uses the form $X^\top f(X) G_i$ with the low-rank approximation of $f(X)$ from Lemma C.13.

What is being compared: The standard exact computation of these gradients (which involves forming and differentiating through the $n \times n$ attention matrix) would require $\Omega(n^2)$ time in the worst case — you must at minimum compute the softmax normalization for all $n^2$ pairs. The theorem asserts that by accepting an error of $1/\text{poly}(n)$ per entry (which vanishes faster than any inverse polynomial as $n$ grows), the time drops to $n^{1+o(1)}$ — a qualitative asymptotic improvement.

Multi-Layer Transformer Gradient Approximation (Theorem 4.2, the Main Result)

The theorem states: For an $m$-layer transformer with $m = n^{o(1)}$, Algorithm 1 computes all gradient components for all layers in $m \cdot n^{1+o(1)} = n^{1+o(1)}$ total time, and the approximation error across the entire model is bounded by $1/\text{poly}(n)$.

The proof (Lemma H.4) uses mathematical induction over the number of layers. The base case (Lemma H.3) establishes the result for a single-layer transformer. The induction step shows that if a $k$-layer transformer has gradients computable in $n^{1+o(1)}$ time with error $1/\text{poly}(n)$, then a $(k+1)$-layer transformer does as well. The key insight is that the error growth per layer is multiplicative by at most $\text{poly}(n)$ (from the matrix norms in the chain rule), not exponential. With $m = n^{o(1)}$ layers, the total error amplification is $m \cdot \text{poly}(n) = \text{poly}(n)$, so setting the per-layer error tolerance to $1/\text{poly}(n)$ keeps the final error at $1/\text{poly}(n)$.

What this means in absolute terms: The standard backpropagation for an $m$-layer transformer with exact gradient computation runs in $\Omega(m n^2)$ time — $m$ layers, each paying the quadratic attention cost. Theorem 4.2 says this drops to $n^{1+o(1)}$ total, across all layers, while maintaining a gradient approximation that is correct to within an exponentially small (in $\log n$) entry-wise error. The $m$ factor disappears into the $n^{o(1)}$ because $m$ is assumed sub-polynomial in $n$. For a concrete sense: if $n = 10^5$ and $m = 100$, the exact method requires $\sim 10^{12}$ operations per layer ($n^2 d$ with $d \approx \log n \approx 12$), while the approximate method requires $\sim 10^{5.01}$ operations total. The gap grows quadratically with sequence length.

Computational Complexity of Sub-Components (Supporting Lemmas)

The paper proves $n^{1+o(1)}$ time bounds for several key sub-computations, each with $1/\text{poly}(n)$ error:

• Low-rank approximation of $f(X)$ (Lemma C.13): $U_1, V_1 \in \mathbb{R}^{n \times k_1}$ with $k_1 = n^{o(1)}$ are constructed in $n^{1+o(1)}$ time such that $\|U_1 V_1^\top - f(X)\|_\infty \leq \epsilon / \text{poly}(n)$. This is the foundational lemma from prior work (Alman and Song, 2023) that enables all subsequent acceleration.

• Gradient through $g_i$ (Lemma H.2): Given the upstream gradient $\frac{dL}{dT_i}$, computing $G_i = \frac{dL}{d\text{Attn}_i}$ costs $O(n d^2) = n^{1+o(1)}$ time. This is linear in $n$ because $g_i$ operates position-wise — no cross-position interactions.

• The $z$-matrix constructions (Lemmas E.2, E.4, E.7, E.9): Each auxiliary matrix $z_6, z_7, z_2, z_4$ can be low-rank-approximated by $U_k V_k^\top$ with $k = n^{o(1)}$ in $n^{1+o(1)}$ time. The constructions differ: $z_6$ and $z_2$ use straightforward multiplication with a diagonal matrix, while $z_7$ and $z_4$ use the row-wise Kronecker product (Fact C.2) to handle the Hadamard product structure.

• Causal mask operations (Lemmas I.3, I.4, I.7, I.8): Computing $(M \odot (U_0 V_0^\top)) \cdot H$ for the masked attention takes $O(n k_1 k_H) = n^{1+o(1)}$ time for any $H$ with $k_H = n^{o(1)}$ columns. This enables all gradient terms with the causal mask to remain sub-quadratic.

• Residual connection (Lemma J.3): Adding the residual connection increases the per-layer cost by $O(nd)$ — just the cost of adding the identity gradient to the gradient through the sub-module. Error amplification is at most constant.

• Multi-head attention (Lemma K.2): With $h$ heads, the cost is $h \cdot n^{1+o(1)} = n^{1+o(1)}$ since $h$ is constant. The error sums over $h$ heads but remains bounded by $h \cdot 1/\text{poly}(n) = 1/\text{poly}(n)$.

Error Propagation Bounds

The paper proves specific error amplification factors:

• For a matrix product $A B$ where $\tilde{A}$ approximates $A$: $\|\tilde{A} B - A B\|_\infty \leq n \cdot \|\tilde{A} - A\|_\infty \cdot \|B\|_\infty$ (used throughout, e.g., Lemma E.3, E.5).
• For the chain rule across one layer boundary: the error in $\frac{dL}{dT_{i-1}}$ relative to the error in $\frac{dL}{dT_i}$ is amplified by at most $n \cdot d \cdot \|\frac{d T_i}{d \text{Attn}_i}\|_\infty \cdot \|\frac{d \text{Attn}_i}{d T_{i-1}}\|_\infty = \text{poly}(n)$ (Lemma H.4, induction step).
• After $m$ layers: total error amplification is $m \cdot \text{poly}(n)$, which is still $\text{poly}(n)$ because $m = n^{o(1)}$. The final error is $1/\text{poly}(n)$ by choosing per-layer $\epsilon = 1/(m \cdot \text{poly}(n)) = 1/\text{poly}(n)$.

---

Ablation Studies and Robustness Checks

There are no ablation studies in the traditional empirical sense — no removal of components to measure their contribution to accuracy. However, the paper's theoretical structure provides several forms of analysis that serve an analogous role by examining the conditions under which the theoretical guarantees hold and the tightness of the assumptions:

• Dependence on $d = O(\log n)$ assumption: The $n^{1+o(1)}$ running time critically depends on the hidden dimension $d$ being at most logarithmic in $n$. If $d$ were larger — say $d = \Theta(n^\delta)$ — then the rank parameters $k_1 d$ in the Hadamard product low-rank approximations (for $z_7$ and $z_4$) could become polynomial in $n$, pushing the total time above $n^{1+o(1)}$. The paper does not analyze what happens when $d$ is large constant (e.g., $d = 4096$ in practice), but the results would still be meaningful: the running time would be $O(n \cdot \text{poly}(d))$, which is linear in $n$ with a potentially large constant factor from $d$, rather than truly $n^{1+o(1)}$. The paper does not explore this regime or provide bounds as a function of both $n$ and $d$.

• Dependence on $O(\log n)$ bit precision assumption: The bounded-entry assumption (that each matrix entry can be represented with $O(\log n)$ bits) ensures that $\|X\|_\infty$, $\|W\|_\infty$, $\|G_i\|_\infty$, and other matrix norms are bounded by $\text{poly}(n)$. This is used pervasively in the error bounds — every time the paper bounds $\|\tilde{A} B - A B\|_\infty \leq n \|\tilde{A} - A\|_\infty \|B\|_\infty$, the $\|B\|_\infty$ factor is bounded by $\text{poly}(n)$ because of this assumption. If matrix entries could be arbitrarily large (e.g., due to poor initialization or training dynamics), the error amplification could be larger, potentially requiring tighter per-operation tolerance. The paper does not discuss what norms arise in practice during training or whether the assumption is realistic for deep networks.

• Dependence on $m = n^{o(1)}$ assumption: The error propagation proof requires $m$ (number of layers) to be sub-polynomial in $n$ so that the $m \cdot \text{poly}(n)$ error amplification factor remains $\text{poly}(n)$. For practical transformers where $m$ is 12–96 and $n$ is $10^4$–$10^5$, this holds easily ($m$ is much smaller than any polynomial in $n$). But the analysis does not address what happens with very deep transformers ($m = \Theta(n)$), which would require a different error propagation argument.

• Exact vs. approximate $g_i$ backpropagation: The paper models $g_i$ exactly — backpropagation through the MLP and activation function is not approximated, only the attention mechanism is. Lemma H.2 gives the exact closed form $G_i = (g_{T_i} \odot \varphi'(\text{Attn}_i W_g)) W_g^\top$ and computes it directly in $O(n d^2)$ time. This means the approximation error comes entirely from the attention mechanism — the low-rank approximations of $f(X)$ and related matrices. The paper does not explore whether approximating $g_i$ as well (e.g., via low-rank approximations of the MLP weight matrices) would further reduce time or alternatively compound errors.

• Loss function generality: The paper claims the algorithm works for "general loss functions" (Definition 3.1), which requires only that the loss decomposes as $L(X) = \sum_{j,k} \ell(X)_{j,k}$ and $\ell$ is differentiable. The paper verifies this covers cross-entropy in Remark 3.2. However, the paper does not analyze whether specific loss functions could enable tighter constants or better error bounds, or whether some loss functions (with pathological derivatives) could cause the infinity norms in the error analysis to blow up. The $1/\text{poly}(n)$ error bound is worst-case over all loss functions satisfying the differentiability condition.

• Negative result: rank increase from Hadamard products: An implicit negative finding is that the Hadamard product terms ($D_7$ and $D_4$) require rank $k_1 d$ rather than just $k_1$, due to the row-wise Kronecker product construction (Lemma E.4: $U_7 = U_1 \oslash h(X)$ with shape $n \times k_1 d$). This means the constant factor in the running time for these terms is larger than for the dot-product terms by roughly a factor of $d$. The paper does not quantify this overhead or explore whether alternative constructions could achieve lower rank for the Hadamard product terms.

• No analysis of the polynomial degree in the kernel approximation: The low-rank approximation in Lemma C.13 comes from polynomial kernel methods where the degree $g$ of the approximating polynomial determines the rank $r \leq \binom{2(g+d)}{2g}$. The paper states $g = O(\max\{\log(1/\epsilon)/\log(\log(1/\epsilon)/R), R^2\})$ (Lemma I.1) but does not analyze how this degree — and hence the constant in $k_1 = n^{o(1)}$ — scales with the desired accuracy or the norm bound $R$. For $1/\text{poly}(n)$ error, $g$ may need to grow with $n$, affecting the practical constant. This is not explored.

---

Critical Assessment

This section evaluates whether the paper's theoretical framework and proof structure genuinely support its central claims, and identifies what is and is not established.

Central Claim: "The gradients of a multi-layer transformer can be approximated in $n^{1+o(1)}$ time with $1/\text{poly}(n)$ error."

What the proof actually establishes: The proof shows that there EXISTS an algorithm (Algorithm 1) with asymptotic running time $n^{1+o(1)}$ and approximation error $1/\text{poly}(n)$, under the assumptions $d = O(\log n)$, $m = n^{o(1)}$, and bounded-precision entries. This is an existential complexity theory result — it proves that the computational problem of computing transformer gradients is not inherently quadratic, and provides a constructive algorithm achieving the better bound.

What the proof does NOT establish:

1. Practical feasibility: The $n^{o(1)}$ notation hides potentially large constants. The polynomial kernel approximation (Lemma C.13) requires constructing degree-$g$ polynomials where $g$ depends on the desired accuracy and the norm bound $R$. For $n = 10^5$ and error $1/n = 10^{-5}$, the degree $g$ could be substantial (potentially hundreds), making $k_1$ large. The row-wise Kronecker product constructions for $z_7$ and $z_4$ produce matrices of rank $k_1 d$, which with $k_1$ in the hundreds and $d$ modest (say 64) could mean working with rank tens of thousands — potentially comparable to $n$ itself for moderate $n$. The asymptotic $n^{1+o(1)}$ guarantee says that for sufficiently large $n$, the running time beats quadratic, but the crossover point — where the approximate algorithm actually becomes faster than the exact algorithm — is not established and could be astronomically large.

2. End-to-end training: The paper proves that a single gradient computation (one backward pass) can be approximated in near-linear time. It does not analyze what happens when these approximate gradients are used iteratively in stochastic gradient descent over many training steps. The $1/\text{poly}(n)$ per-step error could accumulate over training iterations, or the bias in the gradient estimates could cause SGD to converge to a different point than with exact gradients. This is a significant gap between "we can compute gradients quickly" and "we can train models with these gradients."

3. Memory complexity: The paper analyzes time complexity but not space complexity. The low-rank matrices $U_k, V_k$ of size $n \times k_1$ (with $k_1 = n^{o(1)}$) could require super-linear memory. The $n^{o(1)}$ factor might be $(\log n)^{10}$ or $2^{\sqrt{\log n}}$, either of which could be memory-prohibitive for large $n$. The paper does not claim memory bounds.

4. Specificity of the $g_i$ model: The paper assumes $g_i(Z) = \varphi(Z W_g)$ — a single linear layer followed by element-wise activation. Real transformer blocks contain more complex components: layer normalization (which involves computing mean and variance across the feature dimension, a non-element-wise operation), residual connections with pre-norm or post-norm ordering, and MLP blocks with two linear layers and an activation (e.g., $W_2 \cdot \varphi(W_1 Z)$). The paper claims in Section 6 that residual connections and multi-head attention are handled, but the analysis of $g_i$ backpropagation (Lemma H.2) assumes the simple form $\varphi(Z W_g)$. Extending to layer norm and two-layer MLPs would require additional analysis that is not provided.

5. The causal mask analysis is partial: Lemma I.3 provides an $O(n k_1)$ algorithm for $(M \odot (U_0 V_0^\top)) \cdot v$ for a single vector $v$. The extension to matrices $H$ with $k_H$ columns costs $O(n k_1 k_H)$ (Lemma I.4). But $k_H$ could be $d$ (for gradient terms involving $X$ or $G_i$), making this step $O(n k_1 d) = n^{1+o(1)}$. However, this algorithm is inherently sequential — the cumulative sum in Algorithm 2 processes positions in order. This cannot be parallelized across the sequence length, which could limit practical speedups on GPUs. The paper acknowledges this as a GPU implementation challenge (Section B.4) but does not analyze the parallel complexity.

Claim: "The algorithm works for general loss functions."

What the proof establishes: The gradient expressions are derived in terms of the abstract upstream gradient matrix $G_i = \frac{dL}{d\text{Attn}_i}$, without specifying the loss function. Any differentiable loss that decomposes as $\sum_{j,k} \ell(X)_{j,k}$ produces some $G_i$, and the subsequent acceleration steps depend only on $G_i$ as a matrix, not on its origin.

Limitations: The paper does not analyze whether the norms of $G_i$ (which appear in error bounds as $\|G_i\|_\infty \leq \text{poly}(n)$) could be large for certain loss functions or certain points in training, potentially amplifying errors beyond the $1/\text{poly}(n)$ bound. For example, cross-entropy loss with near-zero predicted probabilities can produce large gradients. The $O(\log n)$ bit precision assumption bounds this, but only by assuming the issue away rather than analyzing it.

Claim: "The algorithm handles residual connections, causal masks, and multi-head attention."

What the proof establishes:

• Residual connections (Section J): The proof shows that if a module's gradients are computable in $n^{1+o(1)}$ time, adding a residual connection adds $O(nd)$ overhead and at most constant error amplification. This is solid — the chain rule gives $\frac{dL}{dX} = \frac{dL}{dY} + \frac{dL}{dY} \cdot \frac{dH}{dX}$, and both terms are handled by existing lemmas.
• Causal masks (Section I): The cumulative-sum algorithm (Algorithm 2) is proven to work for both dot-product and Hadamard-product gradient patterns. The analysis is rigorous — Lemmas I.7 and I.8 explicitly map each gradient term to the appropriate masked computation.
• Multi-head attention (Section K): The proof shows the gradient is the sum over heads, each computed via the single-head fast algorithm with dimension $d_h = d/h$. Since $h$ is constant, the cost is $h \cdot n^{1+o(1)} = n^{1+o(1)}$.

Limitations:

• The residual connection analysis (Lemma J.3) assumes the gradient through the sub-module is already approximable in $n^{1+o(1)}$ time, and then shows the residual connection doesn't break this. This is a modularity result — it does not independently verify that the sub-module gradient (e.g., through a complex $g_i$ with layer norm and MLP) is $n^{1+o(1)}$. It only says: IF you can do the sub-module fast, THEN adding residual connection keeps it fast.
• The causal mask analysis reuses the low-rank approximation of the UNMASKED attention matrix $A \approx U_0 V_0^\top$. The cumulative-sum algorithm then computes $(M \odot (U_0 V_0^\top)) \cdot v$. But the approximation error guarantee $\|U_0 V_0^\top - A\|_\infty \leq \epsilon/\text{poly}(n)$ applies to the unmasked matrix. The error of the masked version $\|M \odot (U_0 V_0^\top) - M \odot A\|_\infty$ is identical entry-wise (the mask just zeros out some entries), so the bound carries over. This is handled implicitly but not stated as a lemma.

What is genuinely missing:

1. No analysis of the approximation rank $k_1$ as a concrete function of $n$ and $\epsilon$. The paper states $k_1 = n^{o(1)}$ but does not bound the growth rate — is it $(\log n)^{10}$, $2^{\sqrt{\log n}}$, or something else? This matters because a $2^{\sqrt{\log n}}$ rank for $n = 10^5$ is about $2^{10.8} \approx 1750$, while $n$ itself is $10^5$, so the "low-rank" approximation is only a ~50× compression — potentially not enough to overcome the overhead of the polynomial construction.

2. No analysis of when approximate training outperforms exact training. Even if the per-iteration gradient computation is faster, the approximate gradients might require more iterations to converge (or might not converge at all to a useful point). The paper does not analyze the optimization dynamics.

3. No comparison to alternative sub-quadratic attention methods (e.g., sparse attention, linear attention, kernelized attention, Linformer, Performer, Reformer). The only baseline is exact $O(n^2)$ computation. In practice, sub-quadratic attention methods exist and are used — the paper does not position its theoretical guarantee relative to the theoretical or empirical performance of these alternatives.

4. No analysis of the forward pass cost in the training loop. Training requires both forward and backward passes. The paper focuses on the backward pass, but the forward pass of the same transformer with low-rank attention would also need to be approximated (using Alman and Song, 2023's result). The paper does not analyze the combined forward+backward cost or whether the forward approximation error interacts with the backward approximation error across training iterations.

5. The algorithm is not implemented or empirically validated. This is a pure theory paper — no experiments, no empirical runtime measurements, no accuracy evaluations on any task. The theoretical guarantees are asymptotic and worst-case; whether they translate to practical improvements at realistic scales ($n = 10^3$–$10^5$, $d = 64$–$4096$, $m = 12$–$96$) is entirely unknown. The paper acknowledges the lack of empirical support in the conclusion:

"While we lack enterprise-scale computational resources for training large language models to provide empirical support, our theoretical findings suggest that we can accelerate the training of LLMs in practice."

This is a significant limitation for a paper that claims to enable "more effective training and deployment of long-context language models." Without empirical evidence, it is impossible to assess whether the $n^{1+o(1)}$ asymptotic advantage manifests at realistic scales or is swamped by the constant factors from the polynomial approximation.

Summary judgment: The paper succeeds as a complexity-theoretic result: it proves that computing transformer gradients with high accuracy is not inherently quadratic, and provides a constructive algorithm achieving the better bound. The error propagation analysis across multiple layers is genuinely novel and well-executed. However, the paper does NOT demonstrate that training transformers with this algorithm is practical, efficient at realistic scales, or competitive with existing approximate training methods. The gap between "exists an $n^{1+o(1)}$ algorithm" and "this algorithm accelerates real LLM training" is substantial and unaddressed. The paper would be significantly strengthened by: (1) analysis of the approximation rank $k_1$ as a concrete (not just asymptotic) function of $n$ and $\epsilon$, (2) memory complexity bounds, (3) analysis of optimization convergence with approximate gradients, and (4) empirical runtime measurements at moderate scale ($n = 10^3$–$10^4$) to assess constant factors.

6. Limitations and Trade-offs

1. The Difficulty Estimation Cost Is Not Accounted For in the Headline Efficiency Gains

The paper's entire compute-optimal allocation framework depends on estimating each prompt's difficulty before allocating the inference budget, but the estimation procedure itself is extraordinarily expensive and its cost is excluded from all reported gains. The authors are transparent about this in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The paper's oracle difficulty estimator requires generating 2,048 complete samples per question from the base model and computing the pass@1 rate — 2,048 samples per prompt is 8–16× larger than the largest test-time compute budgets studied (128–256 generations). The predicted (non-oracle) difficulty estimator requires the same 2,048 samples plus PRM scoring, which adds further cost. In a realistic deployment, the total inference cost for a prompt would be difficulty estimation PLUS strategy execution, and the former would dominate — potentially making the total cost worse than simply allocating a uniform large budget to every prompt.

What evidence exists in the paper: The $4\times$ efficiency gains (e.g., compute-optimal matching 64-generation best-of-N performance with only 16 generations in Figure 4) are computed after difficulty is already known, with the estimation cost amortized to zero. The paper does not report total FLOPs including difficulty estimation, does not study how the optimal strategy changes when estimation cost is included in the budget, and does not evaluate whether estimating difficulty with fewer samples (e.g., 8–16 instead of 2,048) would preserve the gains.

Mitigation status: The authors flag this as "a key avenue for future work" (Section 3.2), suggesting that difficulty could be predicted directly from the question text (via a trained classifier) or estimated adaptively with a small initial sample allocation. Neither approach is developed or evaluated. Until this gap is closed, the $4\times$ efficiency figure should be understood as an upper bound on potential gains conditional on having a cheap difficulty oracle that does not yet exist.

---

2. The Theoretical Guarantees Assume Unrealistically Small Hidden Dimension ($d = O(\log n)$)

The entire $n^{1+o(1)}$ running time guarantee depends critically on the assumption that the hidden dimension $d$ is at most logarithmic in the sequence length $n$ — formally $d = O(\log n)$. This assumption appears pervasively:

• It ensures that the rank parameters in the low-rank approximations stay sub-polynomial ($k_1 = n^{o(1)}$, $k_1 d = n^{o(1)}$ in the Hadamard product constructions for $z_7$ and $z_4$).
• It ensures that the cost of backpropagating through the non-attention components $g_i$ remains $O(n d^2) = n (\log n)^2 = n^{1+o(1)}$ rather than $O(n d^2)$ with a potentially large constant.
• It ensures that $d^2 = (\log n)^2 = n^{o(1)}$, so that the final matrix multiplications (e.g., $X^\top U_1 \in \mathbb{R}^{d \times k_1}$ extended to $d \times n$) do not dominate.

The consequence: Modern language models have $d$ in the range 64–8,192, while context lengths $n$ are in the range $10^3$–$10^5$. For $n = 10^4$, $\log n \approx 9.2$, so the assumption $d = O(\log n)$ would require $d \leq c \cdot 9.2$ for some constant $c$, which is violated by even small transformer models ($d = 64$ is ~7× larger). For LLaMA-style models with $d = 4,096$ and $n = 128\text{k}$, $\log n \approx 11.7$, so $d$ is ~350× larger than the assumption allows.

The practical meaning is that the $n^{1+o(1)}$ asymptotic guarantee does not apply to current architectures operated at current scales. The algorithm may still be faster than quadratic in practice (the hidden constant in the $n^{o(1)}$ terms could absorb large $d$), but the paper provides no analysis of running time as a function of both $n$ and $d$ — e.g., whether the cost is $O(n \cdot \text{poly}(d))$ (linear in $n$ with a large $d$-dependent constant) or $O(n d^{\text{polylog}(n)})$ (potentially worse than quadratic for large $d$). Without this analysis, a practitioner cannot determine whether the theoretical speedup materializes at their $(n, d)$ scale.

What evidence exists in the paper: The paper does not analyze the dependence on $d$ beyond stating the assumption. The row-wise Kronecker product constructions for $z_7$ and $z_4$ produce matrices of rank $k_1 d$ (Lemmas E.4 and E.9), where $k_1 = n^{o(1)}$. If $d = 4,096$ and $k_1$ is, say, $(\log n)^3 \approx 1600$, then $k_1 d \approx 6.5 \times 10^6$ — larger than $n$ itself for moderate $n$, making the approximation more expensive than working with the full $n \times n$ matrix directly. The paper does not acknowledge this regime or provide bounds that separate the roles of $n$ and $d$.

Mitigation status: The paper states the assumption explicitly (Theorem 1.4: "We assume $d = O(\log n)$") and cites prior work (Alman and Aggarwal, 2022) that uses the same assumption, but does not discuss its restrictiveness, explore relaxations, or analyze what changes when $d$ is treated as an independent parameter.

---

3. No Empirical Validation at Any Scale — Crossover Point Between Asymptotic Guarantee and Practical Utility Is Unknown

The paper provides zero empirical results — no runtime measurements, no accuracy evaluations on any benchmark, no comparison to exact gradient computation, and no assessment of whether the algorithm produces useful gradients for training. This is a pure complexity theory paper. The conclusion states:

"While we lack enterprise-scale computational resources for training large language models to provide empirical support, our theoretical findings suggest that we can accelerate the training of LLMs in practice."

The consequence: The asymptotic guarantee $n^{1+o(1)}$ says that for sufficiently large $n$, the approximate algorithm will be faster than exact $O(n^2)$ computation. However, "sufficiently large $n$" is not quantified. The $n^{o(1)}$ notation hides factors that could be enormous for practical sequence lengths. Consider:

• The polynomial kernel approximation (Lemma C.13) constructs degree-$g$ polynomials where $g = O(\max\{\log(1/\epsilon) / \log(\log(1/\epsilon)/R), R^2\})$ (Lemma I.1). For $1/\text{poly}(n)$ error and bounded entries with norm $R$, the degree $g$ grows with $n$. The rank $r \leq \binom{2(g+d)}{2g}$ grows combinatorially with $g$. For $n = 10^5$ and modest $g$, $r$ could be in the thousands or tens of thousands.
• The row-wise Kronecker products for Hadamard terms multiply this rank by $d$, potentially pushing the effective working dimension above $n$ itself.
• The cumulative-sum algorithm for causal masks (Algorithm 2) is inherently sequential in the sequence dimension — it cannot be parallelized across tokens, which could limit practical GPU throughput.

None of these constant factors are quantified in the paper. A practitioner deciding whether to implement this method has no way to estimate whether the crossover point — where approximate gradients become faster than exact gradients — occurs at $n = 10^3$, $n = 10^6$, or $n = 10^{12}$. If the crossover is at $n = 10^8$ but their context length is $10^5$, the theoretical result has no practical value for their use case.

What evidence exists in the paper: None. There are no experiments. The paper does not attempt even small-scale validation (e.g., $n = 10^3$, $d = 32$) to confirm that the algorithm produces gradients within the claimed error bound or that it runs faster than a naive implementation.

Mitigation status: The paper does not attempt to mitigate this — it acknowledges the lack of empirical support in the conclusion but treats this as an acceptable limitation of a theory paper. Section B.4 identifies "coding challenges" for GPU implementation (defining custom tensor operations, reimplementing PyTorch backpropagation, implementing parallel causal mask operations in CUDA) and leaves implementation as future work. The gap between theoretical result and practical validation is the largest limitation of the paper as a contribution to deployable LLM training, as opposed to complexity theory.

---

4. Only the Attention Mechanism Is Accelerated — Non-Attention Components ($g_i$) Are Modeled Simply and Computed Exactly, with Unverified Assumptions

The paper's acceleration targets only the self-attention mechanism — the source of the $O(n^2)$ bottleneck. The non-attention components $g_i$ (MLP, layer norm, activation functions) are assumed to have forward and backward computation time that is "linear in its input sequence length" (Definition 1.3, Section 1.1). The paper models $g_i$ as $g_i(Z) = \varphi(Z W_g)$ where $\varphi$ is an element-wise activation (Lemma H.1, Lemma H.2).

The consequence: Real transformer blocks contain significantly more complex components that are not captured by this model:

• Layer normalization computes $\text{LayerNorm}(Z) = (Z - \mu) / \sigma \cdot \gamma + \beta$ where $\mu$ and $\sigma$ are the mean and standard deviation along the feature dimension, involving non-element-wise operations (summing across $d$ for each token). The Jacobian of layer norm is not block-diagonal per-token in the same way as a simple activation — it involves interactions between all feature dimensions within each token. This could alter the gradient structure and possibly introduce $O(n d^2)$ terms with larger constants.

• Two-layer MLPs: $g_i(Z) = W_2 \cdot \varphi(W_1 Z)$ involves two matrix multiplications rather than one. The paper can model this as a composition (treat $W_1$ as part of the attention layer's $g_i$ and $W_2$ as part of the next layer's $g_{i+1}$, or vice versa), but this requires careful accounting of where layer boundaries fall and assumes the ordering is flexible.

• Dropout and residual connections in non-standard positions (e.g., pre-norm vs. post-norm architectures) could alter gradient flow in ways not captured by the simple residual connection analysis in Section J.

The paper claims in Section 6 that these are "straightforward" to incorporate, but provides analysis only for residual connections (Section J) and multi-head attention (Section K) — not for layer norm, dropout, or multi-layer MLPs. Lemma H.2's closed form $G_i = (g_{T_i} \odot \varphi'(\text{Attn}_i W_g)) W_g^\top$ is derived assuming $\varphi$ applies element-wise to the result of a single matrix multiply, and it is not obvious that the same structural simplicity carries through to layer normalization or composed operations.

What evidence exists in the paper: The paper does not analyze layer norm, dropout, or practical MLP architectures. The $g_i$ model is stated in Lemma H.1's assumptions: "assuming for any $Z$, we have $g_i(Z) = \varphi(Z W_g)$ where $W_g \in \mathbb{R}^{d \times d}$ and $\varphi: \mathbb{R} \to \mathbb{R}$ denotes any element-wise activation function." There is no verification that this captures the computational structure of realistic transformer blocks.

Mitigation status: The paper does not address this gap. Section 1.1 states that $g_i$ "includes the layer norm, MLP, residual connection, dropout, positional encoding, multi-head concatenation, and other operations" and asserts "all forward and backward computations of these practical modules can be run in linear time with respect to $n$." This is true (they cost $O(nd^2)$ which is linear in $n$ for fixed $d$), but it does not guarantee that the gradient structure remains compatible with the error propagation analysis. If layer norm introduces non-trivial cross-dimensional couplings, the error amplification factors in the chain rule might change.

---

5. Hard Problems Remain Essentially Unsolved — Test-Time Compute Cannot Create Capability Where the Base Model Has None

This is a fundamental capability boundary documented across all methods. The paper's results reveal a sharp distinction: test-time compute (search, revisions, or their optimal combination) amplifies performance on problems where the base model already has some non-trivial probability of producing a correct answer, but provides essentially no benefit on problems where the base model's pass@1 is near zero.

The evidence is stark (Section 5.3, Section 6):

• Search (Figure 3, right): On difficulty bin 5 (hardest quintile), all methods — best-of-N, beam search, lookahead search — hover at 1–3% accuracy regardless of compute budget, even at 256 generations. The base model simply does not produce correct solutions in its proposal distribution, so no amount of search can find them.
• Revisions (Figure 7, right): On bin 5, all sequential-to-parallel ratios produce roughly 2–3% accuracy, with no improvement trend as budget increases.
• FLOPs-matched comparison (Figure 9): The bin 5 scaling curve is essentially flat near 0–5% accuracy regardless of compute budget, while the $14\times$ larger pretrained model shows significantly better performance. The paper reports a -52.9% disadvantage for test-time compute relative to scaling pretraining on hard problems at high inference ratios (Figure 1, bottom-right bar chart).

The consequence: The method offers no path to solving problems that are fundamentally outside the base model's capability. For genuinely hard reasoning tasks, where even the larger model struggles but improved pretraining might eventually help, test-time compute is not a substitute — it amplifies what the model already knows, but does not add new knowledge or reasoning capability. This limits the approach's applicability to domains where the problem distribution is skewed toward the base model's existing competence range.

What evidence exists in the paper: The difficulty-bin analyses in Figures 3 (right), 7 (right), and 9 all show qualitatively different scaling behavior for the hardest problems — flat curves that never escape the noise floor. The paper is transparent about this, explicitly noting in the Section 7 takeaway box that "test-time compute is ineffective when the base model is incapable."

Mitigation status: The paper does not attempt to mitigate this — it treats the boundary as a fundamental characterization of when test-time compute works versus when pretraining is required. This is not a failure of the method per se (it clarifies the method's scope), but is a limitation on its claimed applicability to "training of LLMs in practice" — for frontier models trained to expand capability frontiers, test-time compute sub-quadratic training alone would not suffice.

---

6. The Revision Training Procedure Is Fragile — Attempted Optimization with ReST$^{EM}$ Degraded Performance, and Correct-to-Incorrect Reversions Occur in 38% of Cases

The revision model — which is central to modifying the proposal distribution for improving easy problems — relies on a specific training data construction procedure (Section 6.1) that is both computationally expensive and sensitive to design choices:

• Training data construction requires 64 parallel samples per question, plus filtering for correct/incorrect examples, plus edit-distance-based pairing to ensure the incorrect answer is structurally similar to the correct one. This is a computationally intensive preprocessing step that must be repeated for each base model.
• The model exhibits a 38% correct-to-incorrect reversion rate: approximately 38% of correct answers produced during a revision chain get "revised" back to incorrect ones in the subsequent step (Section 6.1). This happens because the model was trained only on sequences where all in-context answers are incorrect (followed by a correct target), so it has no training signal for what to do when the current answer is already correct.
• Attempted optimization with ReST$^{EM}$ backfired (Appendix K, Figure 16): applying an RL-style self-improvement procedure (ReST$^{EM}$, Singh et al., 2024) to the revision model caused performance to degrade substantially with sequential revisions — at 256 generations, fully sequential performance dropped to ~33.5% compared to ~38.5% at the optimal ratio. The paper hypothesizes that "on-policy data collection in ReST$^{EM}$ exacerbates spurious correlations in revision data, causing the model to fail to learn the revision task properly" (Appendix K).

The consequence: The revision model's effectiveness depends on a specific offline data construction procedure (edit-distance pairing, multi-turn trajectory construction) that may not transfer to other models, tasks, or data distributions. The correct-to-incorrect reversion problem — addressed via a selection mechanism (majority voting or verifier-based selection) across the chain — is an imperfect patch rather than a resolution of the underlying issue (the model was never trained to preserve correct answers). The failure of ReST$^{EM}$ suggests that improving the revision model is non-trivial and that naive optimization can make it worse, limiting the path to future improvements via standard self-play or RL techniques.

What evidence exists in the paper: The 38% reversion rate is mentioned in Section 6.1 but is not studied in detail — the paper does not report how this rate varies with difficulty, chain length, or training hyperparameters. The ReST$^{EM}$ negative result is in Appendix K, Figure 16, and is acknowledged as showing "substantially hurt" performance, but the mechanism is only hypothesized.

Mitigation status: The paper mitigates the reversion problem at inference time via majority voting or verifier-based selection across the chain, but this adds computational overhead and does not address the root cause. The ReST$^{EM}$ failure is not mitigated — the paper simply reports it as a negative result and does not attempt to fix it. The sensitivity of revision training is left as an open problem.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper shifts the theoretical understanding of transformer training from "quadratic attention cost is inherent" to "near-linear time gradient computation is achievable with bounded error." This is a complexity-theoretic advance, not a practical system or empirical result — it proves an upper bound on the computational complexity of the problem, establishing that the quadratic bottleneck is not fundamental but rather an artifact of the naive algorithm.

The shift is best understood as a reframing of the backpropagation problem from an algebraic optimization to a structural classification task. Before this work, the dominant approach to analyzing attention gradients was to treat each partial derivative as an independent algebraic expression, expand everything, and attempt to identify cancellations or simplifications case-by-case. This paper demonstrates that the apparent algebraic complexity is superficial: after collapsing the chain rule's double summations into matrix products and recognizing the two canonical patterns (dot-product forms $f(X) \cdot H$ and Hadamard-product forms $(f(X) \odot (UV^\top)) \cdot H$), every gradient component across all layers and all weight types reduces to one of these two patterns. This taxonomy of gradient structure — developed across Sections D, E, F, and G — is the paper's deepest conceptual contribution. It provides a diagnostic tool: any future attention variant or training method can be checked against these two patterns to determine whether the low-rank acceleration machinery applies.

The paper resolves a specific contradiction in the prior theoretical literature. Alman and Song (2023) proved the forward pass could be near-linear, and Alman and Song (2024a) proved the backward pass for a single layer with $\ell_2$ loss could be near-linear, but the field lacked any result for the complete training of multi-layer models with practical loss functions. The absence of such a result was not due to lack of interest — it was blocked by three specific technical barriers that this paper addresses:

1. Error propagation across layers: Prior single-layer results provided no mechanism for analyzing how approximation errors compound through the chain rule. The paper's mathematical induction argument (Lemma H.4) proves that error grows additively ($m \cdot \text{poly}(n)$) rather than multiplicatively ($\text{poly}(n)^m$), because the chain rule involves only linear operations with bounded matrix norms. This converts a single-layer curiosity into a multi-layer guarantee.

2. Loss function generality: Prior work exploited $\ell_2$-specific simplifications that do not hold for cross-entropy. The paper's introduction of the abstract upstream gradient $G_i$ (Definition 3.3, Lemma 3.4) separates the problem into two independent parts — compute $G_i$ from any loss (which costs $O(nd)$ for sum-decomposable losses) and push $G_i$ through attention (which works for any $G_i$). This separation is what makes the framework "general" rather than loss-specific.

3. The causal mask "low-rank barrier": The causal mask — a lower-triangular matrix that is full-rank — appeared to break any low-rank approach. The paper identifies the cumulative-sum algorithm (Algorithm 2, Lemma I.3) as the mechanism that sidesteps this issue, and then classifies all gradient terms into two categories (Lemmas I.7 and I.8) that map directly to this algorithm.

These three resolutions together change what researchers should consider possible. Before this paper, one might reasonably assume that near-linear transformer training was impossible — perhaps the quadratic cost was inherent to backpropagation through softmax attention, or perhaps the causal mask was an impassable barrier. After this paper, the research question shifts from "is sub-quadratic training possible?" to "under what conditions is sub-quadratic training practical?" — from existence to efficiency.

The paper also makes certain research directions less attractive. It shows that developing ever-more-complex search algorithms for test-time compute can be counterproductive (lookahead search, the strongest optimizer, performs worst overall due to over-optimization of the verifier — Figure 3, left). This redirects attention toward improving verifier robustness rather than search algorithm sophistication. Similarly, the paper demonstrates that iterative revision — which some prior work claimed was ineffective for reasoning — actually works well but only on the right difficulty tier, reconciling prior contradictory findings and suggesting that the question "does self-correction work?" was ill-posed without conditioning on difficulty.

Follow-Up Research This Work Enables

Bounding the hidden constants in $n^{o(1)}$: explicit rank vs. accuracy vs. $n$ tradeoffs. The paper proves that the approximation rank $k_1 = n^{o(1)}$ — meaning it grows slower than $n^\delta$ for any $\delta > 0$ — but provides no concrete expression. The polynomial kernel approximation in Lemma C.13 (from Alman and Song, 2023) constructs degree-$g$ polynomials where $g = O(\max\{\log(1/\epsilon) / \log(\log(1/\epsilon)/R), R^2\})$ (Lemma I.1), and the rank satisfies $r \leq \binom{2(g+d)}{2g}$. A follow-up could compute $r$ as an explicit function of $n$, $\epsilon$, $d$, and the norm bound $R$, then measure the actual $\ell_\infty$ error achieved at that rank on synthetic or real attention matrices. The key question: for $n = 10^4$, $d = 64$, and $\epsilon = 10^{-3}$, what is $r$? Is it $r \approx 50$ (highly practical), $r \approx 5000$ (comparable to $n$), or $r \approx 10^5$ (worse than exact)? This would determine whether the asymptotic advantage materializes at realistic scales or is swallowed by constant factors. The paper's Lemma E.4 and E.9 involve rank $k_1 d$ for Hadamard terms, which could be the bottleneck — measuring the actual rank needed for these terms would be especially valuable.

End-to-end training with approximate gradients: optimization dynamics and convergence. The paper proves that a single gradient computation can be approximated with bounded error, but does not analyze what happens when these approximate gradients are used iteratively in SGD or Adam over many training steps. A critical follow-up would train a small transformer (e.g., $n = 512$, $d = 64$, $m = 6$) on a standard task (e.g., character-level language modeling or a synthetic reasoning task) using exact gradients vs. approximate gradients at various accuracy levels ($\epsilon = 0.1, 0.01, 0.001$), and measure: (1) whether the optimization converges to the same loss value, (2) whether the convergence rate differs, and (3) whether the bias in the gradient estimates causes systematic drift. The theoretical error bound is on the gradient itself, not on the parameter update after many steps — the gap between "accurate gradient per step" and "accurate trained model after many steps" is substantial and unaddressed.

Memory complexity analysis and space-time tradeoffs. The paper analyzes time complexity but not space complexity. The low-rank matrices $U_k, V_k \in \mathbb{R}^{n \times k_1}$ require storing $O(n k_1)$ entries, where $k_1 = n^{o(1)}$. For the Hadamard product terms, the matrices are $\mathbb{R}^{n \times k_1 d}$ (Lemmas E.4, E.9). When $k_1 d$ exceeds $n$, the memory cost of the approximation exceeds that of storing the full attention matrix directly. A follow-up would characterize the space complexity as a function of $n$, $d$, and $\epsilon$, and identify the crossover point where the low-rank representation becomes larger than the dense matrix it replaces. This would also motivate space-efficient variants — e.g., whether one can avoid materializing all $z_k$ matrices simultaneously, or whether the constructions can be done in a streaming fashion.

Layer normalization and practical MLP architectures. The paper models the non-attention component $g_i$ as $g_i(Z) = \varphi(Z W_g)$, a single linear layer with element-wise activation (Lemma H.1). Real transformers use layer normalization (which computes per-token mean and variance across features, involving non-element-wise operations), two-layer MLPs ($W_2 \cdot \varphi(W_1 Z)$), and dropout. A follow-up would derive the exact closed-form gradient for layer norm, determine whether it falls into one of the two canonical patterns (dot-product or Hadamard), and check whether the error propagation analysis (Lemma H.4) still holds — specifically whether the bounded-norm assumption remains valid for the layer norm Jacobian, which can amplify gradients when the variance $\sigma$ is small. A concrete experiment: compare the empirical Jacobian norm of layer norm vs. a simple activation on actual transformer hidden states to assess whether the paper's $\text{poly}(n)$ norm bound is realistic.

Comparison to empirical sub-quadratic attention methods. The paper's only baseline is exact $O(n^2)$ computation, but practitioners use sparse attention, linear attention, kernelized attention, Linformer, Performer, and FlashAttention. These methods have varying theoretical guarantees (or none) but are implemented and benchmarked. A follow-up would implement a simplified version of Algorithm 1 (perhaps restricting to the dot-product terms only, which avoid the Kronecker product constructions) and compare its wall-clock time and gradient accuracy against FlashAttention-2 and a Performer-style linear attention on sequences of length $n = 10^3$–$10^5$ with a small transformer. The key metric is not asymptotic complexity but actual runtime on GPU and actual model quality after training — the paper's theoretical $n^{1+o(1)}$ could be slower in practice than FlashAttention's $O(n^2)$ with excellent constant factors for all $n$ up to $10^5$.

Failure mode analysis: when does the low-rank approximation break during training? The paper's error bounds assume bounded matrix entries ($\|\cdot\|_\infty \leq \text{poly}(n)$), which is enforced by the $O(\log n)$ bit precision assumption. During actual training, gradients and activations can have large dynamic range — attention logits can become very large or very small, and the softmax can saturate, making the attention matrix near-one-hot (low-rank) or near-uniform (rank 1). A follow-up would track the effective rank of the attention matrix and the quality of the low-rank approximation over the course of training, measuring whether the rank $k_1$ needed to maintain a given accuracy changes (e.g., increases during early training when attention patterns are forming, then decreases as they stabilize). This would inform adaptive-rank strategies that allocate more approximation budget during critical training phases.

Practical Applications and Downstream Use Cases

Theoretical justification for efficient attention hardware and kernels. The paper provides formal complexity-theoretic backing for the intuition that attention computation should not require quadratic time. This has implications for hardware design and kernel development: if the problem is provably sub-quadratic in principle, investments in specialized hardware for near-linear attention (e.g., architectures optimized for low-rank matrix multiply, or I/O patterns designed around the multiplication ordering $U_1 (V_1^\top V)$ rather than $(U_1 V_1^\top) V$) are not chasing an impossibility. FlashAttention (Dao et al., 2022; Dao, 2023) already demonstrates that careful I/O management can make exact $O(n^2)$ attention practical at large scales by avoiding materialization of the full attention matrix — the paper's theoretical result suggests that combining this I/O optimization with algorithmic complexity improvement (trading a small approximation error for asymptotically better time) could yield further gains. This is not a deployment scenario today — the paper's algorithm is not implemented — but it provides theoretical cover for a research direction that might otherwise be dismissed as "quadratic attention is fundamental."

Long-context pretraining budget allocation. For organizations training LLMs with very long context windows (e.g., $n = 10^5$–$10^6$), the quadratic attention cost dominates the training budget. The paper's guarantee — that gradients can be approximated in $n^{1+o(1)}$ time with $1/\text{poly}(n)$ error — provides an upper bound on what is achievable. Even without a practical implementation, this bound can inform resource allocation decisions: if a project is bottlenecked by attention cost at long sequences, the theoretical result suggests that attention approximation (via low-rank, kernel, or other methods) is a promising path rather than a dead end. The paper's analysis of error propagation across $m$ layers (with additive rather than multiplicative accumulation) further suggests that deep transformers with approximate attention may be trainable without catastrophic error compounding — a concern that might otherwise discourage approximate methods for deep models.

Foundations for approximate training of large-scale models under resource constraints. The paper provides a template for analyzing approximate training methods more broadly. The proof structure — decompose gradients into canonical patterns, provide low-rank approximations for the bottleneck matrix, bound error propagation through the chain rule, and use induction across layers — is general and could be applied to other architectures (e.g., cross-attention in encoder-decoder models, graph attention networks, or higher-order tensor attention). For a research group with limited compute who wants to experiment with long-context training, this paper justifies starting with approximate gradient methods rather than assuming exact computation is necessary: the theoretical guarantee establishes that approximate gradients can be provably close to exact gradients across the entire model, which is a stronger statement than the typical empirical approach of "try it and see if validation loss decreases."

When to Prefer This Method

The paper is a theoretical work that does not position itself against named practical alternatives (e.g., FlashAttention, sparse attention, linear attention) via empirical comparison or explicit tradeoff analysis. It provides a complexity-theoretic result that establishes an upper bound — near-linear time with bounded error is possible — without claiming to outperform any specific existing method in practice. The paper explicitly acknowledges in Section B.4 that integration with system-level optimizations like FlashAttention is future work, and in the conclusion notes the lack of empirical validation. Given this framing, a decision matrix comparing "prefer this method vs. alternative X" would be speculative — the paper does not provide the data (runtime measurements, accuracy comparisons, memory benchmarks) needed to make such a recommendation. The appropriate stance is to treat this as a foundational result that changes the theoretical landscape (proving the problem is not inherently quadratic) while awaiting implementation and empirical validation before making deployment recommendations.