ZeroS: Zero-Sum Linear Attention for Efficient Transformers

TL;DR

ZeroS introduces a novel linear attention mechanism that overcomes fundamental limitations of existing methods, enabling more expressive and stable attention with maintained efficiency, and achieves comparable or better results on sequence benchmarks.

Contribution

ZeroS removes the zero-order term and reweights residuals in linear attention, expanding representational capacity and improving stability while preserving $O(N)$ complexity.

Findings

ZeroS matches or exceeds softmax attention on benchmarks.

It enables contrastive operations within a single attention layer.

ZeroS maintains linear complexity with enhanced expressiveness.

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.