ZeroS: Zero-Sum Linear Attention for Efficient Transformers

AI summary

Overview

Research area: Efficient sequence modeling — specifically linear-time attention mechanisms for Transformers, evaluated on in-context learning, language modeling, image classification, and time series forecasting.

Technical level: Advanced. The paper is built on Taylor expansions of softmax, convex/affine hull arguments, Lipschitz bounds, and prefix-sum recurrences.

Scope (one sentence): The paper proposes ZeroS, a linear-time attention mechanism that removes the constant zero-order term from the softmax expansion to create zero-sum (positive and negative) attention weights, and shows it matches or exceeds standard softmax attention across several benchmarks while retaining O(N) complexity.

What This Paper Is About

Linear attention methods reduce the quadratic O(N²) cost of standard softmax attention to O(N), but they consistently underperform it. The authors argue this gap has two causes: attention weights are restricted to convex combinations (only additive blending, no contrastive/subtractive operations), and a roughly uniform 1/N component in the weight expansion dilutes attention in long contexts. Their goal is to remove that uniform term entirely, producing stable, signed, zero-sum weights that let a single attention layer perform contrastive operations without losing linear-time efficiency.

Key Contributions

1. Diagnosis of the zero-order term. The authors identify the constant 1/t term in the softmax Taylor expansion as a source of average-pooling bias and attention dilution, and demonstrate that removing it is both safe and beneficial.

2. The ZeroS mechanism. They develop Zero-Sum Linear Attention, a linear-time attention layer supporting negative weights, built from zero-order subtraction, radial–angular decoupling (separating magnitude from direction with learned gates and reintroducing signed cos θ), and a prefix-sum implementation running in O(Nd²) time and O(d²) memory.

3. Theoretical expressivity and stability results. Proposition 3.1 and Corollary 3.2 show the zero-sum span strictly contains the deviation set achievable by convex weights; Proposition 3.3 shows a first full-attention layer plus stacked zero-sum layers spans exactly the affine hull of the value vectors; Lemma 3.4 shows output norm is O(B), independent of sequence length t; and Proposition 3.5 gives an O(1/√t) Lipschitz constant.

4. Empirical validation. ZeroS and its quadratic variant ZeroS-SM were evaluated on MAD, MQAR, RegBench, WikiText-103, OpenWebText2, ImageNet-1k (DeiT-Tiny), and a time series forecasting benchmark, by directly substituting the multi-head attention module while keeping all other components fixed.

Main Findings

• Two named limitations motivate the method. The paper attributes linear attention's weakness to (1) the bottleneck of convex combination, where weights must be non-negative and can only blend information additively, and (2) uniform weight bias and attention dilution, where a roughly uniform 1/N component weakens focused attention as context grows.

• Zero-sum weights widen the representable set. Proposition 3.1 states that the zero-sum span of {v_i − v_avg} strictly contains the deviations achievable by convex weights, with strictness whenever the v_i are not all identical. Corollary 3.2 frames this as an expressivity gain in a residual block x_t ↦ x_t + Σ w_i v_i.

• Expressivity is recovered by layer stacking. Proposition 3.3 shows a head retaining the zero-order term spans the affine hull Aff{v_i}, a head without it spans Span{Δ_i}, and the Minkowski sum Aff{v_i} + Span{Δ_i} equals Aff{v_i} — so the first full attention layer already covers the affine hull and later zero-sum layers do not shrink it.

• Stability is preserved without convexity. Lemma 3.4 bounds ‖Σ w_{t,i} v_i‖ ≤ B·t·max_i|w_{t,i}| and, with max_i|w_{t,i}| = O(1/t), gives O(B) independent of sequence length. Proposition 3.5 gives a uniform Lipschitz bound of B·L_w/√t, an O(1/√t) decay that the authors note is largely superseded by Layer Normalization in practice.

• MAD benchmark. ZeroS averaged 73.5, ahead of the other linear-time models (Hyena 66.0, MultiHead Hyena 73.2, Mamba 69.3, GLA 60.0, DeltaNet 71.8, LinAttn 62.3) and close to Transformer's 74.5. ZeroS-SM reached 75.4, above Transformer. ZeroS notably scored 99.9 on In-Context Recall and 96.1 on Noisy Recall versus LinAttn's 91.0 and 75.6.

