SLAY: Geometry-Aware Spherical Linearized Attention with Yat-Kernel
Jose Miguel Luna, Taha Bouhsine, Krzysztof Choromanski
TL;DR
SLAY tackles the quadratic time and memory bottleneck of softmax attention by enforcing unit-norm queries/keys and reexpressing the YAT-kernel as a positive mixture of polynomial and exponential kernels via Bernstein’s integral. The method uses Gauss–Laguerre quadrature and positive random features to obtain a strictly positive, tractable approximation, enabling $O(L)$ attention with stable normalization and bounded scores. The paper provides theoretical guarantees (positivity, boundedness, and PD properties) and extensive empirical evidence across synthetic tasks, extreme classification, and full Transformer models, showing SLAY's performance is near that of full softmax while significantly improving scalability over prior linear-time approaches. Practically, SLAY enables scalable long-context Transformers with geometry-aware interactions, reducing computational costs without sacrificing accuracy, and demonstrates strong performance gains in both language and vision-like benchmarks.
Abstract
We propose a new class of linear-time attention mechanisms based on a relaxed and computationally efficient formulation of the recently introduced E-Product, often referred to as the Yat-kernel (Bouhsine, 2025). The resulting interactions are geometry-aware and inspired by inverse-square interactions in physics. Our method, Spherical Linearized Attention with Yat Kernels (SLAY), constrains queries and keys to the unit sphere so that attention depends only on angular alignment. Using Bernstein's theorem, we express the spherical Yat-kernel as a nonnegative mixture of polynomial-exponential product kernels and derive a strictly positive random-feature approximation enabling linear-time O(L) attention. We establish positive definiteness and boundedness on the sphere and show that the estimator yields well-defined, nonnegative attention scores. Empirically, SLAY achieves performance that is nearly indistinguishable from standard softmax attention while retaining linear time and memory scaling, and consistently outperforms prior linear-time attention mechanisms such as Performers and Cosformers. To the best of our knowledge, SLAY represents the closest linear-time approximation to softmax attention reported to date, enabling scalable Transformers without the typical performance trade-offs of attention linearization.