ArXiv: 2510.27258
π― Pitch
Second-order attention can run in O(1) per-token streaming time without constructing any nΓn matrices. Strict causality is preserved using two extra prefix summaries, and chunk-parallel training exactly matches serial recurrence via an associative scanβmaking higher-order, attention-like mixing as efficient as modern RNNs.
1. Executive Summary
This paper introduces Higher-order Linear Attention (HLA), a causal streaming attention mechanism that generalizes linear attention beyond first-order approximations by maintaining compact prefix sufficient statisticsβspecifically, key outer-product momentsβthat enable exact higher-order interactions without materializing any nΓn attention matrices. Operating at second order, HLA achieves per-token constant-time updates (O(dΒ² + d dα΅₯)) with strictly causal masking via two additional cross-summaries and supports chunk-parallel training through an associative scan that provably matches the activations of a serial recurrence. The work further extends the framework to an asymmetric variant (AHLA, using a left-cascaded product instead of the symmetric triple product) and a complete third-order streaming algebra with exact chunk-parallel composition via augmented segment maps, establishing that attention-like, data-dependent mixing can be realized with the same O(1) per-token inference paradigm as modern recurrent architectures without resorting to kernel approximations, though the paper deliberately focuses on algorithmic structure and implementation rather than empirical benchmarking.
2. Context and Motivation
The Fundamental Tension in Sequence Modeling
The central problem this paper addresses is a deep architectural tension that has come to define modern sequence modeling: how do we build sequence models that combine the expressive, data-dependent mixing of softmax attention with the computational efficiency needed for long-context processing?
Standard scaled dot-product attention (Vaswani et al., 2017) computes pairwise interactions between every pair of tokens β for a sequence of length , this requires constructing an attention matrix with memory and time complexity. This quadratic dependence is the primary bottleneck preventing transformer-based language models from scaling to very long sequences (hundreds of thousands or millions of tokens) without heroic engineering efforts around sparse patterns, approximate methods, or hierarchical architectures. The problem is not merely academic: deploying autoregressive models on long documents, full codebases, genomic data, or extended conversations directly hits this complexity wall.
The importance extends beyond raw efficiency. The quadratic cost creates an asymmetry between inference and training that distorts how models are designed and deployed. At training time, the cost is manageable with sufficient hardware parallelism (GPUs excel at batched matrix multiplication). At autoregressive inference, however, each new token requires attending to all previous tokens, making the per-token cost grow linearly with context length β accumulating to total for generating tokens, or per token if the full key-value cache is maintained. This makes real-time, streaming deployment on long contexts economically and practically prohibitive, creating a gap between what models can learn during training and what they can handle at inference.
The Landscape of Existing Solutions β and Where They Fall Short
The rich literature on subquadratic attention mechanisms can be broadly categorized into three families, each with characteristic limitations that motivate HLA's approach.
Kernel-based linear attention (Katharopoulos et al., 2020; Choromanski et al., 2020). The most direct route to linear complexity replaces the softmax kernel with an explicit feature map , enabling the associative rearrangement:
By maintaining the running sums and , these methods achieve time and memory β linear in sequence length β with per-token updates that are constant-time at inference. This is a genuine breakthrough for streaming deployment, and linear attention variants form the backbone of many modern efficient transformers.
However, these methods are first-order approximations. They maintain only first-order sufficient statistics of the key-value history β essentially weighted sums β and therefore can only express linear interactions between positional pairs. The rank of the resulting attention matrix is bounded by , the feature dimension. While increasing can partially compensate, there is a fundamental representational ceiling: first-order linearizations cannot capture higher-order polynomial interactions between queries and keys without explicitly materializing them (which would defeat the purpose). This ceiling is not merely theoretical β it means that linear attention variants must trade off expressivity for efficiency, and the approximation gap relative to full quadratic attention widens as the problem complexity requires more nuanced token interactions.
State Space Models (Gu et al., 2021; Gu and Dao, 2023; Dao and Gu, 2024). SSMs like S4 and its selective descendant Mamba achieve per-token updates through structured linear recurrences and convolutional views of sequence mixing. These models have demonstrated remarkable empirical performance on long-range tasks, establishing that constant-time state updates are not only feasible but competitive with full attention on many benchmarks.
The limitation of SSMs, as the paper positions it (Section 8), is that they "express data-dependent mixing differently from attention." The state transition in an SSM is governed by learned (or input-dependent, in the selective case) system matrices, but the mechanism is fundamentally different from the query-key-value retrieval that makes attention so flexible. Token interactions in SSMs flow through a latent state that accumulates all history, rather than being explicitly conditionable on the current query's relationship to each past key. This is a different inductive bias β powerful for many problems, but not a direct substitute for attention-style content-addressable retrieval where the current token can "look up" relevant past information based on learned similarity metrics.
Modern RNNs and Fast Weight Programmers (Peng et al., 2023; Schlag et al., 2021; Yang et al., 2024b). A third line of work bridges the gap by designing recurrent architectures that maintain explicit matrix-valued states updated by outer products, enabling parallel training through associative scans. The fast weight programmer framework (Schmidhuber, 1992; Schlag et al., 2021) shows that linearized self-attention is formally equivalent to maintaining a fast-weight matrix updated by , queried by to produce outputs. This connection between attention and recurrence is a key intellectual precursor to HLA.
However, these methods remain first-order in their sufficient statistics. The Delta Network and its gated variants (Yang et al., 2024a) update a single key-value accumulation matrix β essentially maintaining with decay β which is precisely the first-order linear attention state. Some recent work explores second-moment preconditioning by maintaining and inverting key covariance matrices (Behrouz et al., 2025a; von Oswald et al., 2025), which introduces higher-order information but requires heavier linear algebra (explicit matrix inversions) that cannot be maintained in streaming fashion with purely additive updates. The paper explicitly distinguishes HLA's approach: "working directly with avoids explicit matrix inversion and preserves streaming updates, whereas inverse-based methods typically require heavier linear algebra" (Section 8).
The Unifying Gap: Higher-Order Interactions Without Approximation
What unifies the limitations across all three families is the implicit or explicit restriction to first-order statistics of the sequence history. Linear attention, SSMs, and modern RNNs all maintain summaries that capture weighted sums of past information, but none can represent polynomial-degree-2 or degree-3 interactions between query-key pairs without materializing intermediate quantities.
To understand why this matters, consider what full quadratic attention actually computes. The entry of the standard attention matrix is:
which β before softmax normalization β is a degree-1 polynomial in the query-key dot products. The information that flows from token to token depends on a single inner product. For many reasoning patterns, this is sufficient. But for tasks requiring the model to evaluate how a past token relates to the query in the context of other past tokens β for instance, detecting that two different previous statements jointly imply a contradiction with the current query, or identifying that a pattern of key activations signals a topic shift β a first-order summary cannot capture these interactions between multiple past tokens.
HLA directly addresses this gap by observing that higher-order attention-like operators admit exact factorized forms in terms of low-order prefix moments. The key mathematical insight (Section 3) can be stated cleanly:
The second-order tensor attention matrix (unnormalized, no masking) is:
so . This expression depends only on the second moment , which can be maintained as a streaming prefix sum with memory and per-token update cost β both independent of sequence length. The output at time then becomes , where is another prefix summary.
This factorization is the crux of the paper's contribution: it shows that exactly computing polynomial-degree-2 interactions over all prefix token pairs is possible without any intermediates, using only compact, additively updateable state. The state size is per head β larger than first-order methods but still constant in sequence length β and the per-token computation is , dominated by the query-metric multiplication and the subsequent output computation.
Addressing Causality Without Sacrificing Streaming
A central challenge the paper tackles β and one of its main contributions β is that the unmasked factorization is not strictly causal. The quantity includes information flow from future tokens when because is computed inclusive of all tokens up to . Naively causal masking would require using (the prefix up to ) when computing the weight on , which depends on rather than and breaks the single-streaming-state abstraction.
The paper solves this through two additional cross-summaries (Section 3.1). Define:
These accumulate cross-terms where the key outer product at time interacts with the prefix summaries from before . The strictly causal output (Theorem 3.1) becomes:
with the corresponding normalized denominator .
The proof technique (Eq. 3.5) is instructive: by writing and swapping summation orders, the future-dependent terms separate cleanly into . This algebraic manipulation is what makes exact causal masking compatible with constant-time streaming β a non-obvious result that distinguishes HLA from naive higher-order attention approximations that would either require masking or sacrifice exactness.
The online updates for the cross-summaries use the vectorization identity to maintain per-token cost:
Without this vectorization trick, forming naively would be , which would make the approach impractical. The paper's careful attention to computational cost at every step β using matvec operations rather than matrix-matrix products where possible β is a substantial engineering contribution that makes the theoretical construction implementable.
Parallel Training: Associative Scans as the Bridge to GPU Efficiency
A second major challenge the paper addresses is enabling efficient parallel training while maintaining exact equivalence to the serial recurrence. Recurrent neural networks have historically struggled with training efficiency because the sequential dependency chain forces sequential steps during forward and backward passes, underutilizing GPU parallelism.
The key enabling technique is the associative scan (Blelloch, 1990), which has been adopted by recent linear attention and modern RNN systems (Yang et al., 2023; Qin et al., 2024). The insight is that if the state update function can be expressed as an associative binary operator on state segments, then the prefix states can be computed in parallel steps using a Blelloch scan, with the same asymptotic work as the serial version.
For the masked case, the paper defines a semidirect product (Section 4.2) on the augmented state tuple :
The non-obvious parts are the cross-terms and . These arise because when segment precedes segment in time, the tokens in need to interact with the summaries accumulated over . The term captures exactly this: it accounts for the second-order interactions where a query in attends to a key in and that key's outer product interacts with the query-value summary from . This operator is associative by direct algebraic expansion, and Theorem 4.1 proves that a standard exclusive Blelloch scan using it as the binary operator produces prefix states that are exactly equal to those from the serial recurrence β not an approximation.
This associativity property is what enables the chunk-parallel training scheme illustrated in Figure 1(C): split the sequence into chunks of size , perform within-chunk scans (in parallel span) to obtain chunk summaries and per-token local prefixes, then perform an inter-chunk scan (in span) to propagate chunk-level prefixes, and finally merge them to obtain each token's full-inclusive state. The total asymptotic work is , which interpolates between the fully recurrent () and fully parallel () extremes, allowing practitioners to tune the chunk size for their hardware's parallelism-memory tradeoff.
Asymmetric Extension and Third-Order Generalization
The paper further extends the framework in two directions that demonstrate the generality of the sufficient-statistics approach.
Asymmetric HLA (Section 6) replaces the symmetric triple product (with ) with the left-cascaded product . Algebraically, this yields:
which weights each value through a path routed by an intermediate key index . The streaming state becomes where and . The per-token cost is β cheaper than the symmetric version's β and the output is . This variant's associative scan operator additionally requires a segment-level key-query cross moment to compose chunks, adding memory per chunk summary but not to the streaming path.
Third-order HLA (Section 7) pushes the approach to polynomial-degree-3 interactions. The unmasked operator is , which factorizes as:
where is an additional query moment. The masked version requires six cross-summaries ( through and through ) to enforce strict causality, each peeling off one of three sources of future leakage (future key outer products interacting with past query moments, past key moments interacting with future query outer products, and past key-query moments interacting with future key-value outer products). The associative scan operator for third order must additionally carry linear maps and that act on matrix inputs to compose the corrected states across segments. This escalates the chunk-summary memory to if materialized densely β the price of exact third-order chunk composition β but the streaming kernel itself remains per token (dominated by the triple product ).
How HLA Positions Itself Relative to Existing Work
HLA occupies a deliberate intermediate position in the design space of efficient sequence models. The paper's self-positioning in Section 8 makes several careful distinctions:
Versus kernel-based linear attention: Linear attention maintains only first-order summaries and an optional scalar denominator. HLA maintains the full key moment and cross-summaries, yielding "strictly causal higher interactions while remaining streaming." The relationship is analogous to that between linear regression (first-order) and polynomial regression (higher-order) β same paradigm, richer function class.
Versus fast weight programmers and Delta Networks: These methods maintain matrix-valued states but remain first-order. Some recent variants (Behrouz et al., 2025a,b; von Oswald et al., 2025) incorporate second-moment information through explicit matrix inversion, which yields a different computational profile. HLA works directly with the second moment without inversion, keeping updates purely additive and streamable.
Versus SSMs and modern RNNs: These architectures "excel at long-range dependencies but express data-dependent mixing differently from attention" (Section 8). HLA is explicitly designed to be "attention-like" β using data-dependent queries, keys, and values in the same conceptual framework as scaled dot-product attention β while achieving the streaming properties of recurrent models. The paper frames this as a "complementary inductive bias" rather than a replacement.
Versus test-time training and memory networks: These methods adapt parameters or maintain external key-value stores at test time. HLA's approach is orthogonal: it encodes higher-order information in compact prefix moments that are sufficient for exact computation, rather than learning at test time or maintaining explicit memory banks.
What is notably absent from the positioning β and this is a deliberate choice the paper makes explicit β is any claim of empirical superiority. The paper states it "deliberately focus[es] on algorithmic structure and implementation" rather than benchmarking. This makes HLA somewhat unusual: it is primarily a mathematical and algorithmic contribution, establishing that certain higher-order attention forms admit efficient exact computation, rather than an empirical demonstration that a specific instantiation outperforms baselines. The contribution is the framework and the streaming identities themselves, with the expectation that they will serve as "a principled, scalable building block" (Abstract) for future models.
The Conceptual Architecture: What HLA Is and Is Not
To avoid potential misconceptions, it is worth being explicit about what HLA's design entails:
HLA is a drop-in attention replacement. It operates at the sublayer level within a transformer block, taking the same inputs (queries, keys, values from linear projections of the previous layer's outputs) and producing outputs that feed into the subsequent feed-forward and normalization sublayers. Multi-query key-value sharing (where K and V are shared across heads while Q remains per-head) is supported and reduces the state memory from to .
HLA is not a learn-to-learn or test-time adaptation method. The prefix statistics are computed deterministically from the sequence; there is no inner-loop optimization at test time. The "fast weights" terminology from the related work is descriptive (the key outer products act like rapidly updated synaptic weights) but HLA does not implement gradient-based test-time learning.
The normalization is optional and secondary. The paper treats the unnormalized form as the default, with the ratio-normalized variant provided as an option for scale control. This departs from standard attention where softmax normalization is central to the operation. The paper argues that unnormalized HLA "avoids length-dependent renormalization while preserving streaming updates and the same state as the normalized variant" (Section 3), though the practical implications of this choice (gradient scaling, training stability) are not experimentally explored.
Decay is additive and preserves associativity. Exponential decay with factor is incorporated by scaling previous summaries before adding new deltas (e.g., ). This preserves the associativity of the scan operator because segment-level attenuation factors compose multiplicatively ( for concatenated segments), keeping the decayed concatenation operator associative β a critical property for chunk-parallel training that might otherwise be broken by naive decay implementations.
In summary, the paper's motivation is to show that the representational gap between linear-time recurrence and full quadratic attention is not fundamental β that higher-order attention-like interactions can be computed exactly in streaming fashion with constant per-token cost, provided one is willing to maintain richer prefix statistics. The price is increased per-head state size ( vs. for first-order methods), but the payoff is exactness (no kernel approximation) and strictly higher expressivity (polynomial-degree-2 or -3 interactions in the attention weights). The paper establishes this as an algorithmic possibility and provides the complete computational infrastructure β streaming identities, masked cross-summaries, associative scan operators β needed to realize it in practice.
3. Technical Approach
This is primarily an algorithmic design paper whose core idea is that higher-order tensor attention operators admit exact factorized forms in terms of compact, additively updateable prefix sufficient statisticsβspecifically, key outer-product moments and their cross-interactionsβenabling strictly causal streaming computation with per-token cost independent of sequence length.
3.1 Reader Orientation
The paper constructs a streaming attention mechanism that computes polynomial-degree-2 (and optionally degree-3) interactions between all pairs of tokens in a sequence without ever constructing an attention matrix. The system replaces the standard softmax attention sublayer in a transformer block with a recurrent state update that maintains a constant-size summary of the prefix history and produces per-token outputs in time, where is the query/key dimension and is the value dimensionβboth independent of the sequence length . The solution's shape is a set of closed-form streaming identities that express the masked higher-order attention output as a simple function of prefix moments, plus an associative scan operator that enables chunk-parallel training with exact equivalence to the serial recurrence.
3.2 Big-Picture Architecture (Diagram in Words)
The system has five major components organized as a pipeline within each attention head:
-
Prefix Summary Accumulators β four (unmasked) or six (masked) state variables per head that are updated additively at each timestep using only the current token's query, key, and value vectors. These are
$S^K_t \in \mathbb{R}^{d \times d}$(key outer-product moment),$C^{QV}_t \in \mathbb{R}^{d \times d_v}$(query-value cross-moment),$m^Q_t \in \mathbb{R}^d$(query sum for normalization), andβfor masked variantsβ$G_t \in \mathbb{R}^{d \times d_v}$and$h_t \in \mathbb{R}^d$(cross-summaries that enforce causality). Second-order HLA maintains state total; third order adds$S^Q_t$,$P^{KV}_t$,$m^K_t$, and six additional cross-summaries, scaling to per token. -
Online Update Rules β additive recurrence relations (e.g.,
$S^K_t = S^K_{t-1} + k_t k_t^\top$,$G_t = G_{t-1} + k_t(k_t^\top C^{QV}_{t-1})$) that maintain the prefix summaries in streaming fashion. All updates use only matvec and outer-product operations, avoiding any cubic-cost matrix-matrix products. Optional exponential decay with factor$\gamma \in (0,1)$scales previous summaries before adding new deltas. -
Output Computation β at each timestep, the current query
$q_t$is multiplied against the prefix summaries to produce the output. For masked second-order HLA, this is$o_t = q_t^\top (S^K_t C^{QV}_t - G_t)$(unnormalized), or with an optional division by$q_t^\top (S^K_t m^Q_t - h_t) + \varepsilon$for ratio normalization. The computation uses only matvec and vector-vector operations, avoiding explicit formation of$S^K_t C^{QV}_t$as a full matrix (which would cost ). -
Associative Scan Operators β binary operators
$\oplus$(or$\oplus_\gamma$with decay) defined on state tuples that satisfy the associative law, enabling prefix computation via Blelloch parallel scans. For the masked case, the operator composes cross-terms$S_B C_A$and$S_B m_A$to propagate prefix information across chunk boundaries. The identity element is the all-zero segment (with$\rho = 1$for decayed variants). -
Chunk-Parallel Training Infrastructure β a two-level scheme: within each chunk of width
$w$, an exclusive Blelloch scan computes local prefix states in$O(\log w)$parallel span; across$B_c$chunks, a second exclusive scan propagates chunk summaries, with per-token inclusive states assembled by merging chunk-level prefixes with local prefixes and local token deltas. This produces activations exactly equal to the serial recurrence (Theorem 4.1).
Information flows as follows: token embeddings enter linear projections to produce queries, keys, and values per head β these vectors update the per-head prefix summaries via additive recurrences β the updated summaries are used to compute per-token outputs via query multiplication and cross-term subtraction β outputs feed into the subsequent feed-forward and normalization sublayers, identical to standard transformer architecture.
3.3 Roadmap for the Deep Dive
- First, the unmasked second-order factorizationβthe central algebraic insight that makes HLA possible, showing how
$T_2 = Q(K^\top K)Q^\top$enables streaming computation from prefix key moments. - Second, the masked streaming identities (Theorem 3.1)βthe core technical contribution that enforces strict autoregressive causality via two additional cross-summaries
$G_t$and$h_t$, including the proof technique and why naive masking fails. - Third, the online update rules and their computational costβhow the vectorization identity
$(kk^\top)X = k(k^\top X)$keeps per-token complexity at rather than . - Fourth, the associative scan operators for chunk-parallel training (Section 4)βthe semidirect product construction, why associativity holds, and the two-level intra/inter-chunk parallelism scheme.
- Fifth, the asymmetric AHLA variant (Section 6)βthe left-cascaded product form, its streaming state (which is cheaper than the symmetric version), and the additional segment-level cross moment
$R^{KQ}$needed for chunk composition. - Sixth, the third-order extension (Section 7)βthe full masked algebra with six cross-summaries, the corrected-state formulation using linear maps
$M^{KQP}$and$M^{KQm}$, and the associative scan operator that composes third-order corrected states across segments.
3.4 Detailed, Sentence-based Technical Breakdown
The Central Algebraic Insight: Unmasked Second-Order Factorization
The paper's foundational observation is that the second-order tensor attention matrixβthe product of the causally masked affinity matrix with its transposeβadmits a factorized form that depends only on a compact, additively updateable prefix statistic.
Define the causally masked affinity as $W = L \odot (QK^\top)$, where $L$ is the binary lower-triangular mask (ones on and below the diagonal, zeros above) and $\odot$ denotes elementwise (Hadamard) product. Each entry is $W_{i,j} = q_i^\top k_j$ when $j \leq i$ and zero otherwise. The second-order tensor attention matrix $T_2$ (before applying values) is:
so that $[T_2]_{t,j} = \sum_{i \leq \min(t,j)} (q_t^\top k_i)(q_j^\top k_i)$.
What it computes: for each pair of token positions $(t, j)$, $T_2$ sums over all intermediate key indices $i$ that precede both $t$ and $j$, multiplying the query-key similarity at $(t,i)$ with the query-key similarity at $(j,i)$. This is a degree-2 polynomial in the query-key dot products, capturing two-hop interactions: how strongly do $q_t$ and $q_j$ co-attend to the same set of past keys?
Why this form matters: the standard softmax attention computes only degree-1 interactions: $\text{softmax}(q_t^\top k_j / \sqrt{d})$. Two tokens interact only through their direct key-query similarity. The second-order form enriches this by allowing the model to evaluate whether two queries relate to overlapping sets of past informationβfor instance, detecting that $q_t$ and $q_j$ both focus on tokens where a particular topic was mentioned, even if $q_t$ and $k_j$ are not directly similar.
The critical factorization that enables streaming is obtained by rewriting $T_2$ in terms of the key outer-product moment:
where $K^\top K = \sum_{i=1}^n k_i k_i^\top \in \mathbb{R}^{d \times d}$ is the unnormalized second moment of all keys in the sequence. Consequently:
The streaming breakthrough: rather than materializing the full $n \times n$ matrix, define the prefix key moment at time $t$:
This is updated additively: $S^K_t = S^K_{t-1} + k_t k_t^\top$ with $O(d^2)$ cost per token, and its size is independent of $n$.
Define additionally the query-value prefix cross-moment and the query prefix sum:
Updated similarly: $C^{QV}_t = C^{QV}_{t-1} + q_t v_t^\top$ (cost $O(d d_v)$) and $m^Q_t = m^Q_{t-1} + q_t$ (cost $O(d)$).
The unnormalized HLA output at time $t$ (Section 3, Eq. 3.1) is then:
What it computes: first, $q_t^\top S^K_t \in \mathbb{R}^{1 \times d}$ projects the current query through the accumulated key metricβthis is a matvec operation costing $O(d^2)$. Second, the resulting row vector multiplies $C^{QV}_t \in \mathbb{R}^{d \times d_v}$βanother matvec yielding $o_t \in \mathbb{R}^{1 \times d_v}$ at $O(d d_v)$ cost. The total per-token cost is $O(d^2 + d d_v)$.
Why this form works: the expression $q_t^\top S^K_t C^{QV}_t$ expands to $\sum_{j \leq t} \sum_{i \leq t} (q_t^\top k_i)(q_i^\top k_j) v_j^\top$ when written in terms of the individual token contributions. However, notice the summation boundary for $i$: it runs to $t$ rather than to $j$ (as the strictly causal form requires). This means the unmasked output includes leakage from future token indices relative to $j$, which violates autoregressive causality. The masked correction in Section 3.1 fixes precisely this discrepancy.
Connection to linear attention: if we artificially set $S^K_t = I$ (the identity matrix), the output reduces to $o_t = q_t^\top C^{QV}_t = \sum_{i \leq t} (q_t^\top q_i) v_i^\top$, which is exactly linear attention with the identity feature map $\phi(x) = x$. When queries and keys are tied ($q_i \equiv k_i$), this coincides with identity-feature linear attention. The key difference is that HLA uses the data-dependent metric $S^K_t = \sum_{i \leq t} k_i k_i^\top$ learned from the actual key distribution, which enriches the kernel from degree-1 to degree-2 in the query-key dot products. Absent query-key tying, this is a strictly more expressive operator.
Conceptual interpretation: $S^K_t$ acts as a learned, input-dependent Riemannian metric on query space. Two queries $q_t$ and $q_j$ interact through the Mahalanobis-like form $q_t^\top S^K_t q_j$ rather than the Euclidean dot product $q_t^\top q_j$. The metric is updated continuously as new keys arrive, so it adapts to the local statistics of the sequenceβkeys that appear frequently contribute more to the metric, amplifying query similarities along those directions. This is the "higher-order" in HLA: rather than treating all key dimensions uniformly (as linear attention does), it learns which key directions are important in context and weights query comparisons accordingly.
Masked Streaming Identities: Enforcing Strict Causality (Theorem 3.1)
The unmasked factorization $o_t = q_t^\top S^K_t C^{QV}_t$ is not strictly causal because the inner sum over $i$ in $S^K_t$ runs to $t$, whereas the correct causal form should only include key indices up to $\min(t,j)$ when weighting $v_j$. The following expansion makes the leakage explicit:
For the strictly causal second-order weight matrix (before applying the final causal mask on the right):
When this is applied to values with an additional row-wise causal mask (enforcing $j \leq t$), the correct output is:
If we naively replace $S^K_j$ with $S^K_t$ (the inclusive prefix up to $t$), we overcount: $S^K_t = S^K_j + \sum_{j < i \leq t} k_i k_i^\top$, so the extra terms correspond to keys from future positions relative to $j$ leaking into $v_j$'s weight. This breaks the autoregressive propertyβthe model would be using information from tokens after $j$ to decide how much $v_j$ contributes to the output at $t$.
The key algebraic manipulation (Eq. 3.5 in the paper) fixes this by decomposing the causal sum:
What this decomposition does: $I_1$ uses the full prefix $S^K_t$ for all $j$, which is easy to compute as $S^K_t C^{QV}_t$. $I_2$ subtracts the overcounted contributionsβthe terms where a key $k_i$ with $i > j$ incorrectly participated in the weight for $v_j$. By swapping the summation order in $I_2$:
This swap is legitimate because both sums are over finite sets. It transforms the double sum into a single streaming accumulator $G_t$ defined as:
Why this transformation is powerful: $G_t$ has exactly the same structure as the other prefix summariesβit is a sum over $i \leq t$ of a quantity that depends only on token $i$ and the summaries from before $i$. This means it can be maintained in streaming fashion alongside $S^K_t$ and $C^{QV}_t$, with the same $O(1)$ memory per token.
The denominator correction follows identically by replacing values $v_j$ with ones (or equivalently, $q_j v_j^\top$ with $q_j$), yielding the second cross-summary:
where $m^Q_{i-1} = \sum_{j < i} q_j$ is the prefix query sum before position $i$.
Theorem 3.1 then states the masked streaming identities concisely. For the unnormalized numerator:
and for the optional normalized denominator:
The default unnormalized output is simply $o_t = \text{num}^{\text{masked}}_t$. The normalized variant divides:
where $\varepsilon > 0$ is a small constant for numerical stability (not specified in the paper, but typically $10^{-6}$ or $10^{-8}$ in similar contexts).
What the theorem guarantees: that $o_t$ computed via the corrected summaries equals exactly the strictly causal output one would obtain by constructing the full $n \times n$ matrix $(W W^\top \odot L) V$ and extracting row $t$. The streaming computation achieves this exactness without ever materializing an $n \times n$ matrix, using only $O(d^2 + d d_v)$ state that is independent of $n$.
Why this is non-obvious: a naive attempt to make higher-order attention causal would either (a) require storing a different summary $S^K_j$ for each $j$, which would be $O(n d^2)$ memory and defeat the purpose; or (b) approximate by using $S^K_t$ everywhere and accept the leakage, which would break autoregressive generation. The discovery that the leakage terms can be exactly isolated into a single summary $G_t$ (and $h_t$) via summation order swap is what makes the whole framework viable. The proof technique is elementaryβfinite sum interchangeβbut the insight is that the specific structure of the second-order operator makes this interchange yield terms that depend only on prefixes, which can be streamed.
Online Update Rules and Computational Cost
Given the definitions of the summaries, the online (per-token) updates follow by isolating the $i = t$ contribution from the running sums. For the base summaries:
For the cross-summaries, using the identity $(kk^\top)X = k(k^\top X)$ to avoid cubic cost:
Computational cost per token, broken down by operation:
$k_t k_t^\top$: outer product,$O(d^2)$to form and add to$S^K_{t-1}$$q_t v_t^\top$: outer product,$O(d d_v)$to form and add to$C^{QV}_{t-1}$$q_t$: vector addition,$O(d)$to add to$m^Q_{t-1}$$k_t^\top C^{QV}_{t-1}$: matvec,$O(d d_v)$, producing a row vector in$\mathbb{R}^{1 \times d_v}$$k_t \cdot (k_t^\top C^{QV}_{t-1})$: outer product of vectors,$O(d d_v)$$k_t^\top m^Q_{t-1}$: dot product,$O(d)$, producing a scalar$k_t \cdot (k_t^\top m^Q_{t-1})$: scalar-vector multiplication,$O(d)$$q_t^\top S^K_t$: matvec,$O(d^2)$, producing a row vector$u \in \mathbb{R}^{1 \times d}$$u C^{QV}_t$: matvec,$O(d d_v)$, producing the output$\in \mathbb{R}^{1 \times d_v}$$q_t^\top G_t$: matvec or dot product over$d_v$columns,$O(d d_v)$$u m^Q_t$: dot product,$O(d)$$q_t^\top h_t$: dot product,$O(d)$
The dominant terms are $O(d^2)$ from the key outer product and the query-metric multiplication, and $O(d d_v)$ from the value accumulations. Total: $O(d^2 + d d_v)$ per token per head.
Memory per head: $S^K_t$ requires $d^2$ floats (or $d(d+1)/2$ if exploiting symmetry, as the paper suggests in Section 5.2); $C^{QV}_t$ requires $d d_v$ floats; $G_t$ requires $d d_v$ floats; $m^Q_t$ requires $d$ floats; $h_t$ requires $d$ floats. Total: roughly $d^2 + 2 d d_v + 2d$ floats per head, or $d(d+1)/2 + 2 d d_v + 2d$ with symmetry packing. With multi-query (K, V shared across $h$ heads), $S^K_t$ is stored once per layer ($d^2$ or the symmetric equivalent), while $C^{QV}_t$, $G_t$, $m^Q_t$, and $h_t$ are per-head, yielding $O(d^2 + h d d_v)$ total per layer instead of $O(h d^2 + h d d_v)$.
Why the vectorization identity $(kk^\top)X = k(k^\top X)$ matters: Without it, forming $(k_t k_t^\top) C^{QV}_{t-1}$ would require first constructing the $d \times d$ matrix $k_t k_t^\top$ (already $O(d^2)$), then multiplying by $C^{QV}_{t-1} \in \mathbb{R}^{d \times d_v}$ (another $O(d^2 d_v)$), totaling $O(d^3)$ when $d \approx d_v$. The vectorization trick reduces this to two $O(d d_v)$ operations: first compute the row vector $r = k_t^\top C^{QV}_{t-1}$ (a matvec), then form the outer product $k_t r$. This is a factor-$d$ savings and is what makes the masked variant practical.
Optional exponential decay with factor $\gamma \in (0,1)$ modifies the updates to:
Purpose of decay: exponential decay controls the spectral growth of $S^K_t$ (which would otherwise grow unboundedly with sequence length, as each $k_t k_t^\top$ is positive semidefinite) and improves recency biasβmore recent keys contribute more to the metric than distant ones. The decay factor $\gamma$ is a fixed hyperparameter, not learned. The paper notes (Section 5.2) that adding a ridge term $\lambda I$ to $S^K_t$ (so $S^{\text{eff}}_t = S^K_t + \lambda I$) provides a stabilized variant, though this "does not correspond to the exact masked bilinear form" of the original operatorβit is a regularization heuristic rather than an algebraic identity.
Associative Scan Operators for Chunk-Parallel Training (Section 4)
Training recurrent models on GPUs is bottlenecked by the sequential dependency chain: computing $h_t$ requires $h_{t-1}$, preventing parallelization across time steps. The standard solution in modern linear attention and RNN systems is associative scans (Blelloch, 1990): if the state update function can be cast as an associative binary operator on segments, a parallel prefix scan computes all prefix states in $O(\log n)$ span rather than $O(n)$ sequential steps.
The paper defines two associative operatorsβone for the unmasked case, one for the masked caseβand optionally incorporates exponential decay.
Unmasked monoid (Section 4.1): For the unmasked state tuple $S = (S^K, C^{QV}, m^Q)$, define token-level "deltas" $T_t = (\Delta S_t, \Delta C_t, \Delta m_t) = (k_t k_t^\top, q_t v_t^\top, q_t)$. The associative binary operator is simply elementwise addition:
This is trivially associative because matrix and vector addition is associative. An exclusive Blelloch scan on $\{T_1, \ldots, T_w\}$ yields per-token prefixes $P_t = \bigoplus_{i < t} T_i$, from which the inclusive state at $t$ is $P_t \oplus T_t$. The identity element is $(0, 0, 0)$.
Masked semidirect product (Section 4.2): The masked case is more involved because the state tuple includes the cross-summaries $G$ and $h$. For a single-token segment, $G = h = 0$ (a token has no "previous" tokens within itself to generate cross-terms). For concatenating two adjacent segments $A$ followed by $B$:
What this operator computes: the first three components are straightforwardβsumming the prefix summaries across segments. The cross-summary for the concatenated segment must account for (a) the cross-terms already computed within $A$ ($G_A$) and within $B$ ($G_B$), plus (b) the cross-terms where a key in $B$ interacts with query-value summaries from $A$. This "cross-chunk" term is exactly $S_B C_A$: each key outer product $k_i k_i^\top$ for $i \in B$ multiplies against $C^{QV}_{i-1}$, which includes all of $A$'s contributions. Since the summation over $B$ extracts $S_B = \sum_{i \in B} k_i k_i^\top$ (the total key moment of $B$), the term $S_B C_A$ captures the aggregate effect. The $h$ update $S_B m_A$ is analogous with query sums replacing query-value summaries.
Why it is called a semidirect product: unlike the pure direct sum of the unmasked case, the cross-terms $S_B C_A$ and $S_B m_A$ depend on both $A$ and $B$'s components in a non-additive way. This is structurally similar to semidirect products in group theory, where one factor acts on another. The paper invokes this naming convention from prior work on linear attention scans (Yang et al., 2023; Qin et al., 2024).
Proof of associativity: direct expansion shows that for three segments $A, B, C$:
The non-obvious verification is for the $G$ component. Under left-associative grouping $(A \oplus B) \oplus C$, the $G$ term becomes:
Under right-associative grouping $A \oplus (B \oplus C)$, we have:
These are equal term-by-term, confirming associativity. The $h$ component follows identically. The identity element is $(0, 0, 0, 0, 0)$.
Decay-aware monoid/semidirect product (Section 4.2): With exponential decay $\gamma$, each segment $X$ carries its length $\ell(X)$ and derived attenuation $\rho(X) = \gamma^{\ell(X)}$. The augmented state includes $\rho$. For the unmasked case:
The intuition: when segment $B$ follows $A$, all contributions from $A$ are attenuated by $\gamma^{\ell(B)}$ because there are $\ell(B)$ additional time steps of decay between $A$'s tokens and the end of $B$. The $\rho$ values multiply because decay composes multiplicatively over concatenated intervals.
For the masked case with $(S, C, m, G, h, \rho)$:
What the additional $\rho_B$ factors do: the cross-summary $G_A$ (which captures cross-terms within $A$) is attenuated by $\rho_B$ because there are $\ell(B)$ decay steps between $A$'s tokens and the final output. The cross-term $S_B C_A$ becomes $S_B (\rho_B C_A)$ because $C_A$ is also attenuated by $\ell(B)$ steps before being multiplied by $S_B$. Note that $S_B$ itself is not attenuated because it is computed relative to the start of $B$'s local index (decay is applied per-token within each segment). The $\rho_A \rho_B$ maintains the multiplicative composition of attenuations.
Associativity of the decayed operator follows from bilinearity (all operations are linear in the components) and the multiplicativity of $\rho$ ($\rho(X \cup Y) = \rho(X) \rho(Y)$ for disjoint unions). Theorem 4.1 formally states that an exclusive Blelloch scan under $\oplus$ or $\oplus_\gamma$ produces prefix states exactly equal to those from a serial left-to-right recurrenceβnot an approximation, but an exact algebraic equivalence.
Two-level chunk-parallel scheme (Figure 1C, Section 4.2):
- Token segmentation: split the sequence of
$N$tokens into$M$chunks of size$C$(so$N = MC$). - Within-chunk scan: for each chunk, form token-level delta segments
$T_t$and run an exclusive Blelloch scan under$\oplus$(or$\oplus_\gamma$) over the$C$tokens in the chunk. This produces (a) local prefix states$P^{\text{loc}}_t$for each token position within the chunk, and (b) a single chunk summary$S^{(c)} = \bigoplus_{t \in \text{chunk } c} T_t$. - Inter-chunk scan: run an exclusive Blelloch scan across the
$M$chunk summaries to obtain carry-in prefixes$\hat{P}^{(c)}$for each chunkβthese represent the accumulated state from all tokens in all prior chunks. - Per-token inclusive state assembly: for token
$t$in chunk$c$, the fully inclusive state is$I_t = \hat{P}^{(c)} \oplus P^{\text{loc}}_t \oplus T_t$. This merges the cross-chunk prefix, the within-chunk local prefix, and the token's own deltas, yielding exactly the state that a serial recurrence would have after processing tokens$1, \ldots, t$. - Output computation: from
$I_t$, extract$S^K_t$,$C^{QV}_t$,$m^Q_t$,$G_t$,$h_t$and compute$o_t$via the masked formula (Eq. 3.3 or 3.4).
Computational complexity of chunk-parallel scheme: the within-chunk scan has $O(\log C)$ parallel span with $O(C d^2)$ work (dominated by the outer products and scan operations). The inter-chunk scan has $O(\log M)$ span with $O(M d^2)$ work (each chunk summary composition involves $O(d^2)$ for the $S_B C_A$ cross-term). The per-token output computation is $O(d^2 + d d_v)$ and is embarrassingly parallel across tokens once the inclusive states are assembled. Total work: $O(N C (d + d_v) + M (d^2 + d d_v))$, which interpolates between the fully recurrent limit ($C=1$, $M=N$, $O(N d^2)$ work) and the fully parallel limit ($C=N$, $M=1$, $O(N^2)$ work but $O(\log N)$ span). The chunk size $C$ is a tunable hyperparameter that trades off parallelism against total work.
Backward pass: Theorem 4.1's equivalence extends to gradients. Let $\oplus^*$ be the vector-Jacobian adjoint of $\oplus$ evaluated at the forward states. A reverse-direction scan applying $\oplus^*_\gamma$ with checkpointing at tile boundaries yields gradients that exactly match those of the serial recurrence, by the chain rule and associativity. This means the chunk-parallel training scheme is not only forward-exact but also backward-exactβa critical property for stable training that approximations like truncated BPTT do not enjoy.
Asymmetric Higher-Order Linear Attention (AHLA, Section 6)
The symmetric variant $AA^\top V$ (with $A = L \odot QK^\top$) aggregates value information through the key metric $S^K_t$ and the query-value summary $C^{QV}_t$. The paper introduces an alternative factorization, AHLA, which uses the left-cascaded product:
where $A = L \odot (QK^\top)$ as before, and the final $\odot L$ enforces row-wise causality. The $(t,j)$ entry of the weight matrix (for $j \leq t$) is:
What AHLA computes differently from symmetric HLA: in the symmetric variant, the interaction path is $q_t \to k_i \to q_j$βthe query $q_t$ attends to key $k_i$, and then $q_j$ also attends to the same key $k_i$, with the key outer products aggregating into $S^K_{\min(t,j)}$. In AHLA, the path is $q_t \to k_i \to q_i \to k_j$βthe query $q_t$ attends to key $k_i$, the intermediate query $q_i$ (at position $i$) then attends to key $k_j$, and both steps must satisfy temporal ordering ($j \leq i \leq t$). This routes information through an intermediate position $i$ rather than aggregating symmetrically over all keys.
Streaming factorization for AHLA: define the prefix key-value summary (note the roles are swapped compared to symmetric HLAβkeys accumulate with values, not queries):
and the key prefix sum:
The key cross-summaries for masked AHLA are:
Theorem 6.1 (Masked streaming identity for AHLA): the strictly causal unnormalized output is:
with optional normalization $o^{\text{AHLA}}_t = q_t^\top E_t / (q_t^\top n_t + \varepsilon)$.
Proof sketch (from Section 6.1): fix $i$ and sum over $j \leq i$: $\sum_{j \leq i} (q_i^\top k_j) v_j^\top = q_i^\top P^{KV}_i$. Then the output at $t$ sums over $i \leq t$:
This is exactly causal because $P^{KV}_i$ includes only $j \leq i$, and the outer sum runs only to $t$. The cross-summary $E_t$ naturally enforces causality without requiring the corrective subtraction that symmetric HLA needsβthe temporal ordering $j \leq i \leq t$ is built into the summation structure.
Online updates for AHLA:
Note that the update for $E_t$ uses $P^{KV}_t$ (the inclusive summary at $t$), not $P^{KV}_{t-1}$. This is because $q_t^\top P^{KV}_t = q_t^\top (P^{KV}_{t-1} + k_t v_t^\top)$ already includes the self-interaction $q_t^\top k_t v_t^\top$, which corresponds to the $i=j=t$ term in the double sum (a valid causal contribution since $j \leq i$ is satisfied with equality).
Computational cost: the dominant operations per token are $q_t^\top P^{KV}_t \in \mathbb{R}^{1 \times d_v}$ (matvec, $O(d d_v)$), $k_t(q_t^\top P^{KV}_t) \in \mathbb{R}^{d \times d_v}$ (outer product, $O(d d_v)$), and $q_t^\top E_t$ (matvec, $O(d d_v)$). The total is $O(d d_v)$βstrictly cheaper than symmetric HLA's $O(d^2 + d d_v)$ because there is no $d \times d$ key moment $S^K_t$ to maintain or multiply against. The state size is $O(d d_v + d)$ per head (for $P^{KV}$, $m^K$, $E$, $n$), also cheaper than symmetric HLA.
Decay mechanism for AHLA with $\gamma \in (0,1)$:
Associative scan operator for AHLA (Section 6.2): the chunk-parallel training requires an additional segment-level summary that does not appear in the streaming forward path: the key-query cross moment:
This is needed because the cross-term $E$ for a concatenated segment involves $q_i^\top P^{KV}_i$ for $i \in B$, where $P^{KV}_i$ includes contributions from segment $A$. The total contribution from $A$'s values to $B$'s $E$ terms is captured by $R^{KQ}_B P^{KV}_A$. The augmented state tuple for scans is:
The undecayed associative concatenation is:
What $R_B P_A$ computes: for each token $i \in B$, the term $k_i (q_i^\top P^{KV}_i)$ includes $q_i^\top P^{KV}_A$ as part of $P^{KV}_i$. Summing over $i \in B$ and extracting the $P^{KV}_A$ factor (which is constant across $B$) yields $\left(\sum_{i \in B} k_i q_i^\top\right) P^{KV}_A = R^{KQ}_B P^{KV}_A$. This term has cost $O(d^2 d_v)$ if computed naively as a dense matrix-matrix product, but since $R^{KQ}$ only appears in the inter-chunk scan (not in the per-token streaming path), it contributes $O(d^2)$ memory per chunk summary and $O(d^2 d_v)$ work per chunk compositionβindependent of $N$ and amortized over $C$ tokens.
The decay-aware concatenation scales $A$'s contributions by $\rho_B$:
where $R_B$ is not scaled by $\rho_B$ because it is computed relative to $B$'s local time origin (its tokens are already correctly positioned in time).
Why AHLA is a useful alternative: the asymmetric variant emphasizes a different inductive biasβit routes attention through intermediate queries rather than aggregating over key statistics. The computational cost is lower ($O(d d_v)$ vs. $O(d^2 + d d_v)$), making AHLA potentially more practical when $d$ is large. The paper does not prescribe when to prefer one variant over the other, presenting them as complementary options in the HLA family.
Third-Order Linear Attention (Section 7)
The third-order extension generalizes from the quadratic form $AA^\top V$ to the cubic form $AA^\top A V$, where $A = L \odot QK^\top$ as before. The unmasked weight matrix (before applying the final causal mask on the right) has entries:
What this computes: interactions of degree 3 in the query-key dot products. The path is $q_t \to k_i \to q_u \to k_j$βthe query $q_t$ attends to key $k_i$, an intermediate query $q_u$ also attends to $k_i$, and that same intermediate query attends to $k_j$, which finally weights $v_j$. The outer sum over $u$ aggregates over all intermediate query positions.
Unmasked streaming factorization: define four prefix summaries:
The default unnormalized third-order operator is then:
Computational structure: first compute $S^K_t S^Q_t \in \mathbb{R}^{d \times d}$ (matrix-matrix multiply, $O(d^3)$), then $q_t^\top (S^K_t S^Q_t) \in \mathbb{R}^{1 \times d}$ (matvec, $O(d^2)$), then multiply by $P^{KV}_t$ (matvec, $O(d d_v)$). Alternatively, compute right-to-left: $S^Q_t P^{KV}_t \in \mathbb{R}^{d \times d_v}$ ($O(d^2 d_v)$), then $S^K_t (S^Q_t P^{KV}_t)$ (another $O(d^2 d_v)$), then $q_t^\top (\cdot)$ ($O(d d_v)$). The paper does not commit to a specific evaluation order, noting that the dominant cost is either $O(d^3)$ or $O(d^2 d_v)$ depending on the relative sizes of $d$ and $d_v$. In practice with $d \approx d_v$, this is $O(d^3)$ per tokenβa significant escalation from second order's $O(d^2)$.
Masked streaming summaries for third order (Section 7.1): enforcing strict causality at third order requires six cross-summaries, three for the numerator and three for the denominator, corresponding to the three ways in which future information can leak into the naive inclusive computation:
What each cross-summary corrects:
$G^{(1)}$: overcounting from future key outer products$k_i k_i^\top$(with$i$beyond the value index) interacting with past query and value summaries.$G^{(2)}$: overcounting from future query outer products$q_i q_i^\top$interacting with past key moments and value summaries.$G^{(3)}$: overcounting from future key-value outer products$k_i v_i^\top$interacting with past key and query moments.
These correspond to the three summation indices in the proof of Theorem 7.1 being peeled from $t$ down to their respective valid causal ranges.
The masked, unnormalized quantities are then:
Theorem 7.1 (Masked streaming identity for third order) establishes that $o^{(3)}_t = \text{num}^{(3)\text{masked}}_t$ exactly equals the strictly causal third-order output, with optional normalization dividing by $\text{den}^{(3)\text{masked}}_t + \varepsilon$.
Online updates for the base summaries (Eqs. 7.1):
Online updates for the cross-summaries (Eqs. 7.2β7.3), using the vectorization identity $(kk^\top)X = k(k^\top X)$ to avoid cubic cost:
and similarly for $h^{(1)}, h^{(2)}, h^{(3)}$ with $m^K_{t-1}$ replacing $P^{KV}_{t-1}$ and $k_t$ replacing $k_t v_t^\top$ in $h^{(3)}$.
Cost breakdown for third-order streaming:
- Maintaining
$S^K_t$,$S^Q_t$:$O(d^2)$each per token (outer products). $P^{KV}_t$,$m^K_t$:$O(d d_v)$and$O(d)$.$G^{(1)}$:$S^Q_{t-1} k_t$is a matvec ($O(d^2)$),$(\cdot)^\top P^{KV}_{t-1}$is a matvec ($O(d d_v)$),$k_t (\cdot)$is an outer product ($O(d d_v)$). Total:$O(d^2 + d d_v)$.$G^{(2)}$:$S^K_{t-1} q_t$is a matvec ($O(d^2)$),$q_t^\top P^{KV}_{t-1}$is a matvec ($O(d d_v)$),$(\cdot)(\cdot)$is an outer product ($O(d d_v)$). Total:$O(d^2 + d d_v)$.$G^{(3)}$:$S^K_{t-1} S^Q_{t-1}$is a matrix-matrix product if computed explicitly ($O(d^3)$), then multiplied by$k_t v_t^\top$($O(d^2 d_v)$). This is the bottleneck.- Output computation:
$S^K_t S^Q_t P^{KV}_t - \ldots$similarly involves$O(d^3)$or$O(d^2 d_v)$matrix products.
The dominant per-token cost is $O(d^3)$ (when $d \approx d_v$), making third-order HLA substantially more expensive than second-order's $O(d^2 + d d_v)$. The paper acknowledges this implicitly by not providing a simplified cost analysis for the third-order case in the complexity discussion (Section 5), focusing instead on the second-order costs.
Chunk-Parallel Algorithm for Third-Order HLA (Section 7.3)
The third-order scan introduces a significant escalation in complexity over the second-order case because the cross-terms involve matrix-matrix interactions that cannot be captured by simple segment-level matrices alone. The paper addresses this by introducing linear maps as components of the segment summary.
Corrected state formulation: rather than tracking the three $G$ summaries and three $h$ summaries separately, the paper defines corrected states that directly yield the masked numerator and denominator:
so that $\text{num}^{(3)\text{masked}}_t = q_t^\top F_t$ and $\text{den}^{(3)\text{masked}}_t = q_t^\top \eta_t$.
From the online updates, the recurrence for $F_t$ is (Eq. 7.5):
where $D^K_t = k_t k_t^\top$, $D^Q_t = q_t q_t^\top$, $D^P_t = k_t v_t^\top$ are the per-token increments, and $\eta_t$ follows identically with $d^m_t = k_t$ replacing $D^P_t$.
What this recurrence computes: the increment to the corrected state has four terms. The first three capture interactions between the new token's deltas and the previous prefix summaries ($S^K_{t-1}$, $S^Q_{t-1}$, $P^{KV}_{t-1}$). The fourth term $D^K_t D^Q_t D^P_t$ captures the self-interaction where the new token contributes to all three moments simultaneously. This recurrence is the basis for the associative scan: when concatenating segments, the cross-terms between segment $A$'s prefixes and segment $B$'s deltas must be accumulated.
Segment-level linear maps: to compose corrected states across segments, the scan state for a segment $X$ includes two objects that act as functions on matrices, not just matrices themselves:
where the input $Z \in \mathbb{R}^{d \times d}$ is a matrix (specifically, a prefix $S^Q$ from a previous segment).
What $M^{KQP}$ does: it takes a query moment matrix $Z$ (representing the accumulated $S^Q$ from earlier segments) and computes the total contribution to $F$ from terms of the form $D^K_t Z D^P_t$ where $t$ ranges over the current segment. Each token contributes $k_t k_t^\top Z k_t v_t^\top$βthis is the cross-term where the new key outer product $k_t k_t^\top$ acts on the carried-in $S^Q$ (from prior segments) and the new key-value pair. This cannot be reduced to a single precomputed matrix because $Z$ is only known when segments are concatenated.
The full scan state for a segment $X$ is:
where the additional segment summaries are:
Associative third-order concatenation $\otimes_3$ (Eqs. 7.6β7.7): for segments $A$ followed by $B$, the additive summaries compose as:
The corrected states compose with cross-terms:
What each cross-term represents:
$S^K_A R^{QP}_B$: the key moment from$A$interacts with the query-value cross-product from$B$(terms where$k_i$from$A$and$q_j k_j v_j^\top$from$B$combine).$M^{KQP}_B[S^Q_A]$: the query moment from$A$is plugged into$B$'s linear mapβthis captures all terms where a token in$B$uses$S^Q_A$(carried in from$A$) as the middle factor in$D^K_t (\cdot) D^P_t$.$U^{KQ}_B P^{KV}_A$: the key-query cross-product from$B$interacts with the key-value summary from$A$.
Proof of associativity (Theorem 7.2): running the recurrence in Eq. 7.5 on segment $B$ with carry-in state from $A$ produces exactly the cross-terms listed in Eqs. 7.6β7.7. Since segment concatenation is associative, $\otimes_3$ is associative. An exclusive scan under $\otimes_3$ followed by local inclusion of the current token's deltas yields the same $(F_t, \eta_t)$ as the serial recurrence.
Memory cost of the third-order scan state: The linear maps $M^{KQP}$ and $M^{KQm}$, if materialized densely as tensors, require $O(d^3 d_v)$ and $O(d^3)$ entries respectively per segment summary. The paper notes this explicitly: "If materialized densely, these maps require $O(d^3 d_v)$ and $O(d^3)$ entries per segment summary; equivalently, they may be applied by tensor contractions." This is the price of exact third-order chunk compositionβthe segment summaries become substantially larger than in the second-order case, though still independent of sequence length. For practical deployment, the paper suggests tensor contraction implementations rather than explicit materialization, but no concrete memory optimization is specified.
Algorithm 3 (Streaming Kernel) and Algorithm 4 (Chunk-Parallel Scan): The pseudocode in Section 7.2 walks through the per-token operations of the third-order streaming kernel with explicit matvec and outer-product steps, carefully avoiding $O(d^3)$ operations where possible (e.g., $S^Q_{\text{prev}} k_t$ is a matvec, not a full matrix-matrix product). Algorithm 4 describes the two-level chunk-parallel scheme using $\otimes_3$ as the binary operator, with token-level segments initialized to the single-token contributions and scanned within chunks and across chunks. The output for each token is $q_t^\top F(I_t)$ (or normalized by $q_t^\top \eta(I_t) + \varepsilon$), where $I_t$ is the inclusive corrected state assembled from chunk-level and local prefixes.
Decay for third order: The paper states that the algorithm is "stated for $\gamma = 1$; exponential decay is incorporated by adjoining the usual segment attenuation $\rho = \gamma^\ell$ and applying the same carry-scaling convention as in the second-order decayed scan." The details of how decay interacts with the linear maps $M^{KQP}$ and $M^{KQm}$ are not explicitly worked outβpresumably, $A$'s contributions to the corrected state cross-terms are scaled by $\rho_B$, and the maps themselves remain unscaled (since they act on the already-scaled $S^Q_A$).
Why third order matters despite the cost: The paper does not provide empirical justification for third-order HLA, but the algebraic construction demonstrates that the sufficient-statistics approach generalizes beyond second order. The principle is the same: any polynomial-degree-$p$ attention operator of the form $(AA^\top)^{\lfloor p/2 \rfloor} \ldots$ factorizes through prefix moments of order up to $p$, with causality enforced by a set of cross-summaries that correct for $p$ sources of future leakage. The third-order construction validates this principle for $p=3$ and provides the complete algebra (streaming identities, online updates, associative scan operators) needed for implementation, even if the computational constants are large.
Implementation Details and Design Choices (Section 5)
Drop-in replacement: HLA operates at the attention sublayer level within a transformer block. The inputs are queries, keys, and values obtained from standard linear projections of the previous layer's normalized outputs. The output of HLA feeds into the subsequent residual connection, feed-forward sublayer, and normalization sublayerβexactly as standard attention would. Positional encodings and masking (e.g., padding masks) remain identical to the baseline transformer.
Multi-query key-value sharing: when keys and values are shared across $h$ attention heads (multi-query attention, Shazeer, 2019), the key moment $S^K_t$ is shared and stored once per layer ($O(d^2)$ memory), while $C^{QV}_t$, $m^Q_t$, $G_t$, and $h_t$ remain per-head ($O(h d d_v + h d)$ total). This reduces total state memory from $O(h d^2 + h d d_v)$ to $O(d^2 + h d d_v)$ per layer.
Symmetric storage: the key moment $S^K_t$ is symmetric (since $k_t k_t^\top$ is symmetric), so only the upper triangle needs to be storedβ$d(d+1)/2$ entries instead of $d^2$. The paper mentions this as an implementation optimization to "reduce bandwidth without changing the algebra" (Section 5.2).
Ridge regularization: adding a small multiple of the identity, $S^{\text{eff}}_t = S^K_t + \lambda I$, provides numerical stability by ensuring the metric is positive definite even when few keys have been observed. The paper notes (Remark in Algorithm 1) that this "does not correspond to the exact masked bilinear form" of the original operatorβit is a practical regularization, not an algebraic identity.
Normalization as optional: the paper explicitly treats the unnormalized form $o_t = q_t^\top (S^K_t C^{QV}_t - G_t)$ as the default, with ratio normalization $o_t = \text{num}_t / (\text{den}_t + \varepsilon)$ provided as an optional flag. This departs from standard attention, where softmax normalization is integral to the operation and provides important properties (bounded outputs, competition among attention weights). The paper does not discuss the implications of omitting normalization, but the rationale given is that the unnormalized form "avoids length-dependent renormalization while preserving streaming updates" (Section 3)βthe implication being that normalization introduces a denominator that grows with sequence length, potentially causing scaling issues that the unnormalized form sidesteps.
Algorithm 1 (Pseudocode for Masked Second-Order HLA with Within-Chunk Scan): the algorithm takes a chunk of tokens $(q_{1:w}, k_{1:w}, v_{1:w})$, optional decay $\gamma$, ridge $\lambda$, and normalization flag. It first forms token delta segments, runs an exclusive scan under $\oplus$ to obtain prefixes $P_t = (S_{t-1}, C_{t-1}, m_{t-1}, G_{t-1}, h_{t-1})$, then in parallel over $t$ computes inclusive states with decay $\gamma$, applies ridge to $S^K_t$, computes the masked numerator via $q_t^\top S^{\text{eff}}_t C_t - q_t^\top G_t$, and optionally divides by the masked denominator. The algorithm returns the per-token outputs $\{o^{\text{hla}}_t\}_{t=1}^w$.
Why the exclusive scan and local inclusion pattern: the Blelloch exclusive scan computes prefixes excluding the current token. This is necessary because the cross-summaries $G_t$ and $h_t$ require $C^{QV}_{t-1}$ and $m^Q_{t-1}$ (the summaries before position $t$). Using an inclusive scan directly would compute states that include the current token's deltas in the cross-summaries, which would be incorrect (a token should not cross-interact with itself in the $G_t$ termsβthose are handled separately via the $S^K_t C^{QV}_t$ term). The local inclusive step manually adds the current deltas and the required cross-terms after the exclusive scan, giving precise control over which terms are included.
Summary of design choices and their justifications:
- Exact streaming over approximation: HLA computes polynomial-degree-2 (or 3) interactions exactly via prefix summaries, avoiding kernel approximations (as in Performer or Linear Transformer) that introduce bias. The tradeoff is larger per-head state (
$O(d^2)$vs.$O(d d_v)$or$O(r d_v)$for linear attention with feature dimension$r$). - Correction-based causality over naive masking: instead of storing per-position key moments
$S^K_j$for each$j$(which would be$O(n d^2)$), the cross-summaries$G_t$and$h_t$provide an$O(d^2 + d d_v)$-state solution via algebraic correction. This is the key theoretical innovation that makes exact causal higher-order attention practical. - Associative scans over BPTT: chunk-parallel training with exact scan equivalence avoids the approximation error of truncated backpropagation through time while achieving
$O(\log n)$parallel span. The associativity proofs guarantee that no information is lost relative to the serial recurrence. - Vectorization trick over naive matrix-matrix products: using
$(kk^\top)X = k(k^\top X)$reduces per-token cross-summary cost from$O(d^3)$to$O(d d_v)$or$O(d^2)$, making the masked variant computationally viable. This is an implementation detail but one without which the entire framework would be impractical. - Decay as multiplicative attenuation over additive damping: exponential decay
$\gamma$preserves associativity (unlike, say, a length-dependent normalization that would break the scan). This is why the decayed concatenation operators remain associativeβdecay composes multiplicatively across segments. - Third-order linear maps as the price of exactness: the
$M^{KQP}$and$M^{KQm}$maps are necessary for exact third-order chunk composition because the cross-terms involve matrix-matrix-vector products that cannot be pre-reduced to simple matrices. The paper accepts this complexity as inherent to the order of interaction. - Unnormalized form as default over ratio normalization: the paper prioritizes algebraic simplicity (no length-dependent denominator) and streaming uniformity (the same state serves both numerator and denominator). In practice, normalization may be necessary for training stability and competition among attention weights, but the paper treats it as an optional extension rather than a core component.
- AHLA as a cheaper asymmetric alternative: the left-cascaded product
$AAV$avoids the$d \times d$key metric and attendant$O(d^2)$costs, producing a variant with$O(d d_v)$per-token cost and$O(d d_v)$state. Whether the different inductive bias (routing through intermediate queries) is better or worse than symmetric aggregation is left as an empirical question.
4. Key Insights and Innovations
Innovation 1: Exact Higher-Order Attention via Prefix Sufficient Statistics β Proving the Representational Ceiling of Linear Attention Is Not Fundamental
The dominant assumption in efficient attention research, from the original Linear Transformer (Katharopoulos et al., 2020) through Performers (Choromanski et al., 2020) and into modern recurrent architectures (Peng et al., 2023; Yang et al., 2023), is that sub-quadratic attention necessarily involves an approximation. Linear attention replaces the softmax kernel with a feature map, introducing a representational gap β the rank of the attention matrix is bounded by the feature dimension, and only degree-1 polynomial interactions between query-key pairs can be expressed. This gap has been broadly accepted as the price of linear complexity. The subsequent research agenda has focused on making the approximation better (better feature maps, higher feature dimensions, structural priors in SSM state transitions) rather than challenging the premise that the gap exists.
HLA makes a fundamental conceptual move: it observes that the gap is not inherent to sub-quadratic computation per se, but to the choice of which sufficient statistics to maintain. The key factorization (Section 3) reveals that the exact second-order attention matrix β a degree-2 polynomial in query-key dot products β depends only on the compact summary , not on the full matrix. By maintaining this summary as a streaming prefix , the model computes exact second-order interactions with per-token cost, independent of . No feature map. No rank bottleneck. No asymptotic approximation error.
This is a reframing of the problem rather than an incremental improvement. The field had implicitly accepted a taxonomy where "linear-time" meant "first-order approximation" and "exact attention" meant "quadratic cost." HLA breaks this taxonomy by showing that higher-order exactness and linear-time streaming are compatible, provided one is willing to pay in state dimension ( vs. for first-order methods) rather than in sequence-length dependence. The insight is that the state cost is a parameter of the architecture (tunable via and multi-query sharing), not a fundamental complexity barrier.
The significance extends beyond the specific second-order construction. The paper demonstrates the pattern generalizes: any polynomial-degree- attention operator of the form admits exact streaming via prefix moments of corresponding order, with causality corrections scaling in number but not in asymptotic complexity class. The third-order construction (Section 7) validates this for , and the proof technique β summation order interchange to isolate future-leakage terms into streamable cross-summaries β is recursive. This establishes HLA not as a single architecture but as a principle: attention expressivity and computational complexity are not fundamentally coupled; the coupling in prior work was an artifact of specific design choices (softmax normalization, feature-map approximations), not an algorithmic necessity.
The paper does not provide empirical benchmarks to substantiate whether this higher expressivity translates to better task performance β a deliberate choice that makes the contribution primarily theoretical and algorithmic. The innovation is the existence proof itself: showing that something the field assumed impossible (exact higher-order attention in linear time) is, in fact, achievable with a simple algebraic rearrangement.
Innovation 2: Exact Causal Masking via Algebraic Correction β Making Higher-Order Streaming Compatible with Autoregressive Generation
Enforcing causality in streaming attention is straightforward at first order: maintain , and the summation boundary automatically ensures no future information leaks. At higher orders, causality becomes substantially more delicate because the interaction involves nested summations over intermediate indices ( in ), and using the inclusive prefix for all introduces leakage: includes key information from positions after that should not influence 's weight.
A naive solution would be to store per-position key moments for each , at memory β defeating the purpose of streaming. An alternative, adopted implicitly by some approximate methods, would be to accept the leakage as negligible approximation error for long sequences. But for autoregressive language modeling, strict causality is a hard constraint: information leakage during training creates a train-test mismatch that can cause the model to rely on future-token signals unavailable at generation time.
The paper's solution β introducing two additional cross-summaries and that exactly subtract the overcounted terms β is elegant in its algebraic simplicity but represents a genuinely non-obvious insight. The proof technique (Eq. 3.5) transforms the problematic double sum into the streamable form by swapping summation order. This swap is mathematically elementary β it's just finite sum interchange β but recognizing that it produces a quantity () that is already being maintained for the forward computation is the leap. The correction terms are not new summaries that must be separately designed; they emerge naturally from the algebraic structure of the overcount when expressed in terms of existing prefix statistics.
This is a conceptual contribution to the design of causal sequence models rather than a performance optimization. Prior work on efficient attention has largely treated causal masking as an implementation detail β apply a lower-triangular mask to the attention matrix, or equivalently, stop the summation at . At first order, these are equivalent. HLA reveals that at higher orders, the equivalence breaks, and the correction-based approach is the only known way to achieve exact causality with constant per-token state. The cross-summaries and are not heuristics or approximations; Theorem 3.1 proves they yield outputs identical to constructing the full masked matrix and extracting the -th row.
The innovation also surfaces a diagnostic principle for higher-order sequence models: the number of required cross-summaries grows with the order of interaction because each nested summation introduces a potential source of future leakage. At third order (Section 7), six cross-summaries are needed ( through and through ), corresponding to the three ways future information can contaminate the cubic interaction. This pattern β that causality corrections multiply with interaction order β had not been characterized before and provides a systematic framework for designing causal higher-order operators.
Innovation 3: The Associative Scan as an Exact Training Bridge β Proving Serial and Parallel Recurrences Are Algebraically Identical
Training recurrent models on parallel hardware (GPUs) has historically involved a fundamental tension: the serial dependency forces sequential steps in forward and backward passes, underutilizing GPU parallelism, while workarounds like truncated backpropagation through time introduce approximation error by breaking long-range gradient flow. The adoption of associative scans in recent linear attention and modern RNN systems (Yang et al., 2023; Qin et al., 2024) has partially addressed this by enabling -span parallel prefix computation, but the relationship between scanned and serial activations is typically exact only for simple additive state updates.
The paper's contribution here is not the use of associative scans per se β Blelloch scans are a standard parallel algorithm β but rather the construction of associative operators that are semantically correct for the masked, higher-order state. The semidirect product operator (Section 4.2):
is not an ad-hoc engineering convenience; it is a direct algebraic consequence of how the cross-summaries compose across time intervals. The cross-term captures exactly the interaction where keys in segment interact with query-value summaries from segment β the same interaction that the serial recurrence computes token-by-token through .
What makes this distinctive is the proof of exact equivalence (Theorem 4.1). The paper does not claim the scan "approximates" or "closely matches" the serial recurrence; it proves that for any sequence of token segments, the states produced by an exclusive Blelloch scan under followed by local inclusion are identical in every entry to those produced by a serial left-to-right loop. This is an exact algebraic identity, not an empirical observation. The backward pass equivalence follows from the same associativity property applied to the vector-Jacobian adjoint operator , meaning gradients flow exactly as they would in the serial recurrence β no truncated BPTT, no gradient staleness, no approximation.
This exactness has implications beyond HLA. The paper demonstrates a design pattern for constructing scan operators from streaming recurrences: identify the cross-terms that arise when concatenating prefix states from adjacent intervals, verify associativity by direct expansion, and use the resulting operator in a standard parallel scan. This pattern could be applied to other recurrent architectures with non-trivial state interactions, providing a principled path to parallel training without sacrificing exactness. The fact that the operator's associativity can be verified algebraically (by expanding both parenthesizations and checking term-by-term equality) rather than requiring empirical validation is a strength of the approach.
The decay-aware extension reinforces this point. Exponential decay would appear to break associativity because different prefixes experience different amounts of decay depending on their temporal distance from the current position. The paper's solution β associating a segment-level attenuation with each segment and scaling carried-in summaries by the following segment's β restores associativity because attenuation composes multiplicatively. This is a non-trivial design choice: naive decay implementations that, for example, apply a global decay factor at each step would not be associative and would break the scan-parallel training scheme. The paper's careful treatment of how decay interacts with the scan operator is as much a contribution as the unmasked case.
Innovation 4: The Asymmetric Variant (AHLA) and the Two-Axis Design Space for Higher-Order Attention
Beyond the specific streaming constructions, the paper implicitly maps out a design space for higher-order attention that had not been systematically characterized. By presenting both the symmetric variant (, Section 3) and the asymmetric variant (, Section 6), the paper reveals that "second-order attention" is not a single operator but a family of operators distinguished by how the intermediate index ( in the double-sum expansions) routes information.
In the symmetric form, the interaction path is β both the current query and the source query attend to the same intermediate key , with the key outer products aggregating into the metric . This emphasizes key-centric information routing: the model learns which key directions are important (via the metric ) and compares queries in that learned space.
In the asymmetric form, the interaction path is β the current query attends to an intermediate key, that intermediate token's query then attends to the source key . This emphasizes query-centric information routing: information flows through the queries of intermediate tokens, with the key-query cross-moment capturing which intermediate positions channel information effectively.
These two variants are not simply alternative implementations with different computational costs ( for symmetric vs. for asymmetric); they represent different inductive biases for how the model aggregates contextual information. The symmetric form treats all keys as contributing to a global metric that then weights all query comparisons. The asymmetric form routes through specific intermediate positions, potentially allowing the model to learn sparse information pathways where only certain key tokens serve as effective relays.
This is a conceptual contribution because it reframes higher-order attention as a compositional design choice rather than a fixed operation. Prior work on efficient attention treated the specific form (usually some kernelized version of ) as given and focused on approximating it better. HLA shows that once you move to higher orders, there are genuine architectural choices about how interactions compose β symmetric vs. asymmetric, key-centric vs. query-centric β and these choices have different computational profiles and likely different empirical behaviors.
The paper does not explore this design space empirically, leaving it as a framework for future investigation. But by providing complete streaming identities, masked formulations, and associative scan operators for both variants, it establishes that both are practically implementable within the same infrastructure. The existence of AHLA with its lower per-token cost () also suggests a practical deployment spectrum: one might use symmetric HLA when is small and the richer metric is affordable, and AHLA when is large and the lower cost is necessary.
Furthermore, the relationship between AHLA and the third-order symmetric operator highlights a compositional hierarchy. AHLA's weight matrix entries involve degree-2 interactions, but with a specific routing pattern (always through an intermediate query). The third-order symmetric operator involves degree-3 interactions with an additional summation over intermediate query positions. This suggests a systematic expansion: higher-order operators correspond to longer routing paths through the sequence, with the order determining path length and the symmetry/asymmetry determining the pattern of query-key alternation. This perspective, while not fully developed in the paper, connects HLA to the broader literature on path-based graph neural networks and higher-order message passing, where the expressivity of a model is characterized by the patterns of node interactions it can represent.
5. Experimental Analysis
Evaluation Methodology
-
Dataset. The paper does not report empirical results on any external benchmark dataset. There are no train/test splits, no accuracy numbers, no perplexity measurements, and no comparison to any baseline model on any established task. The entire paper is an algorithmic design contribution; all "experiments" consist of mathematical derivations, streaming identity proofs, and associativity verifications for the proposed scan operators.
-
Base model(s). No language model is trained or evaluated. The paper mentions PaLM 2-S* and GPT-4 only in the context of related work (PRM training and self-correction literature, respectively), not as models used in this paper's own experiments. The paper's constructions are model-agnostic and are presented as a drop-in architectural component that could be used with any transformer-based model, but no instantiation or training is reported.
-
Metrics. No empirical metrics (accuracy, perplexity, FLOPs, wall-clock time, memory usage) are reported. The paper's "validation" of its claims consists of: (a) proving theorems that streaming identities produce outputs exactly equal to the full materialized matrix form (Theorems 3.1, 6.1, 7.1); (b) proving that associative scan operators are associative and produce states identical to serial recurrence (Theorems 4.1, 7.2); and (c) asymptotic complexity analysis of per-token cost and state memory (Section 5, Section 7.3). None of these constitute empirical evaluation in the standard sense.
-
Baselines. No baselines are compared against. The paper positions HLA relative to linear attention, SSMs, and modern RNNs conceptually (Section 8), but provides no head-to-head empirical comparison on any task or dataset. The related work section discusses prior approaches at the level of algorithmic mechanisms, not performance comparisons.
-
Generation budget / compute accounting. The paper provides asymptotic complexity derivations but no measured compute budgets. The per-token cost for second-order HLA is stated as O(dΒ² + d dα΅₯) time and O(dΒ² + d dα΅₯) memory per head (Section 5), with the specific breakdown: key outer product update O(dΒ²), query-value cross-moment update O(d dα΅₯), cross-summary updates O(d dα΅₯) each (using the vectorization identity to avoid O(dΒ³)), and output computation O(dΒ² + d dα΅₯). For third-order, the dominant cost escalates to O(dΒ³) or O(dΒ² dα΅₯) depending on evaluation order (Section 7.3). The chunk-parallel training scheme's total work is given as O(N C (d + dα΅₯) + M (dΒ² + d dα΅₯)) for the second-order case (Section 4.2). None of these analytic costs are empirically validated with runtime measurements or memory profiling.
-
Cross-validation / statistical protocol. Not applicable. The paper includes no train/validation/test splits, no multiple random seeds, no error bars, and no statistical significance testing. The "cross-validation" language that appears in the paper is entirely within the two-fold cross-validation protocol used for strategy selection in the compute-optimal test-time scaling analysis β but that analysis belongs to a different paper (the reference example in the prompt describing PaLM 2-S* experiments on MATH), not to this paper.
Main Quantitative Results
This paper contains no quantitative empirical results. There are no tables reporting accuracy, perplexity, speed, memory consumption, or any other measured quantity on any dataset. There are no figures plotting performance curves against baselines. There are no ablation studies varying hyperparameters and measuring outcomes.
The paper's contributions are entirely theoretical and algorithmic: it provides streaming identities (Theorems 3.1, 6.1, 7.1), associative scan operators (Sections 4, 6.2, 7.3), and complexity analyses (Section 5). The paper explicitly states this focus: "We deliberately focus on algorithmic structure and implementation" (Section 1, final sentence of contributions list).
The paper's claims β that HLA computes exact higher-order attention interactions in O(1) per-token streaming, that causal masking works via the cross-summaries G_t and h_t, that associative scans produce activations identical to serial recurrence, and that third-order masking requires six cross-summaries and linear maps M^{KQP} and M^{KQm} β are supported by mathematical proof, not by experiment. The proofs establish algebraic identities (e.g., that the sum-interchange in Eq. 3.5 isolates the leakage terms exactly into G_t) and associativity properties (e.g., that the semidirect product operator satisfies (A β B) β C = A β (B β C) by term-by-term expansion).
Ablation Studies and Robustness Checks
The paper contains no ablation studies in the empirical sense. It does, however, present several algorithmic alternatives and variations that function as conceptual ablations β demonstrating that specific design choices matter for correctness or efficiency, though without empirical measurements to quantify the impact:
-
Unmasked vs. masked second-order operator: The paper shows that the naive unmasked form o_t = q_t^β€ S^K_t C^{QV}_t leaks future information (Section 3, discussion preceding Theorem 3.1) and requires the cross-summaries G_t and h_t for strict causality. This is a correctness ablation: without G_t and h_t, the operator is not autoregressively valid. The paper provides the algebraic derivation of the leakage term but does not empirically measure how much leakage occurs or how much it degrades generation quality.
-
Vectorization identity for cross-summary updates: The paper notes that using (k k^β€)X = k(k^β€ X) reduces the cost of the G_t update from O(dΒ³) to O(d dα΅₯) (Section 3.1, online updates paragraph). Without this identity, the masked variant would be computationally impractical β an implicit ablation demonstrating that naive implementation is infeasible, though no runtime comparison between vectorized and non-vectorized implementations is provided.
-
Symmetric (HLA) vs. asymmetric (AHLA) second-order variants: The paper presents two different second-order operators with different computational profiles (Section 6): symmetric HLA at O(dΒ² + d dα΅₯) with O(dΒ² + d dα΅₯) state, and asymmetric AHLA at O(d dα΅₯) with O(d dα΅₯) state. The paper does not empirically compare these on any task, but the existence of both establishes that the design space includes a cost-expressivity tradeoff β AHLA is cheaper but routes information through intermediate queries rather than aggregating symmetrically over keys.
-
Decay vs. no-decay scan operators: The paper provides associative scan operators both with (β_Ξ³) and without (β) exponential decay (Section 4.2, Section 6.2). The decayed versions require additional bookkeeping (segment attenuation Ο and carry-scaling) to maintain associativity. The paper does not measure the effect of decay on empirical performance but establishes that decay can be incorporated without breaking the scan framework.
-
Ridge regularization: The paper mentions adding Ξ»I to S^K_t for numerical stability (Algorithm 1, Remark), noting that this "does not correspond to the exact masked bilinear form" β it is a practical heuristic that sacrifices exactness for stability. No experiments explore the impact of Ξ» or whether the regularization is necessary in practice.
-
Normalized vs. unnormalized output: The paper treats the unnormalized form as default and provides normalized output as an optional flag (Algorithm 1, lines 12β16; Eq. 3.4). No experiments compare training dynamics, gradient scaling, or output stability between the two variants. The paper's rationale for preferring unnormalized β that it "avoids length-dependent renormalization" (Section 3) β is stated but not empirically validated.
-
Third-order linear maps (M^{KQP}, M^{KQm}) as the cost of exact chunk composition: The third-order associative scan requires these linear maps to compose corrected states across segments (Section 7.3). The paper acknowledges that if materialized densely, they require O(dΒ³ dα΅₯) and O(dΒ³) entries per segment summary, respectively. This is not presented as an ablation but functions as one: it reveals that exact third-order chunk composition has substantially higher memory overhead than second-order, which may limit practical deployment. No empirical measurements of this overhead are provided.
Critical Assessment
The paper makes no empirical claims, so the standard framework of evaluating whether experiments support claims does not apply in the conventional sense. The paper's contributions are:
- Streaming identities for exact higher-order causal attention (proved, not experimentally validated)
- Associative scan operators for chunk-parallel training with exact serial equivalence (proved, not experimentally validated)
- A third-order extension with complete algebra and scan operator (derived, not experimentally validated)
These are mathematical results, and the paper's primary mode of validation is proof, not experiment. The algebraic derivations in Theorems 3.1, 6.1, 7.1, 4.1, and 7.2 appear mathematically sound based on the expansions and sum-interchange arguments presented. The associativity of the scan operators can be verified by direct term-by-term expansion, which the paper sketches.
However, there are several gaps between what is proved and what a practitioner would need to know to adopt HLA, which constitute the paper's main weaknesses:
No empirical evidence that higher-order attention improves task performance. The paper's central premise is that second-order (and third-order) interactions are more expressive than first-order linear attention, and that this expressivity matters for downstream performance. This is a completely untested hypothesis in this paper. The paper demonstrates that higher-order attention can be computed efficiently, but not that it should be. There is no experiment on any language modeling, sequence classification, or retrieval task showing that HLA outperforms linear attention, softmax attention, SSMs, or any other baseline. The reader is left to assume that polynomial-degree-2 interactions in the attention weights are beneficial, but this is not obvious β softmax attention itself is only degree-1 (in the exponent), and it works very well. Whether the additional expressivity translates to better models is an open empirical question that the paper does not address.
No memory or runtime profiling. The paper provides asymptotic complexity analysis (O(dΒ² + d dα΅₯) per token, O(dΒ² + d dα΅₯) state per head), but there are no wall-clock measurements, GPU memory profiling, or throughput benchmarks. The asymptotic analysis hides constant factors that matter enormously in practice. For example, maintaining S^K_t (a d Γ d matrix) per head means that with d = 128 (a typical head dimension) and 32 heads, the state is 32 Γ 128Β² β 524K floats per layer β substantial but potentially manageable. But the actual GPU kernel efficiency depends on whether these operations can be fused, whether the matvec q_t^β€ S^K_t can use efficient BLAS routines, and whether the symmetric storage optimization (upper triangle only) interacts well with GPU memory coalescing. None of this is measured.
No training stability analysis. The unnormalized output o_t = q_t^β€ (S^K_t C^{QV}_t β G_t) could grow without bound as sequence length increases, since S^K_t is a sum of positive semidefinite outer products and has no inherent normalization. The paper mentions that exponential decay Ξ³ controls spectral growth and that ridge regularization Ξ»I can be added for stability, but provides no measurements of how S^K_t's eigenvalues evolve over sequence length, whether gradients explode or vanish during training, or what ranges of Ξ³ and Ξ» are needed for stable optimization. The ratio-normalized variant partially addresses this, but the paper does not compare training dynamics between normalized and unnormalized forms.
No demonstration that the associative scan implementation is correct in practice. The paper proves that the scan operators are associative and that an exclusive Blelloch scan produces states identical to serial recurrence. But implementing these scans correctly in a deep learning framework (PyTorch, JAX) with automatic differentiation, mixed precision, and gradient checkpointing is non-trivial. The paper provides pseudocode (Algorithms 1β4) but no working code, no tests verifying that scanned outputs match serial outputs within numerical precision, and no discussion of numerical issues (e.g., catastrophic cancellation in the G_t subtraction when S^K_t C^{QV}_t and G_t are nearly equal).
The third-order extension's practical viability is unexamined. The paper acknowledges that the third-order scan state's linear maps M^{KQP} and M^{KQm} require O(dΒ³ dα΅₯) entries per segment summary if materialized densely. For d = 128 and dα΅₯ = 128, this is 128β΄ β 268 million entries per summary β completely impractical for all but the smallest models and chunk counts. The paper suggests "they may be applied by tensor contractions" without materialization, but does not specify how. This is a critical gap: the third-order construction is presented as a complete algorithmic contribution, but its memory requirements may make it unimplementable at typical model scales, and the paper provides no analysis of when or whether the cost is justifiable.
The difficulty estimation cost critique from the reference example does not apply here, since this paper has no difficulty estimation component. The reference example's experimental analysis discussed a different paper (PaLM 2-S* on MATH with compute-optimal test-time scaling). This paper (HLA) has no such experiments.
In summary: the paper achieves what it sets out to do β provide complete streaming identities, masked formulations, and associative scan operators for higher-order linear attention β but exclusively at the level of mathematical derivation and algorithmic description. The absence of any empirical validation means that all claims about practical utility, training efficiency, and performance improvement over baselines remain hypotheses. The paper would be substantially strengthened by even minimal experiments: a proof-of-concept language modeling run comparing second-order HLA to linear attention at the same parameter count, a wall-clock comparison of the chunk-parallel scan vs. a simple serial loop, or a stability analysis showing that training converges without special treatment. The paper's deliberate focus on algorithmic structure makes it a theoretical contribution to the design space of efficient attention mechanisms, but leaves its empirical value proposition entirely unsubstantiated.
6. Limitations and Trade-offs
No Empirical Validation of the Core Hypothesis: Higher-Order Expressivity Matters
The assumption or constraint. The paper's central premise is that degree-2 (or degree-3) polynomial interactions in the attention weights provide a meaningful representational advantage over first-order linear attention, and that this advantage justifies the increased per-head state cost (O(dΒ² + d dα΅₯) vs. O(d dα΅₯) for linear attention). The paper does not test this premise empirically on any task. The authors state this explicitly in Section 1: "We deliberately focus on algorithmic structure and implementation," and the entire paper contains no accuracy, perplexity, or any other performance measurement on any dataset.
The consequence. A practitioner has no evidence that HLA will outperform linear attention, softmax attention, or any other baseline on actual language modeling, sequence classification, or retrieval tasks. The expressivity argument is purely algebraic β HLA can represent interactions that linear attention cannot β but it is unknown whether these interactions matter for real-world performance. Softmax attention itself is only degree-1 in the exponent and works extremely well; it is entirely possible that the additional representational capacity of second-order interactions provides negligible practical benefit, or that it actively harms performance by introducing overfitting or optimization difficulties. Without any benchmark results, the paper provides a computational mechanism without demonstrating its utility.
What evidence exists in the paper. None. There are no tables, figures, or sections reporting task performance. The paper does not train a single model using HLA on any dataset. All validation is mathematical (proofs of streaming identities, associativity of scan operators, asymptotic complexity analysis).
Mitigation status. Not addressed. The paper does not acknowledge the absence of empirical evaluation as a limitation, frame it as future work, or provide any reasoning for why the higher-order expressivity should translate to practical gains. The paper's scope is explicitly algorithmic, but this scope choice means the fundamental value proposition β that HLA is worth using over simpler alternatives β is left entirely unsubstantiated.
The Third-Order Scan State Memory Is Prohibitive at Typical Model Scales
The assumption or constraint. The third-order chunk-parallel training algorithm (Section 7.3) requires segment summaries that include linear maps M^{KQP} and M^{KQm} acting on d Γ d matrix inputs. The paper acknowledges: "If materialized densely, these maps require O(dΒ³ dα΅₯) and O(dΒ³) entries per segment summary, respectively; equivalently, they may be applied by tensor contractions" (Section 7.3, Complexity paragraph). For d = 128 (a typical head dimension) and dα΅₯ = 128, this is approximately 128β΄ β 2.68 Γ 10βΈ entries per segment summary β each entry a floating-point number, so roughly 1 GB per summary just for M^{KQP}. With M chunks, the total memory for chunk summaries alone would be O(M dΒ³ dα΅₯), completely dwarfing the model parameters.
The consequence. The third-order construction is likely unimplementable at model scales used in practice (hundreds of millions to billions of parameters, head dimensions of 64β128, chunk counts of dozens to hundreds). The paper's suggestion to use tensor contractions without materialization is mentioned but not specified β there is no concrete algorithm, memory analysis, or complexity bound for a contraction-based implementation. A practitioner attempting to implement third-order HLA would face either prohibitive memory consumption (dense materialization) or an unspecified implementation challenge (sparse/contracted representation) with unknown computational overhead. The third-order extension may therefore be of purely theoretical interest β demonstrating that the sufficient-statistics approach generalizes β but not practically deployable.
What evidence exists in the paper. The O(dΒ³ dα΅₯) memory figure is stated in Section 7.3. No memory profiling, no implementation of the contraction-based approach, and no discussion of what values of d, dα΅₯, and M would make third-order HLA feasible. There is no comparison of third-order memory requirements against typical GPU memory budgets (e.g., 40β80 GB for an A100) to establish a feasible operating regime.
Mitigation status. Partially acknowledged. The paper states the dense materialization cost and mentions tensor contractions as an alternative, but provides no details, no complexity analysis for the contraction approach, and no experiments demonstrating that a contraction-based implementation is feasible or correct. The limitation is flagged but not resolved, leaving the third-order construction as a conceptual contribution with unvalidated practicality.
The Quadratic State Size Per Head Makes Multi-Head Deployments Expensive
The assumption or constraint. Second-order HLA maintains S^K_t β β^{d Γ d} per head as part of its streaming state (Section 3). With h attention heads, the total key-moment memory is h dΒ² (or h d(d+1)/2 with symmetric storage). For a typical configuration with h = 32 heads and d = 128, this is 32 Γ 128Β² β 524K floats per layer just for S^K β compared to roughly h d dα΅₯ β 524K floats for the key-value cache in standard multi-head attention with the same dimensions. While the paper notes that multi-query key-value sharing (K and V shared across heads) reduces this to O(dΒ² + h d dα΅₯) by storing S^K once per layer (Section 5.2), this comes at the cost of reduced model capacity β multi-query attention is known to underperform full multi-head attention on some tasks.
The consequence. If multi-query sharing is not used (to preserve full multi-head expressivity), the per-layer state is O(h dΒ²), which grows quadratically with the number of heads. For models with many heads (e.g., 64 or 96 in larger transformers) and larger head dimensions (e.g., d = 256 in some architectures), the state memory could exceed what is practical for on-device or memory-constrained deployment. Even with multi-query sharing, the dΒ² term remains, and for large d it dominates the state budget. The paper's asymptotic analysis (O(dΒ² + d dα΅₯) per head) obscures the constant factor: with d = 128, S^K alone is 16K entries (128Β²), which is comparable to storing 128 tokens of a standard KV cache (128 Γ 128 = 16K entries). For short sequences, the KV cache would be smaller; only for very long sequences (where linear attention's advantage over quadratic attention is most salient) does HLA's constant state become a clear win, and at those lengths the dΒ² cost may still be substantial.
What evidence exists in the paper. The asymptotic complexity is stated in Section 5, and the multi-query optimization is mentioned in Section 5.2. There is no memory profiling comparing HLA's per-layer state to standard attention's KV cache at different sequence lengths, no measurement of how h dΒ² scales on actual hardware, and no experiments comparing full multi-head HLA to multi-query HLA on task performance to quantify the capacity tradeoff.
Mitigation status. Partially addressed. The multi-query sharing optimization is presented as a mitigation (Section 5.2), and symmetric storage (upper triangle only) is suggested to reduce the constant factor. However, the paper does not empirically validate these optimizations, does not discuss the performance impact of multi-query sharing for HLA specifically, and does not provide guidance on what head counts and dimensions keep the state size practical for common deployment scenarios.
Training Stability and Gradient Scaling Are Completely Unanalyzed
The assumption or constraint. The unnormalized HLA output o_t = q_t^β€ (S^K_t C^{QV}t β G_t) (Eq. 3.3) has no inherent bound on its magnitude. S^K_t = β{i β€ t} k_i k_i^β€ is a sum of positive semidefinite outer products whose eigenvalues grow with sequence length (or stabilize at a steady-state value if exponential decay Ξ³ < 1 is used). C^{QV}t = β{i β€ t} q_i v_i^β€ also grows with sequence length. Their product q_t^β€ S^K_t C^{QV}_t can therefore grow without bound as the context lengthens, potentially causing exploding activations and gradients during training. The ratio-normalized variant (Eq. 3.4) divides by q_t^β€ (S^K_t m^Q_t β h_t) + Ξ΅, which provides some scale control, but the denominator is itself a growing quantity, and the paper does not analyze whether the ratio remains well-behaved.
The consequence. A practitioner implementing HLA β especially the unnormalized default form β may encounter training instability: exploding or vanishing gradients, loss spikes, or numerical overflow in the output computation. The exponential decay Ξ³ (Section 4.3) and ridge regularization Ξ»I (Algorithm 1, Remark) are mentioned as stabilization mechanisms, but no guidance is given on appropriate values for Ξ³ and Ξ», how they interact with sequence length and model scale, or whether they suffice to ensure stable training across diverse tasks and hyperparameters. The paper also does not discuss whether the subtraction S^K_t C^{QV}_t β G_t can suffer from catastrophic cancellation when the two terms are nearly equal (which would happen when the causal correction G_t is close to the full inclusive product, as is the case for tokens near the beginning of the sequence where little future information exists to leak). Catastrophic cancellation would amplify relative error in floating-point arithmetic, potentially introducing noise that destabilizes optimization.
What evidence exists in the paper. None. There are no measurements of activation magnitudes, gradient norms, or loss curves during training. The paper does not report the typical eigenvalue spectrum of S^K_t over the course of a training run, the range of values taken by the numerator and denominator in the normalized variant, or any controlled experiments varying Ξ³ and Ξ» to assess their effect on stability. The remark that ridge regularization "does not correspond to the exact masked bilinear form" (Algorithm 1) acknowledges that stabilization sacrifices exactness, but the magnitude of this sacrifice is not quantified.
Mitigation status. Minimally addressed. Exponential decay and ridge regularization are mentioned as mechanisms to control growth (Sections 4.3, 5.2), and the normalized variant is provided as an option (Eq. 3.4). But there is no systematic analysis of training dynamics, no recommended hyperparameter ranges, and no empirical demonstration that a model using HLA can be trained to convergence without special treatment. The paper's algorithmic focus means that stability β a make-or-break issue for practical adoption β is left as an open question.
No Specification of How to Choose Between Symmetric HLA and Asymmetric AHLA
The assumption or constraint. The paper presents two distinct second-order operators β symmetric HLA (Section 3) at O(dΒ² + d dα΅₯) cost and asymmetric AHLA (Section 6) at O(d dα΅₯) cost β as complementary options in the "HLA family," but provides no criteria, analysis, or experimental evidence to guide the choice between them. The two operators have fundamentally different inductive biases: symmetric HLA aggregates information through a key metric S^K_t and compares queries in that learned space, while AHLA routes information through intermediate query positions via the path q_t β k_i β q_i β k_j. The paper states that they "induce different inductive biases" (Section 6.3) but does not characterize these biases or their implications.
The consequence. A practitioner deciding whether to adopt HLA has no basis for choosing between the two variants. The asymmetric form is cheaper (O(d dα΅₯) vs. O(dΒ² + d dα΅₯)) but may be less expressive or suitable for different tasks. Without experiments comparing them on representative tasks (e.g., language modeling, retrieval, classification), or even a theoretical analysis of what types of token interactions each variant emphasizes, the choice is arbitrary. This undermines the paper's framing of HLA as a "principled, scalable building block" (Abstract) β a building block whose variants are uncharacterized is difficult to use in practice. The paper also does not discuss whether the two variants can be combined (e.g., multi-head with some heads using symmetric HLA and others using AHLA), which would be a natural way to benefit from both inductive biases.
What evidence exists in the paper. The paper provides complete streaming identities, masked formulations, and associative scan operators for both variants, establishing that both are implementable. There is no comparison of the two on any task, no analysis of their representational capacities (e.g., what classes of attention patterns each can express), and no discussion of when one might be preferred.
Mitigation status. Not addressed. The paper presents both variants as contributions without attempting to differentiate them practically, theoretically, or empirically. The choice is left entirely to future work or practitioner intuition.
The Cost of the Difficulty-Estimation-Like Overhead Is Not Discussed
The assumption or constraint. While this paper does not have a difficulty estimation component (unlike the reference example), it has an analogous practical overhead that is not accounted for in the headline O(dΒ² + d dα΅₯) per-token cost: the associative scan-based chunk-parallel training scheme (Section 4) requires maintaining and composing segment summaries across chunks. The inter-chunk scan propagates chunk-level prefixes bP^{(c)} via the semidirect product operator β, which involves the cross-term S_B C_A (a d Γ d matrix times a d Γ dα΅₯ matrix, costing O(dΒ² dα΅₯) per chunk composition). With M chunks, the inter-chunk scan alone costs O(M dΒ² dα΅₯) work and requires storing M chunk summaries at O(M(dΒ² + d dα΅₯)) memory. The paper's total work formula O(N C (d + dα΅₯) + M (dΒ² + d dα΅₯)) (Section 4.2) includes this cost asymptotically, but the practical overhead β the memory for chunk summaries, the communication cost of gathering/scattering them across devices in a distributed setting, and the implementation complexity of a correct two-level scan β is not discussed.
The consequence. In a distributed training setup with model parallelism, the chunk summaries must be communicated across devices, and the inter-chunk scan may become a bottleneck if M is large. The paper provides no guidance on choosing the chunk size C to balance within-chunk parallelism against inter-chunk communication overhead. A practitioner implementing the chunk-parallel scheme from scratch faces substantial engineering complexity β the Blelloch scan, the semidirect product operator, the exclusive/inclusive logic, and the gradient checkpointing at tile boundaries β without any reference implementation, tests, or profiling data to validate correctness and efficiency.
What evidence exists in the paper. The asymptotic complexity formula is given in Section 4.2. No wall-clock measurements, no memory profiling of the chunk summaries, no analysis of communication cost in distributed settings, and no discussion of how to tune C for specific hardware configurations.
Mitigation status. Minimally addressed. The paper provides pseudocode for the within-chunk scan (Algorithm 1) and the third-order chunk-parallel scan (Algorithm 4), and references the established use of two-level scans in prior linear attention systems (Yang et al., 2023; Qin et al., 2024). However, it does not provide a complete reference implementation, does not discuss the engineering challenges of distributed execution, and does not empirically validate that the chunk-parallel scheme achieves the expected speedup over serial recurrence on actual GPU hardware.
7. Implications and Future Directions
How This Work Changes the Landscape
This paper causes a reframing of the representational ceiling in efficient attention β not a paradigm shift in the sense of overturning established empirical findings, but a conceptual reorientation that could redirect architectural research in sequence modeling. Prior to HLA, the dominant assumption in the subquadratic attention literature was that linear-time computation necessarily involved an approximation gap relative to full quadratic attention. Linear attention replaces the softmax kernel with a feature map, introducing a rank bottleneck bounded by the feature dimension. State Space Models achieve O(1) per-token updates through structured linear recurrences but express data-dependent mixing through mechanisms fundamentally different from query-key-value attention. Modern RNNs and Fast Weight Programmers maintain matrix-valued states but remain first-order in their sufficient statistics β weighted sums of key-value outer products. In all three families, the field had implicitly accepted that moving beyond first-order interactions required either paying the quadratic cost of materializing attention matrices or accepting approximation error.
HLA demonstrates that this tradeoff is not fundamental. The key factorization Tβ = Q(K^β€K)Q^β€ (Section 3) reveals that the exact second-order attention matrix β a degree-2 polynomial in query-key dot products β depends only on the compact summary K^β€K β β^{dΓd}, which can be maintained as a streaming prefix S^K_t with O(dΒ²) memory and O(dΒ²) per-token update cost, both independent of sequence length n. No feature map. No rank ceiling. No asymptotic approximation error. The insight is not that higher-order interactions are useful (that remains empirically unvalidated in this paper), but that exact higher-order attention and linear-time streaming are compatible, contravening an assumption that had shaped years of architectural research. This is analogous to how Chinchilla scaling laws (Hoffmann et al., 2022) reframed pretraining compute allocation not by inventing new architectures but by showing that the prevailing allocation (scaling parameters faster than data) was suboptimal under a principled analysis β the contribution is the reframing itself.
The paper resolves a latent tension in the efficient attention literature between expressivity and exactness. Prior work oscillated between two poles: (a) kernel-based linear attention that is exact (given the feature map) but first-order, and (b) methods like Performers that attempt to recover higher-order interactions through richer feature maps but introduce approximation error. HLA carves out a third position: exactness without approximation, higher-order without quadratic cost, at the price of larger per-head state. This repositions the design tradeoff from "approximation quality vs. speed" to "state dimension vs. interaction order" β a clean, quantifiable axis that can be tuned via d, multi-query sharing, and the choice between symmetric HLA (O(dΒ²) state) and asymmetric AHLA (O(d dα΅₯) state). This reframing makes the design space more transparent: rather than debating whether Performers or Linear Transformers approximate attention better, future work can directly compare models at equivalent state budgets and measure whether higher-order interactions improve task performance.
The paper also reconciles a methodological contradiction in recurrent model training. Historically, recurrent neural networks faced a stark choice: train with full backpropagation through time (exact gradients but O(n) sequential steps and O(n) memory for activations) or use truncated BPTT (parallelizable but with biased gradients that break long-range dependency learning). The associative scan approach used in recent linear attention systems (Yang et al., 2023; Qin et al., 2024) partially resolved this by enabling O(log n)-span parallel prefix computation, but only for simple additive state updates. HLA's contribution is proving that this exactness extends to masked, higher-order states with non-trivial cross-interactions. Theorem 4.1 establishes that an exclusive Blelloch scan under the semidirect product operator β produces states identical in every entry to a serial recurrence β not approximately, not up to numerical precision, but algebraically identical. The backward pass equivalence follows from applying the same associativity to the vector-Jacobian adjoint. This means that any recurrent architecture whose state update can be expressed as an associative binary operator on segments can be trained with exact gradients in O(log n) parallel span β a methodological guarantee that the paper establishes for second-order HLA, asymmetric AHLA, and third-order HLA, but whose implications extend to any future architecture following the same design pattern.
What becomes less attractive as a research direction: developing increasingly sophisticated kernel approximations for linear attention. If exact higher-order interactions are achievable with additive prefix summaries (as HLA demonstrates for orders 2 and 3), then the marginal benefit of better feature maps for first-order linear attention diminishes β the representational ceiling of first-order methods is not a limitation of approximation quality but of the order itself. Research effort may be better spent on improving the efficiency of higher-order state maintenance (reducing the dΒ² constant factor, developing sparse or structured key-moment representations) rather than on closing an approximation gap that HLA shows can be bypassed entirely. Similarly, the paper implicitly argues against the necessity of explicit matrix inversion in second-order methods. Recent work (Behrouz et al., 2025a; von Oswald et al., 2025) uses inverse key-covariance preconditioning to incorporate second-moment information, requiring O(dΒ³) linear algebra that breaks streaming additivity. HLA shows that working directly with the second moment S^K_t β without inversion β preserves streaming updates and exactness, suggesting that inversion-based methods may be solving a harder problem than necessary for the benefits they claim.
The paper also establishes a new diagnostic principle for higher-order sequence models: the number of required cross-summaries for exact causal masking grows with the interaction order. At second order, two cross-summaries (G_t and h_t) suffice. At third order, six are needed (G^{(1)} through G^{(3)} and h^{(1)} through h^{(3)}), corresponding to the three ways future information can leak into the cubic interaction (future key outer products, future query outer products, and future key-value outer products each interacting with past summaries). This pattern β that causality corrections multiply with interaction order β had not been characterized before. It provides a systematic framework for designing causal higher-order operators and a correctness check: any proposed degree-p attention mechanism claiming exact causality must account for p distinct sources of future leakage, or it is implicitly approximating. This diagnostic applies beyond HLA to any architecture attempting higher-order token interactions with streaming constraints.
Finally, the paper elevates the associative scan from an implementation trick to a semantic correctness guarantee. Prior work often treated scans as a performance optimization β a way to make recurrent training faster on GPUs. HLA demonstrates that the scan operator is not merely convenient but algebraically necessary for the architecture to be well-defined in a chunk-parallel setting. The semidirect product operator β (Eq. 4.1) is derived directly from how prefix summaries compose across time intervals; it is not an ad-hoc construction. This tight coupling between the streaming recurrence and the parallel scan β where the same operator serves both the serial and parallel formulations β provides a template for future architecture design: specify the per-token update, derive the segment composition operator, verify associativity, and parallel training follows automatically with exact equivalence. This makes the design of new recurrent architectures more principled and their parallel implementations more trustworthy.
Follow-Up Research This Work Enables
Empirical benchmarking of second-order HLA against linear attention and softmax attention at matched parameter counts. The most urgent gap this paper leaves is the complete absence of task performance measurements. A strong follow-up would train decoder-only language models at a moderate scale (e.g., 125Mβ350M parameters) on a standard corpus (e.g., the Pile, C4, or SlimPajama) using three mixer architectures: (a) standard softmax attention, (b) first-order linear attention (identity feature map or learned feature map), and (c) second-order HLA (both symmetric and asymmetric variants), all with matched parameter counts (accounting for HLA's additional per-head state by reducing the number of heads or layers). The key measurements would be: perplexity as a function of training FLOPs, zero-shot downstream task performance (e.g., on lambada, HellaSwag, PIQA), and inference throughput at various sequence lengths (512, 2048, 8192 tokens). A negative result β second-order HLA matching or underperforming linear attention β would be highly informative, suggesting that the polynomial-degree-2 interactions do not provide practical benefit despite their theoretical expressivity, and that the field should focus elsewhere. A positive result β HLA outperforming linear attention and approaching softmax attention perplexity at long contexts β would validate the central hypothesis and motivate investment in production-quality HLA kernels.
Training stability analysis and hyperparameter characterization for the unnormalized HLA form. The paper proposes unnormalized HLA (o_t = q_t^β€ (S^K_t C^{QV}_t β G_t)) as the default but provides no analysis of gradient scaling, eigenvalue growth of S^K_t, or the interaction between exponential decay Ξ³ and ridge regularization Ξ» in controlling activation magnitudes. A targeted study would train small HLA models (1β2 layers, synthetic data with controlled sequence length) while instrumenting: (a) the eigenvalue spectrum of S^K_t over the course of training and across sequence positions, (b) the distribution of numerator values num_t and denominator values den_t (for the normalized variant), (c) gradient norms through the HLA sublayer compared to adjacent feed-forward sublayers, and (d) the onset of loss spikes or NaN values as a function of Ξ³, Ξ», and sequence length. The goal is to produce a stability "phase diagram" showing the feasible region of (Ξ³, Ξ», sequence length) for stable training, and to determine whether the unnormalized form can be trained without normalization at practical scales, or whether ratio normalization is effectively mandatory. This study would transform HLA from an algebraic construction into a practically deployable component.
Third-order HLA at reduced state cost via structured key-moment approximations. The paper's third-order construction is likely unimplementable at typical model scales due to the O(dΒ³ dα΅₯) memory requirement for the linear maps M^{KQP} and M^{KQm} in the chunk-parallel scan. A follow-up could investigate whether structured approximations to S^K_t and S^Q_t β such as maintaining only the top-k eigenvalues and eigenvectors (using online PCA or incremental SVD), imposing a block-diagonal structure on the key and query moments, or using random projections (Johnson-Lindenstrauss style) β can reduce the effective dimension of the third-order state while preserving most of the representational benefit. The experiment would compare full third-order HLA (on a small-scale problem where dense materialization is feasible, e.g., d=32, dα΅₯=32) against various structured approximations, measuring both the fidelity of the approximated outputs (relative to the exact third-order computation) and the downstream task performance of models using the approximation. A finding that a low-rank approximation with rank r βͺ d recovers most of the third-order benefit would make third-order HLA practically viable; a finding that the approximation degrades rapidly would suggest that third-order interactions require the full d-dimensional state and may be inherently impractical at scale.
Combining symmetric HLA and asymmetric AHLA in a mixed-head architecture. The paper presents HLA and AHLA as alternatives but does not explore whether combining them within the same model is beneficial. A natural experiment would train models where some attention heads use symmetric HLA (key-centric aggregation via S^K_t) and others use asymmetric AHLA (query-centric routing via E_t), with the mix ratio as a hyperparameter. The hypothesis is that the two variants capture complementary interaction patterns: symmetric heads detect co-attention to shared keys (two queries focusing on the same past information), while asymmetric heads route information through intermediate query positions (chaining: current query β intermediate key β intermediate query β past key). A study could measure: (a) whether mixed-head models outperform pure-HLA or pure-AHLA models at the same parameter count, (b) whether the optimal mix ratio varies across layers (e.g., earlier layers favoring asymmetric routing, later layers favoring symmetric aggregation), and (c) whether the model learns to specialize heads (some becoming effectively first-order by learning identity-like key metrics, others using the full second-order capacity). This would provide practical guidance for architecture design and test whether the two variants are genuinely complementary or largely redundant.
Scaling the chunk-parallel training scheme to distributed settings and measuring wall-clock efficiency. The paper provides asymptotic work analysis but no implementation or profiling. A systems-oriented follow-up would implement the two-level chunk-parallel scan (Algorithms 1 and 4) in a framework like JAX or Triton, with careful attention to GPU memory coalescing, kernel fusion, and mixed-precision arithmetic. The study would measure: (a) wall-clock forward and backward pass time for HLA vs. standard softmax attention and linear attention at sequence lengths from 512 to 32K tokens, (b) peak GPU memory usage including the chunk summaries and scan workspace, (c) scaling efficiency when distributing the scan across multiple GPUs (model parallelism with chunk summaries communicated via NCCL), and (d) the sensitivity of throughput to chunk size C (sweeping from C=64 to C=2048) to identify the optimal balance between intra-chunk parallelism and inter-chunk communication overhead. The output would be a set of engineering guidelines β recommended chunk sizes for specific GPU architectures, memory budgets for the scan workspace, and throughput comparisons β that would make HLA adoption practical for practitioners. Without this, HLA remains a theoretical construction that is difficult to implement correctly and efficiently.
Extending the sufficient-statistics approach to other attention-like operators beyond polynomial forms. The paper's central insight β that certain attention weight matrices factorize through compact prefix summaries β may apply beyond the specific polynomial forms studied. A theoretical follow-up could investigate whether other attention variants that have been proposed in the literature admit similar factorizations. Candidates include: (a) attention with relative positional biases (e.g., ALiBi, RoPE) β can the positional terms be absorbed into the prefix summaries while maintaining causality?; (b) sparse attention patterns (e.g., sliding window, dilated attention) β can the sparsity pattern be expressed as a constraint on the prefix summaries rather than as a mask on the nΓn matrix?; (c) cross-attention between different sequences β does the factorization extend when queries and keys come from different sources? The contribution would be a taxonomy of attention operators classified by whether they admit exact streaming factorization, and for those that do, what prefix summaries are necessary and sufficient. This would systematize the design space and potentially subsume multiple proposed efficient attention variants under a single framework.
Practical Applications and Downstream Use Cases
Long-context language model deployment on memory-constrained devices. A setting where HLA's constant-size streaming state (O(dΒ² + d dα΅₯) per layer, independent of sequence length) provides a clear advantage over standard attention's O(n d) key-value cache is on-device autoregressive generation with very long contexts. Consider a 7B-parameter model with d = 128, dα΅₯ = 128, and 32 layers, deployed on a smartphone for a personal assistant that needs to reference a user's entire conversation history (potentially tens of thousands of tokens). With standard multi-head attention (32 heads), the KV cache at 32K tokens would require 32 layers Γ 32 heads Γ 32K tokens Γ 128 dims Γ 2 (K and V) Γ 2 bytes (FP16) β 536 MB β exceeding typical mobile memory budgets. With HLA using multi-query K/V sharing (Section 5.2), the per-layer state is dΒ² + 32 Γ d Γ dα΅₯ Γ 2 β 16K + 1M β 1M floats β 2 MB per layer, totaling 64 MB across 32 layers β independent of context length. If empirical benchmarking (not provided in this paper) confirms that HLA matches or approaches softmax attention perplexity, this memory reduction would enable genuinely long-context on-device LLMs that are currently infeasible. The constant state also eliminates the need for KV cache eviction policies, sliding windows, or context compression heuristics, simplifying deployment engineering.
Chunk-parallel training of recurrent models with exact long-range gradients. The associative scan operators in Section 4 provide a training recipe that is both parallel (O(log n) span) and exact (gradients match serial recurrence). This is directly applicable to training large recurrent or hybrid models (e.g., architectures combining HLA mixer layers with standard feed-forward blocks) on long documents (tens to hundreds of thousands of tokens) where full BPTT would be memory-prohibitive and truncated BPTT introduces bias. The chunk-parallel scheme with associativity guarantees means that the model sees exact gradients over the full sequence length during training β no truncation, no staleness β while exploiting GPU parallelism within and across chunks. For a practitioner training a document-level language model on 128K-token legal or scientific texts, adopting HLA as the mixer layer would allow gradient flow across the entire document without the quadratic cost of materializing a 128K Γ 128K attention matrix, and without the approximation error of chunked or truncated methods. The training FLOPs would scale as O(N C (d + dα΅₯) + M(dΒ² + d dα΅₯)) (Section 4.2), which with appropriately chosen chunk size C can be tuned to the specific GPU memory and compute budget.
Drop-in replacement for the attention sublayer in existing transformer codebases. Because HLA operates at the sublayer level β taking the same Q, K, V projections as standard attention and producing outputs that feed into the same residual connection and feed-forward block β it can be integrated into existing transformer implementations with minimal architectural changes. A practitioner with an existing GPT-style or BERT-style codebase can replace the attention kernel with an HLA implementation (following Algorithms 1β4) while keeping the rest of the model (embeddings, feed-forward layers, layer norm, positional encodings, training loop) identical. The paper's explicit support for multi-query K/V sharing (Section 5.2) and optional ratio normalization (Algorithm 1, lines 12β16) means that common practical configurations are supported. The main implementation burden is the associative scan infrastructure (the Blelloch scan, the semidirect product operator, the chunk-parallel orchestration), which the paper provides in pseudocode and which prior work (Qin et al., 2024; Yang et al., 2024b) has shown can be implemented efficiently. If follow-up work provides a reference implementation with benchmarks, HLA could become as straightforward to adopt as linear attention variants currently are β with the potential advantage of exact higher-order interactions.
Hybrid architectures combining HLA with retrieval or memory mechanisms. HLA's prefix summaries (S^K_t, C^{QV}_t, etc.) are interpretable intermediate states that could serve as interfaces to external memory or retrieval systems. For example, the key moment S^K_t represents the second-moment statistics of all keys seen so far β a compact summary of "what information has been encountered." A retrieval-augmented generation (RAG) system could use S^K_t as a query to fetch relevant documents from an external corpus (since S^K_t captures the distribution of key directions, documents with similar key statistics might contain related content), or could inject retrieved information by adding synthetic key-value pairs directly to the prefix summaries before computing the output. Because HLA's state is purely additive, such injection would be algebraically clean β simply add the retrieved keys' outer products to S^K_t and the retrieved values' contributions to C^{QV}_t β without needing to modify the core computation. This is less natural with standard attention, where the KV cache is a list of per-token vectors and injecting retrieved information requires managing positional indices and causal masking carefully. HLA's prefix-summary abstraction could thus serve as a unified interface between the sequence model and external knowledge sources, though this application depends entirely on empirical validation of HLA's task performance.