• Differential gains are task-dependent. The authors report minimal gains on Compress and Memorize, tasks less reliant on complex representations, and note ZeroS maintains strong memorization (88.1) despite using negative weights, unlike DeltaNet (52.8).

• WikiText-103. ZeroS reached 23.91 validation and 24.61 test perplexity with 46.31M parameters, beating vanilla Transformer (24.4 / 24.78, 44.65M). ZeroS-SM reached 23.62 / 24.17 with 44.69M parameters. Mamba (22.58 / 23.19) and HGRN2 (23.1 / 23.73) still achieve lower perplexity.

• Ablation on WikiText-103. Reintroducing the 0-th order term worsened perplexity to 24.05 / 24.74 and removing the reweighted zero-sum softmax worsened it to 24.21 / 24.97, versus 23.91 / 24.61 for ZeroS and 23.62 / 24.17 for ZeroS-SM. The paper reports that reintroducing the 0-th order term reduces performance on In-Context Recall, Noisy Recall, and WikiText.

• Image modeling. With encoder-only ZeroS replacing DeiT-Tiny's softmax attention, ZeroS reached 75.51 top-1 accuracy on ImageNet-1k with 6.0M parameters, above HGRN2 (75.39, 6.1M), HGRN1 (74.40, 6.1M), TNN (72.29, 6.4M), and DeiT (72.20, 5.7M).

• Time series forecasting. On Weather, Solar, ETTh1, ETTh2, ETTm1, and ETTm2, ZeroS recorded lower MSE than GLA, iTransformer, and DLinear on the reported entries — for example Weather MSE 0.218 and Solar MSE 0.192 — and the authors state it outperforms GLA, AFT, iTransformer, and PatchTST on most datasets. PatchTST is better on some entries (ETTh2 MSE 0.330 vs. ZeroS 0.341; Weather MAE 0.261 vs. 0.265) and AFT is better on ETTh1 (0.409 vs. 0.414).

• MQAR and RegBench. ZeroS performs comparably to vanilla attention across most MQAR configurations (Figure 3 plots model dimension against accuracy, with no numeric table reported), and on RegBench it outperforms linear-time baselines including GLA, RetNet, and RWKV (Figure 2, no numeric table in the text).

• OpenWebText2. Using a 12-layer, 768-dimensional GPT-2 architecture, ZeroS tracks much closer to vanilla Transformer than AFT and GLA, with ZeroS-SM further improving on vanilla Transformer.

• Deviations from past work are deliberate. Unlike prior Taylor-based approximations that require O(d³) for second-order terms, ZeroS uses logits dependent only on step i, computes full softmax, zero and first-order terms, and derives higher-order terms by differentiation, with t-dependent gates supplying (t, i) interaction.

Methodology in Plain English

The authors start from the Taylor expansion of the attention weight exp(q·k) = 1 + ⟨q,k⟩ + ½⟨q,k⟩² + …, where the leading "1" is what makes weights non-negative and stable but also forces uniform averaging. Instead of designing a more complicated kernel to approximate softmax while keeping that constant, they subtract it. Subtracting it from each softmax row gives weights that sum to zero, which are allowed to be negative and therefore can express contrasts between tokens rather than only blends.

To make this computationally practical, they split the weight into a radial part (magnitude, computed from logits that depend only on position i) and an angular part (direction, a signed cos θ between normalized query and key, optionally rotated by RoPE). The radial component is decomposed into the full softmax term, the 0th-order baseline 1/t, and the 1st-order term; the residual above first order is what captures higher-order competition. Learned per-step gates σ¹_t and σ^h_t (and optionally σ⁰_t in the first layer) decide how much of each order to keep — this is the "reweighted zero-sum softmax." The deviation logits s_i are computed as the negative scaled inner product between each token's projected vector u_i and its cumulative average ū_i, smoothed by trainable parameters μ and τ, so that each position is scored by how much it deviates from prior context.

Because everything depends only on i for the logits and only on t for the gates, all sums become prefix scans: E_t = Σ e^{s_i}, P_t = Σ s_i, and the matrix accumulations F_t, G_t, H_t. The output is a gated combination of these scans multiplied by the current angular vector, requiring only O(d²) state and O(d²) updates per step, for O(Nd²) total time and O(d²) memory. The same reweighting trick can also be applied to standard softmax attention, yielding the quadratic-time variant ZeroS-SM. In experiments, the multi-head attention module is swapped out and everything else — MLP/GLU, embeddings, hyperparameters — is held fixed.

