Higher-order Linear Attention

1 Introduction

• HLA uses compact prefix summaries to incorporate higher interactions while enabling exact causal streaming without constructing attention matrices.In the second-order case, it maintains a constant-size state per head and computes each token in linear time.Source: §1 · p. 1
• The second-order masked formulation adds two summaries to enforce strict autoregressive causality while retaining streaming updates and avoiding n×n intermediates.Its per-token cost remains O(d^2 + d dv).Source: §1 · p. 2
• Associative scans for unmasked and masked settings reproduce serial-recurrence forward activations exactly while enabling intra- and inter-chunk parallelism.The scan formulation also covers settings with and without exponential decay.Source: §1 · p. 2
• The third-order extension provides masked streaming updates and an exact chunk-parallel algorithm whose associative composition reproduces the serial recurrence.Source: §1 · p. 2
• HLA is intended as a drop-in, attention-like mixer for long-context models, combining data-dependent weighting, constant per-token streaming memory, and scan-based parallel training.The stated aim is to avoid approximate backpropagation through time.Source: §1 · p. 2

2 Background

• The notation uses row-vector outputs and distinguishes unnormalized HLA outputs from a ratio-normalized variant that divides by a masked scalar denominator.Prefix summaries are computable in streaming fashion with O(1) memory per token and per head.Source: §2 · p. 2

2.1 Scaled dot-product attention

• Scaled dot-product attention is introduced with queries, keys, and values; the causal mask retains positions on and before the current token.Source: §2.1 · pp. 2–3

2.2 Linear attention

• Linear attention approximates the softmax kernel with a feature map, and maintaining running key-feature sums yields O(n r(d+dv)) time and O(r dv) memory.Source: §2.2 · p. 3

3 Higher-order Linear Attention

• Second-order tensor attention depends on the key second moment K^T K, and prefix summaries make these moments streamable.The summaries are updated online, with per-token costs O(d^2) for the key moment and O(d d_v) for the value accumulator.Source: §3 · p. 3
• Unnormalized HLA computes a numerator-style bilinear form from prefix moments, avoiding length-dependent renormalization while retaining the normalized variant’s state.Optional normalization divides by a masked denominator.Source: §3 · p. 3, §3.1 · p. 6

3.1 Causal masking via extended summaries

• The masked streaming identity uses two additional prefix summaries to enforce strict causality, while the per-token cost remains O(d^2 + d d_v).The binary causal mask restricts attention to positions j ≤ t.Source: §3.1 · pp. 5–6

4.1 Unmasked monoid

• The unmasked monoid adds token-level state deltas, with inclusive states formed by combining each exclusive prefix with that token’s deltas.Adjacent segments are composed in time order.Source: §4.1 · p. 7

4.2 Masked semidirect product

• The masked semidirect product adds cross-terms to the state composition, and HLA uses the same scan infrastructure as first-order linear attention with a changed state tuple.For the second-order masked state, the tuple includes S, C, m, G, and h.Source: §4.2 · pp. 7–8

4 Chunk-parallel training via associative scans

• Associative scans reproduce serial per-token masked outputs exactly, and a reverse-mode scan yields gradients matching the serial recurrence.The equivalence applies with or without exponential decay.Source: §4.2 · pp. 7–8
• Within each chunk, an exclusive Blelloch scan computes prefixes and local updates form inclusive states; chunk summaries are then scanned to provide carry-in prefixes.Within-chunk scans have O(log w) span, and the same operator supports inter-chunk scans.Source: §4 · p. 7, §4.2 · p. 8

4.3 Adding decay and regularization

• With fixed exponential decay, segment summaries carry an attenuation factor, and decay controls spectral growth and improves recency bias while preserving associativity.The decay-aware updates are the decayed analogues of the online updates.Source: §4.2 · p. 7, §4.3 · p. 8

5 Implementation details and complexity

• Chunk token deltas feed an exclusive scan; each token then forms its inclusive state and computes its output from the resulting summaries.The pseudocode supports optional decay, ridge stabilization, and normalization.Source: §5 · p. 9

5.1 Pseudocode

• Optional normalization divides by the masked denominator, while adding λI stabilizes the causal variant but does not yield the exact masked bilinear form.The default implementation returns unnormalized output.Source: §5.1 · p. 9

5.2 Implementation considerations

• Drop-in use swaps only the attention kernel; multi-query sharing reduces state from O(h d^2) to O(d^2 + h d d_v), while packed symmetric storage further cuts bandwidth.The feed-forward and normalization sublayers, positional encodings, and masking remain unchanged.Source: §5.2 · pp. 9–10

6 Asymmetric Higher-order Linear Attention

• AHLA complements symmetric second-order HLA with a left-cascaded product and strictly causal streaming at O(d^2+d dv) per-token cost.Source: §6 · p. 10

6.1 Definition and masked streaming identity

• The masked streaming identity computes strictly causal outputs online, using O(d dv) time and O(d dv+d) state per head; chunk summaries additionally use O(d^2) memory.Source: §6.1 · pp. 10, 12

6.2 Chunk-parallel scans for AHLA

• An exclusive associative scan followed by local inclusion reproduces the serial recurrence's activations exactly.Source: §6.2 · p. 13

6.3 Pseudocode

• The pseudocode specifies masked streaming with optional decay and denominator normalization.Source: §6.3 · p. 13

7 Third-Order Linear Attention

• The section introduces third-order HLA as an extension of the earlier attention formulation.Source: §7 · p. 13

7.1 Streaming form of Causal HLA

• The masked streaming identity computes strictly causal third-order outputs from prefix summaries, with an optional normalized variant.Source: §7.1 · p. 14

7.2 Pseudocode

• The pseudocode presents a strictly causal streaming kernel for inference and an associative scan operator for chunk-parallel training.Source: §7.2 · p. 15

7.3 Chunk-parallel algorithm for third-order HLA

• The third-order scan is associative and, for γ = 1, produces outputs matching the serial algorithm; its exact chunk composition additionally stores segment maps costing O(d^3 dv) and O(d^3) entries when dense.Source: §7.3 · p. 17

• Fast-weight work connects associative outer-product memories to attention, while HLA directly uses higher sufficient statistics without explicit matrix inversion.Source: §8 · p. 18
• Feature-map linear attention methods typically maintain first-order statistics, whereas HLA maintains higher-order key moments and masked cross-summaries for strictly causal interactions.Source: §8 · p. 18
• SSMs and modern RNNs use recurrent updates but express data-dependent mixing differently or typically maintain first-order statistics, unlike HLA's attention-style queries and keys.Source: §8 · p. 18
• Test-time adaptation, external memory, and Hopfield views are described as distinct from HLA's compact prefix moments and exact masked streaming.Source: §8 · p. 19

• HLA combines exact streaming updates, strict masked causality, and associative scans that match serial recurrences.Source: §9 · p. 19
• At second order, HLA uses O(d^2) state per head and O(d^2) time per token, with optional normalization and decay preserving associativity.Source: §9 · p. 19
• The paper also develops an asymmetric variant and a masked third-order algebra with streaming formulas and an exact chunk-parallel scan.Source: §9 · p. 19