Title: Outrageously Large Context Windows via RACE Attention -- A Family of Non-Linear Attention that can be calculated in Strictly Linear-Time
Primary area: foundation or frontier models, including LLMs
Keywords: Sketching, Locality Sensitive Hashing, RACE, Attention, Linear, Transformers

Abstract: Quadratic attention strains memory and time, and even FlashAttention on a GH200 (96 GB) cannot complete a a single forward–backward step of a multi-head attention layer at sequence lengths over one million tokens. We introduce RACE Attention, a kernel-inspired alternative to softmax that runs in linear time in sequence length and embedding dimension. RACE builds on angular (cosine-based, scale-invariant) similarity, replacing the softmax exponential with powers of cosine similarity, and computes a differentiable, non-linear approximation via randomized projections and soft Locality Sensitive Hashing (LSH). The computation is carried out efficiently using Repeated Arrays-of-Count Estimators (RACE) sketches, yielding constant per-token working state. Across language modeling, masked language modeling, and classification, RACE Attention achieves performance comparable to state-of-the-art attention baselines while reducing runtime and memory. In a rigorous scale test, our proposal can successfully scale to 12 million tokens on a GH200 GPU and to 75 million tokens on an Intel Xeon® Gold 5220R CPU—well beyond the limits of highly optimized attention implementations. RACE offers a practical and theoretically grounded alternative to softmax attention for dealing with outrageously large context windows on today's infrastructure, we hope it gets adopted in practice and opens the door for orders of magnitude larger context windows than that of now.

TL;DR: We introduce an efficient way of doing attention by using LSH and notice that it scales up to 75 million tokens on CPU and 12 million tokens on GPU.