PLASH: Provably Linear-Time Attention with Selective Higher-Order Feature Sketching

Abstract

Standard softmax attention scales quadratically with sequence length, which makes long-context training and inference expensive. We introduce PLASH, an attention block whose cost grows linearly in the number of keys. PLASH compresses the original keys and values into $M$ learned prototypes, where $M\in\mathbb{Z}_{>0}$ is much smaller than the number of keys. The compressed prototypes are then enriched with randomized polynomial features that recover inter-token information lost to compression. The output is computed by exact scaled dot-product softmax attention from each query to the enriched prototypes, so PLASH preserves the standard attention interface.The construction applies to self- and cross-attention. We prove sketch-error bounds for the enrichment step, a per-input certificate that upper-bounds the deviation from standard softmax attention on each forward pass, and a runtime bound linear in the number of queries and keys. Experiments on long-context language modeling (Qwen3-4B on PG-19) and time-series forecasting (ETT, ECL, Weather) show competitive accuracy and favorable scaling against efficient-attention baselines.