Why This Matters

Impact on research. The paper reframes why linear attention underperforms: not a failure to approximate softmax well enough, but an inherited constraint from convexity. It offers a constructive alternative — remove the zero-order term — supported by hull-based expressivity arguments and Lipschitz stability bounds, and it works on both linear and standard softmax attention. This gives the efficient-attention literature a new lever beyond kernel design, gating, and hybrid local windows, and it connects to prior work on negative weights, differential attention, and memory-deletion mechanisms like DeltaNet.

Real-world applications.

• Long-context language modeling and retrieval, where the O(N²) attention cost is prohibitive.
• Time series forecasting for weather, solar, and electricity load, where ZeroS achieved lower MSE than GLA, iTransformer, and DLinear on the reported datasets.
• Vision backbones at small parameter budgets, where ZeroS improved DeiT-Tiny top-1 accuracy on ImageNet-1k.
• In-context and associative recall systems (e.g., MQAR-style induction heads), where signed contrastive weights improved In-Context and Noisy Recall over LinAttn on MAD.

Industry relevance. ZeroS is a drop-in attention replacement with demonstrated scaling behavior in the 44–48M parameter range and a reported linear O(Nd²) runtime with O(d²) memory — the profile that matters for inference serving on long sequences. Its stability guarantees that are independent of sequence length (O(B) output norm, O(1/√t) Lipschitz constant) target exactly the numerical concerns that block deployment of signed-weight attention in production training. The authors are affiliated with Georgia Institute of Technology and Amazon Web Services, and code is available at a link given in the paper.

Future Directions

• The 0th-order trade-off is unresolved. The paper keeps the zero-order term optionally in the first layer to avoid losing the v_avg direction, but reports that reintroducing it hurts In-Context Recall, Noisy Recall, and WikiText. Determining when a first-layer 0th-order restore is worth its cost remains an open question.

• Scaling beyond the studied sizes. Experiments cover roughly 44–48M-parameter language models, 6.0M-parameter DeiT-Tiny, and a 12-layer, 768-dimensional GPT-2. Whether ZeroS's advantage holds at much larger parameter counts and context lengths is not reported.

• Comparison against the strongest perplexity baselines. On WikiText-103, Mamba (23.19 test) and HGRN2 (23.73 test) still beat ZeroS (24.61) and even ZeroS-SM (24.17); closing that gap, or understanding why recurrences win on this task, is a natural next step.

• Layer Normalization versus the √t decay. Proposition 3.5's O(1/√t) Lipschitz decay is described as largely superseded by LayerNorm. Whether the decay factor can be dropped entirely, or should be combined with normalization, is left open.

• Generalization of the reweighted zero-sum softmax. The paper notes the reweighting can be applied directly to standard softmax attention (Section A.1.6), suggesting further exploration of quadratically-attended variants is intended.

Target Audience

Researchers and engineers working on efficient Transformer architectures, long-context sequence modeling, and kernel-based or recurrent linear attention. The paper assumes comfort with attention formulations, Taylor expansions, convex/affine hull arguments, and Lipschitz analysis, so the theoretical sections are best suited to readers with a strong machine learning or optimization background. Practitioners looking for a drop-in replacement for multi-head attention in 40–50M-parameter models will find the experimental sections and ablation directly actionable, while those interested in the theory of what linear attention can and cannot represent will find Propositions 3.1 and 3.3 the most distinctive contribution.

Authors’ abstract

Linear attention methods offer Transformers $O(N)$ complexity but typically underperform standard softmax attention. We identify two fundamental limitations affecting these approaches: the restriction to convex combinations that only permits additive information blending, and uniform accumulated weight bias that dilutes attention in long contexts. We propose Zero-Sum Linear Attention (ZeroS), which addresses these limitations by removing the constant zero-order term $1/t$ and reweighting the remaining zero-sum softmax residuals. This modification creates mathematically stable weights, enabling both positive and negative values and allowing a single attention layer to perform contrastive operations. While maintaining $O(N)$ complexity, ZeroS theoretically expands the set of representable functions compared to convex combinations. Empirically, it matches or exceeds standard softmax attention across various sequence modeling benchmarks.