SLAY: Geometry-Aware Spherical Linearized Attention with Yat-Kernel

AI summary

Overview

• Research area: Efficient Transformer attention mechanisms; kernel methods and random-feature approximation for linear-time sequence modeling.
• Technical level: Advanced. The paper combines kernel theory (Bernstein's theorem, complete monotonicity, positive definite kernels on spheres), random-feature approximation theory, numerical quadrature, and empirical Transformer training.
• Scope: The paper derives, implements, and empirically validates SLAY, a linear-time attention mechanism built on a spherical version of the Yat-kernel (the ⵟ-Product from Neural Matter Networks), and benchmarks it against softmax attention, Performers, Cosformers, and other linear baselines.

What This Paper Is About

Standard Transformer attention builds an explicit L×L matrix of pairwise query–key similarities, costing quadratic time and memory in sequence length L. The Yat-kernel (ⵟ-Product) is a geometry-aware similarity that couples alignment and proximity, but its denominator entangles query and key terms, so it cannot be factorized and therefore also costs quadratic time. This paper shows that by constraining queries and keys to the unit sphere and rewriting the kernel through Bernstein's theorem, the Yat-kernel becomes expressible as a nonnegative mixture of factorizable product kernels that can be approximated with strictly positive random features — yielding linear-time attention.

Key Contributions

1. Spherical reformulation of the Yat-kernel. Normalizing queries and keys to unit norm reduces the kernel to a function of angular alignment alone: ⵟ_sph(q̂, k̂) = x²/(C − 2x), where x = q̂ᵀk̂ ∈ [−1, 1] and C = 2 + ε. This is interpreted as an ε-regularized squared chordal distance interaction on the sphere S^(d−1).

2. Integral linearization via Bernstein's theorem. Since 1/y is completely monotone on (0, ∞), the factor 1/(C − 2x) is expressed as a Laplace integral, rewriting the spherical kernel as an integral over s of e^(−sC)·[x²e^(2sx)] ds — a positively weighted mixture of a degree-2 polynomial kernel (q̂ᵀk̂)² and an exponential dot-product kernel e^(2s q̂ᵀk̂).

3. Strictly positive random-feature estimator. The integral is discretized with R-point Gauss–Laguerre quadrature, the polynomial factor is approximated with a nonnegative feature map (anchor features by default), and the exponential factor is approximated with positive random features (PRFs) from Choromanski et al. (2021). The two are fused via a sketching operator, giving a mechanism that produces well-defined, nonnegative attention scores with linear-time O(L) scaling.

4. Empirical demonstration at scale. Across polynomial-approximation ablations, scaling benchmarks, 22 synthetic tasks, extreme classification on Eurlex, and full GPT-2 Small-scale training runs under Chinchilla-budget-matched conditions, SLAY is reported as the closest linear-time approximation to softmax attention to date, consistently outperforming Performers, Cosformers, and Linear (ELU+1) attention.

Main Findings

• Polynomial factor ablation (Table 2): Anchor features achieve the lowest error among practical approximations, with relative L2 error 0.527, cosine similarity 0.850, MSE 4.55e-03, and forward-pass latency 489.42 ms. Laplace-only is comparably accurate (rel. L2 0.544, cos 0.839, MSE 4.84e-03) but much slower at 1905.80 ms. Signed approximations degrade severely: Nystrom reaches rel. L2 70.291, TensorSketch 24823.685, and Random Maclaurin 15826.841, which the authors attribute to negative approximate inner products causing denominator cancellation and instability.

• Nonnegativity constraint matters: Table 1 classifies approximation options by whether ⟨φ(x), φ(y)⟩ ≥ 0 is guaranteed. Anchor features (P dimensions, O(dP) cost) and the exact vec(uuᵀ) map (d² dimensions, O(d²) cost) preserve nonnegativity; TensorSketch, Random Maclaurin, and Nystrom do not.

• Scaling behavior: In an isolated causal attention benchmark with embedding dimension 256, 8 heads, batch size 1, on a single NVIDIA A100-SXM4 GPU (80 GB), quadratic mechanisms (standard softmax and exact YAT) fail beyond 16K tokens due to out-of-memory, while linear methods remain stable up to 128K tokens. SLAY closely tracks other linear mechanisms and uses orders of magnitude less memory than exact methods at long lengths.

• Synthetic tasks (Table 3): Across 22 synthetic tasks, SLAY scores 0.57 on Basic, 0.57 on Arithmetic, 0.68 on Long-Range, 0.73 on Memory, 0.86 on Patterns, 0.57 on Reasoning, and 0.80 on Robustness. It matches or exceeds other linear baselines and closes part of the gap to standard softmax on Basic operations (standard: 0.60).

• Extreme classification on Eurlex (Table 4): SLAY outperforms Performer FAVOR+ on every reported metric: P@1 0.4978 vs. 0.3442, P@3 0.3953 vs. 0.2703, P@5 0.3261 vs. 0.2263, PSP@1 0.9391 vs. 0.5970, PSP@3 0.7862 vs. 0.4785, PSP@5 0.6693 vs. 0.4141.

• Full-scale language modeling (Table 5): With identical architecture, optimization, and data, at 125M parameters and 2.5B tokens under the Chinchilla scaling law, SLAYformer (O(n)) reaches validation loss 4.6760 and perplexity 107.35, close to Standard Softmax (4.6417 / 103.73) and Yat (Spherical) (4.7180 / 112.00), and ahead of Linear ELU+1 (5.0884 / 161.99), Cosformer (5.1983 / 180.97), and FAVOR+ Performer (5.4524 / 233.32). Exact Yat achieves the best numbers overall (4.5747 / 97.03).

• Training dynamics (Figure 3): The gap between SLAY and softmax stays small throughout the full optimization trajectory, while the gap between SLAY and other linear methods stays consistently large, indicating the advantage is not a transient effect of early optimization.

• Claimed context reach: The conclusion states SLAY can process sequences up to 30× longer than standard softmax attention.

Methodology in Plain English

The starting point is a similarity function borrowed from physics-style inverse-square interactions: the squared inner product of a query and a key, divided by the squared distance between them plus a small stabilizer. This has the appealing property that it rewards vectors that are both well-aligned and close together, but the distance term mixes the query and key inside a denominator, which blocks the algebraic reordering that makes linear attention possible.

The authors fix this by forcing every query and key onto the unit sphere. Once both vectors have length one, the denominator collapses to a simple linear function of the single number x = q̂ᵀk̂, so the whole kernel becomes x²/(C − 2x). They then use a classical result (Bernstein's theorem on completely monotone functions) to rewrite the reciprocal 1/(C − 2x) as an integral of decaying exponentials. This turns a non-factorizable kernel into a weighted average of terms that each consist of a degree-2 polynomial kernel times an exponential dot-product kernel — and each of those pieces can be factorized.

Two approximation steps follow. The integral is replaced by a finite Gauss–Laguerre quadrature with R nodes. For each node, the polynomial piece is approximated with a nonnegative feature map (anchor features are the default, chosen for nonnegativity and low latency), and the exponential piece is approximated with positive random features — exponentials of random projections — which are known to give nonnegative, unbiased estimates of exponential kernels. The two feature maps are combined by a sketching operator so that the full tensor product is never materialized. Finally, standard linear-attention reordering is applied: instead of building an L×L attention matrix, the mechanism computes Ψ(K)ᵀV and Ψ(K)ᵀ1 first and then multiplies by Ψ(Q), giving O(L) time and memory.

Evaluation proceeds in five stages: isolate and compare polynomial-approximation options with matched feature budgets; measure latency, peak memory, and throughput versus sequence length; run 22 synthetic tasks; run an extreme classification benchmark on Eurlex; and train full GPT-2 Small-scale Transformer language models that differ only in their attention mechanism.

Why This Matters

Research impact. The paper attacks the long-standing quality gap between quadratic softmax attention and linear-time alternatives. Its central move — using a geometric kernel on the sphere plus Bernstein's theorem to make the kernel factorizable while preserving nonnegativity — is a template that could be reused for other non-factorizable kernels, not just the Yat-kernel. It also connects efficient attention back to classical kernel theory (isotropic spherical kernels, positive definite functions, random features), positioning "geometry-aware" kernel design as a research direction rather than a one-off trick. The reported result that a linear-time mechanism stays within a narrow margin of softmax in a fully trained language model, if it holds up, is a meaningful data point for the efficient-attention literature.

Real-world applications.

• Long-context document processing: summarizing, searching, or question-answering over book-length or repository-length inputs where quadratic attention runs out of memory.
• Retrieval-augmented generation: attending over very large retrieved evidence sets without the memory blow-up of explicit attention matrices.
• Domain-specific corpora such as legal, medical, or scientific literature, where long inputs are the norm and where the kernel's geometry - favoring tokens that are both aligned and close in representation space - may aid precision.
• Long-sequence multimodal inputs such as video or long audio, where sequence lengths routinely exceed the 16K-token regime where the paper's benchmarks showed quadratic attention failing.

Industry relevance. If a linear-time mechanism can match softmax quality, the practical payoff is lower GPU memory and compute per token at long context, which translates directly into cheaper training and serving of long-context models. The paper also frames this in terms of reduced energy cost for training and deployment, and it reports a publicly available code repository for reproducibility.

Future Directions

• Closing the remaining gap to softmax: SLAY's validation loss (4.6760) and perplexity (107.35) remain slightly above standard softmax (4.6417 / 103.73) and exact Yat (4.5747 / 97.03); narrowing this further, and understanding whether the residual gap is from quadrature node count, feature budget, or the sketching step, is a natural next step.
• Scaling beyond GPT-2 Small scale: The reported full-model experiments use 125M parameters and 2.5B tokens; whether the near-softmax behavior persists at substantially larger parameter counts and token budgets is not established by the paper.
• Improving or replacing the polynomial approximation: The default anchor features introduce bias (Table 1 marks them as not unbiased), and signed approximations such as TensorSketch and Random Maclaurin proved numerically unstable at the tested budgets. Finding unbiased, nonnegative, and cheap alternatives remains open.
• Extending the framework to other kernels: The Bernstein/Laplace linearization recipe is presented specifically for the spherical Yat-kernel; whether the same approach yields positive, factorizable estimators for other non-factorizable geometric kernels is an open question.

Target Audience

This paper is most useful to machine learning researchers and engineers working on efficient attention, long-context Transformers, and kernel methods. It suits readers already comfortable with attention mechanics, random-feature approximations, and basic kernel theory; readers without that background will find the theoretical sections (Bernstein's theorem, RKHS tensor products, quadrature bounds) demanding. Practitioners evaluating whether to swap softmax attention for a linear alternative in a long-context system will find the scaling benchmarks and the GPT-2 Small-scale comparison tables the most directly actionable parts, while theoreticians will find the spherical positive-definiteness and nonnegativity arguments the most novel.

Authors’ 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.