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

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.

AI systems that read and write text compare every word in a passage with every other word. For a long document this resembles asking each person in a huge crowd to shake hands with everyone else: the effort explodes as the crowd grows, making long inputs slow and costly to process. Our method, PLASH, takes a shortcut. It condenses the text into a small set of representative summaries, enriches them so they keep the important detail, and lets each word consult only those summaries instead of the whole crowd. The effort then grows in proportion to the text length, not its square. For every input, PLASH also computes a guarantee: a number showing how far its fast answer can stray from the exact one, so users know when to trust it. In tests on long books and real forecasting data, PLASH ran more than twice as fast as the leading method while matching its accuracy, and its guarantee held on all 35,000 real test cases. On long inputs, it makes AI cheaper, faster, and more reliable.