PaperlayerView exact revision on arXiv

Source-linked AI summary

Quasi Linear Kernel Attention with Infinite Capacity

Nicolaj Rux, Johannes Hertrich, Sebastian Neumayer

arXiv:2609.35349v1cs.LGcs.NAmath.NA

TL;DR

Transformers’ softmax attention has quadratic sequence-length cost, while finite-feature-map approaches can limit expressivity. The paper defines capacity, proves infinite capacity for additive kernels that admit quasi-linear sorting computation, and reports speed advantages over softmax implementations for long sequences.

  • Problem

    Softmax attention scales quadratically with sequence length, while finite-feature-map kernels can have limited expressivity and incur costs that grow with feature dimension.

    Source: §1 · p. 1

  • Method

    The paper defines capacity through identity approximation and constructs additive kernels from univariate kernels whose sums can be evaluated in quasi-linear time by sorting.

    Source: §1 · p. 2, §3.2 · p. 5, §3.3 · p. 6

  • Results

    N ≳2,048: the CUDA implementation is faster than conventional single-precision PyTorch attention, with break-even against half-precision FlashAttention at N ≳20,000.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Additive Laplace and bump kernels provide infinite capacity with quasi-linear attention computation.

    Source: §3.3 · p. 6

    Capacity captures only one aspect of expressiveness; transformer expressivity also depends on properties such as stationarity and locality.

    Source: §5 · p. 10

Abstract

from arXiv · show

The evaluation cost of transformers with softmax attention scales quadratically with sequence length. Kernel attention addresses this by replacing softmax with a more general kernel function. In this paper, we aim to identify kernels that retain the expressivity of attention while enabling quasi linear computation. To quantify expressivity, we introduce a capacity for each kernel, measuring the maximum sequence length for which the attention matrix can approximate the identity. A higher capacity thus indicates greater expressivity. We show that expressive kernels like softmax, Gauss, and Laplace have infinite capacity. In contrast, common quasi linear kernels, such as those derived from finite dimensional feature maps, exhibit finite capacity. As a solution, we propose additive kernels constructed from univariate spline and polynomial exponential kernels. We prove that these maintain infinite capacity while allowing quasi linear computation via sorting. Finally, we implement additive sorting kernels efficiently and benchmark them against modern softmax backends, demonstrating advantages for long sequences.

1 INTRODUCTION

The paper seeks attention kernels that retain expressivity while allowing quasi-linear computation, and introduces capacity to measure expressivity through identity approximation. It proposes additive kernels with infinite capacity and quasi-linear evaluation, with implementations benchmarked against softmax attention.

1 INTRODUCTION

  • The softmax attention computation scales quadratically with sequence length, creating a practical bottleneck as context windows reach 128k tokens and beyond.The paper distinguishes its focus on exact subquadratic attention from approximation and windowed-pattern approaches.Source: §1 · p. 1
  • Univariate Laplace and piecewise-linear kernels support quasi-linear sorting algorithms, but most such algorithms do not generalize to higher dimensions; arbitrary-dimensional Laplace evaluation has exponential dependence on D and is impractical above D > 6.The cited examples include Laplace and piecewise-linear kernels [Hertrich, 2024; Hertrich et al., 2024; Vialard & Boufad`ene, 2025].Source: §1 · pp. 1–2
  • Finite-feature-map kernels admit linear-time evaluation, but their capacity is capped by feature count L, and their evaluation cost also grows with L.The introduction relates this limit to previously reported expressivity constraints compared with softmax attention [Schlag et al., 2021].Source: §1 · pp. 1–2
  • Capacity measures the maximum sequence length for which attention can approximate the identity, with higher capacity indicating greater expressivity.The paper proves infinite capacity for softmax, Gauss, and Laplace kernels, while finite-feature-map kernels have capacity at most L.Source: §1 · p. 2
  • N ≳2,048: the CUDA implementation is faster than conventional single-precision PyTorch attention, and it reaches break-even with half-precision FlashAttention at N ≳20,000.The advantage over FlashAttention increases rapidly with sequence length.Source: §1 · p. 2
  • Additive kernels formed from univariate Laplace or bump kernels have infinite capacity and support quasi-linear attention computation.The paper reports numerical experiments and an efficient CUDA implementation of these kernels.Source: §1 · p. 2

2 CAPACITY OF KERNEL ATTENTION

The section defines kernel capacity by whether attention can approximate an identity matrix at a given sequence length. It establishes infinite capacity for stationary C0 kernels and softmax, while bounding capacity for spd kernels by their RKHS dimension.

2 CAPACITY OF KERNEL ATTENTION

  • Capacity is the largest sequence length N for which some queries and keys make the attention matrix approximate the N × N identity; larger capacity indicates greater expressivity.Above capacity, the attention mechanism cannot represent the identity.Source: §2 · pp. 2–3
  • Stationary C0 kernels with F(0) > 0 have infinite capacity, including Gauss, Laplace, and Matérn kernels; softmax also has infinite capacity for D ≥ 2.The stationary-kernel condition requires F to be continuous and to tend to zero as the norm of its input tends to infinity.Source: §2 · p. 3
  • For any spd kernel, capacity is at most the dimension of its reproducing kernel Hilbert space: Cap(Φ) ≤ dim(HΦ).The argument bounds the rank of the attention matrix by the RKHS dimension.Source: §2 · p. 3
  • Infinite RKHS dimension does not guarantee infinite capacity: some kernels have dim(HΦ) = ∞ but Cap(Φ) < ∞.Source: §2 · p. 3

3 CAPACITY OF QUASI LINEAR ATTENTION

The section relates attention capacity to expressivity and develops sorting-based quasi-linear kernels with infinite capacity. It contrasts their properties with finite-feature-map limits and extends univariate kernels to additive multivariate forms.

3 CAPACITY OF QUASI LINEAR ATTENTION

  • The section targets exact attention computation in O((N+M) log^ν(N+M)) operations, avoiding the generally O(NM) cost of infinite-capacity kernels.The existing multivariate Laplace sorting algorithm has complexity exponential in dimension, limiting its practical use.Source: §3 · pp. 3–4

3.1 FINITE FEATURE MAPS

  • Finite feature maps have capacity at most their feature dimension L, while evaluation cost scales linearly with L.Thus, increasing L trades greater capacity for slower evaluation.Source: §3.1 · p. 4

3.2 QUASI LINEAR ATTENTION WITH INFINITE CAPACITY IN ONE DIMENSION

  • Sorting-based algorithms compute sums for univariate piecewise-linear kernels in O((N+M)(log2 N+C)) work; bump and Laplace are quasi linear and have infinite capacity.The sorting approach also applies to polynomial exponential kernels, while the Riesz kernel illustrates that an infinite-dimensional RKHS alone does not ensure infinite capacity.Source: §3.2 · pp. 4–5

3.3 ADDITIVE KERNELS FOR QUASI LINEAR ATTENTION WITH INFINITE CAPACITY

  • Additive kernels inherit quasi-linear computation and at least the capacity of their univariate kernel; additive Laplace and bump have infinite capacity.Theorem 9 also states that positive definiteness and stationarity carry over from the univariate kernel.Source: §3.3 · pp. 5–6

4 IMPLEMENTATION AND EXPERIMENTS

The experiments test kernel capacity, runtime, and text-embedding performance, with the final comparison examining retrieval quality against encoding cost. Across these tests, additive kernels combine high capacity with quasi-linear computation, while practical speed and retrieval outcomes depend on the kernel and setting.

4 IMPLEMENTATION AND EXPERIMENTS

  • The experiments cover associative-recall capacity, runtime break-even points, and text-embedding performance; time measurements use a single NVIDIA GeForce RTX 5090.The paper also describes a vision-transformer experiment in Appendix E.Source: §4 · p. 6

4.1 ASSOCIATIVE RECALL CAPACITY

  • Finite-feature-map kernels achieve roughly zero recall loss below feature dimension L but positive loss above L, whereas Gauss, Laplace, additive Laplace, and additive bump pass the reported capacity tests.Additive Laplace and additive bump pass up to N = 600; additive Riesz optimization fails around N = 40 despite a theoretical capacity lower bound of 128.Source: §4.1 · p. 7

4.2 IMPLEMENTATIONS AND SPEED BENCHMARKS

  • At N = 131072, GLOBAL is roughly 31× faster; fp16 backends are 2–6× faster than fp32 counterparts, placing break-even against FLASH16 slightly beyond N = 16384.At equal precision, FUSED becomes faster than MEM EFF32 at N ≈ 2048; the backend allocates roughly 3× MEM EFF32’s forward peak memory.Source: §4.2 · p. 8

4.3 TEXT EMBEDDING TRANSFORMER

  • Students are initialized from a pretrained softmax teacher and trained by layer-by-layer then end-to-end distillation, before evaluation on short MTEB sequences and longer LoCo retrieval tasks.The distillation data are truncated to 512 tokens; LoCo evaluation spans lengths 2048, 4096, and 8192.Source: §4.3 · pp. 8–9
  • Additive bump and add laplace remain close to the softmax teacher; on LoCo, tri and add laplace are faster than softmax with moderate performance loss, with tri slightly faster and add laplace slightly better.The comparison plots mean nDCG@10 over five LoCo tasks against GPU time per document at 8192 and 16384 tokens.Source: §4.3 · p. 9, Figure 2 · p. 9

5 CONCLUSIONS, LIMITATIONS AND FUTURE WORK

The conclusion presents additive kernels with one-dimensional sorting as a route to infinite capacity and quasi-linear complexity, and reports a speed advantage over softmax for long fp32 sequences. It also cautions that capacity captures only one aspect of expressiveness.

5 CONCLUSIONS, LIMITATIONS AND FUTURE WORK

  • The authors report that their CUDA kernel attention is faster than softmax attention for sequence lengths above 2048 in fp32.The conclusion connects this result to additive kernels and one-dimensional sorting, which provide infinite capacity with quasi-linear complexity.Source: §5 · p. 9
  • Capacity is only one aspect of expressiveness; stationarity and locality also influence transformer expressive power.The authors identify fp16 implementation, causal attention, and larger-scale benchmarks as future directions.Source: §5 · p. 10

AI USE STATEMENT

  • The authors report using generative AI to extend and optimize code and polish wording, while retaining responsibility for the paper and verifying modified code and suggested references.They state that the research ideas, methodology, experimental design, and result interpretation are their own.Source: AI USE STATEMENT · p. 10

REPRODUCIBILITY STATEMENT

A BACKGROUND ON FFMS

The background describes several finite feature maps (FFMs) and their induced kernels, including maps designed to approximate the Gauss or softmax kernel. It also relates these FFMs to additive kernels and contrasts their finite-dimensional univariate components with infinite-dimensional kernels.

A BACKGROUND ON FFMS

  • Peng et al. use Bochner’s theorem to construct a Gauss-kernel feature map; the associated random-feature kernel can be nonnegative neither and may cause numerical instability.The passage states that the kernel induced by the random feature map is not nonnegative, which can lead to numerical instability.Source: A BACKGROUND ON FFMS · p. 14
  • For numerical comparisons, the authors use Choromanski et al.’s ReLU feature map because it generally performed best among their FFMs.The softmax-oriented exponential feature map is nonnegative but not stationary.Source: A BACKGROUND ON FFMS · p. 14
  • FFMs in Table 4 use coordinate-wise feature functions, and their induced kernels are additive sums of univariate kernels.The feature vector stacks the functions applied to each coordinate; the resulting kernel sums the univariate inner products.Source: A BACKGROUND ON FFMS · p. 14
  • The univariate kernel underlying these FFMs has dimension J, typically 1 or 2, whereas Laplace, Riesz, and bump kernels have infinite dimension.Source: A BACKGROUND ON FFMS · p. 15
  • DPFP uses ν < D to keep feature components linearly independent; the experiments use ν = 3 and D = 64, giving dim H = 384.The passage says ν ≥ D is wasteful because it no longer increases the RKHS dimension.Source: A BACKGROUND ON FFMS · p. 15

B CONNECTION TO SLICING

The section connects additive kernels to slicing and develops sorting-based computation for spline-kernel attention. It also describes a fused implementation that avoids large intermediates while trading off some reduction efficiency.

B CONNECTION TO SLICING

  • Fixing the slicing directions to coordinate axes yields the additive kernel, while random directions give a Monte Carlo approximation with cost O(L(N+M) log N) and error rate O(L^-1/2).Choosing the D coordinate axes can also be viewed as quasi-Monte Carlo design, which sometimes achieves better error rates than standard Monte Carlo.Source: B CONNECTION TO SLICING · p. 15
  • Additive Riesz approximates multivariate Riesz by Monte Carlo, but additive Laplace instead approximates a power series, not ℓ1 or ℓ2 Laplace.Source: B CONNECTION TO SLICING · p. 15

C FAST COMPUTATION OF SPLINE-KERNEL ATTENTION

  • For N, M ≫ D, C, the naive absolute-value sum costs O((D+C)NM), while sorting reduces it to O(D(N+M)(log2 N+C)).Finitely-noded spline kernels can be computed by applying the one-dimensional sorting algorithm at each node.Source: C FAST COMPUTATION OF SPLINE-KERNEL… · pp. 15–16

C.1 PARALLELIZATION AND MEMORY OVERHEAD

  • A practical implementation must balance parallelism and memory efficiency; serializing over D reduces memory use but sacrifices parallelism.Source: §C.1 · p. 16
  • The implementation sorts queries and keys, scans prefix sums on the fly, and parallelizes over P, D, and C without materializing large intermediates.The final reduction over D requires atomic stores, identified as the implementation’s main bottleneck.Source: §C.1 · p. 16

C.2 GRADIENTS AND AUTODIFF

The section extends sorting-based attention to backward computation and variable-length batches. Its forward-and-backward benchmark reports GLOBAL beating FLASH16 before sequence length 16,384.

C.2 GRADIENTS AND AUTODIFF

  • All gradients run in O((N+M) log(N+M)); the backward pass reuses sorted tensors and recomputes the prefix sum because it is formed on the fly.The v-gradient follows the forward computation with q and k swapped; the work for the sums is O((N+M)(log2 N+C)).Source: §C.2 · p. 16
  • GLOBAL beats FLASH16 before N = 16384 in forward-plus-backward timing, with relative memory overhead around 1.5 versus the fp32 softmax backend.The table reports mean runtime over 10 runs after 5 warm-up iterations, with forward and backward timed together.Source: Table 5 · p. 17, §C.2 · p. 17

C.3 PADDING

  • Zero-padding values preserves kernel sums, but normalization needs an O(ND) correction for padded keys; this permits variable-length samples to share a batch.Outputs for padded queries are discarded.Source: §C.3 · p. 17

C.4 CAUSAL MASKING

Causal masking is handled by recursively partitioning the masked kernel matrix into smaller causal and unmasked sums. The same section describes a numerically stable weighted Laplace sum and notes that autoregressive decoding remains a full-cache operation.

C.4 CAUSAL MASKING

  • Algorithm 2 decomposes the triangular mask into 2N − 1 blocks and computes the causal sum in O(N log^2(N)) work; blocks can run independently with atomic additions at accumulation.The recursion splits the matrix into four blocks: diagonal blocks remain causal, the upper-right block is zero, and the lower-left block is an ordinary kernel sum.Source: §C.4 · p. 18, Figure 3 · p. 18

C.5 AUTOREGRESSIVE DECODING

  • During autoregressive decoding, each new query reads the full KV cache at cost O(N(D+C)); the passage says this procedure works for any kernel.The stated setting is memory-bound rather than compute-bound.Source: §C.5 · p. 19

D WEIGHTED LAPLACE SUMS

  • The weighted Laplace algorithm expresses prefix sums relative to the current node, keeping exponentials in (0, 1] and avoiding overflow from separate factors.Its work is O((N+M)(log2 N+C)), and intermediate quantities remain bounded independently of input magnitude; gradients and masking are out of scope for this implementation.Source: D WEIGHTED LAPLACE SUMS · p. 19

E VISION TRANSFORMER

The CIFAR-10 experiment is a controlled sanity check because its sequence length does not exceed the head dimension, so the capacity mechanism cannot bind. After full distillation, all kernels reach within 1.0% accuracy of softmax.

E VISION TRANSFORMER

  • The 12-layer CIFAR-10 model uses three heads of dimension 64 and patch size 4, yielding sequence length 64.Because sequence length does not exceed head dimension, a 64-dimensional RKHS suffices to separate all tokens, so capacity cannot bind.Source: E VISION TRANSFORMER · p. 19
  • Each model is trained from scratch with random initialization and also distilled from a softmax teacher with 96.57% accuracy.Distillation copies teacher weights, finetunes one attention module at a time to match its attention output, then finetunes all parameters against teacher logits.Source: E VISION TRANSFORMER · pp. 19–20
  • Scratch training uses AdamW for 1000 epochs, while vision end-to-end distillation runs 50k steps with KL temperature T = 4.The layer-by-layer stage runs 10k steps per layer.Source: E VISION TRANSFORMER · p. 20
  • After full distillation, all kernels train to within 1.0% accuracy of softmax.The reported result is from the CIFAR-10 vision-transformer experiment.Source: Table 6 · p. 20, E VISION TRANSFORMER · p. 20
  • The copied Gauss student reaches 96.10%, while Laplace with τ = 6 reaches 94.92% without finetuning.All other kernels are at least 19% below the copied Gauss student, according to the passage.Source: E VISION TRANSFORMER · p. 20
  • Distillation matches or improves training from scratch for every kernel in both accuracy and training time.Layer-by-layer finetuning does most of the work; end-to-end finetuning adds less than 1%.Source: E VISION TRANSFORMER · p. 20

F ASSOCIATIVE RECALL

The associative-recall setup assigns each token a learned embedding and a one-hot value, then forms keys from both. The paper's version uses standard-transformer keys and differs from Schlag et al. (2021) only in the key matrix.

F ASSOCIATIVE RECALL

  • Associative recall tests whether attention can retrieve the value assigned to a queried token from a sequence of embedded tokens.The setup samples two random permutations to arrange token identities and their one-hot values.Source: F ASSOCIATIVE RECALL · p. 20
  • Each key combines the token embedding and its one-hot value, while the query is formed from the token embedding.The setup uses learnable key and query projection matrices.Source: F ASSOCIATIVE RECALL · p. 20
  • The paper uses standard-transformer keys kn = WKun with WK ∈ RD×D, unlike Schlag et al. (2021), and is contained in their setting by zero-padding WK.The paper also states that its loss is, up to a factor of 1/2, the inner term of (13).Source: F ASSOCIATIVE RECALL · p. 20

G PROOFS

The proofs establish infinite capacity for several kernel constructions and show how additive kernels retain that capacity while remaining quasi-linear for fixed dimension. They also analyze finite-feature-map limits and fast summation procedures.

G PROOFS

  • For translation-invariant kernels Φ(q, k) = F(q − k) with F in C0, scaling distinct inputs makes attention approach the identity for any sequence length.The proof concludes that Cap(Φ) = ∞ and notes that the scaling is realizable through query and key weight matrices.Source: G PROOFS · pp. 20–21
  • The Gauss kernel has infinite capacity by correspondence with the softmax kernel when keys have equal norms.The proof selects infinitely many unit vectors when D ≥ 2.Source: G PROOFS · p. 21
  • A finite feature map forces the RKHS to be finite dimensional, which in the proof implies a finite-support spectral measure.The argument represents the stationary kernel through Bochner’s theorem and an even/odd function-space feature map.Source: G PROOFS · p. 21
  • Piecewise-linear kernels with finitely many nodes admit a sorting-based summation algorithm with complexity O((N+M)(log(N+M) + CJ)).The proof reuses sorted keys and queries rather than repeating sorting and bisection for each node; it also states an extension to noncontinuous piecewise-linear functions.Source: G PROOFS · pp. 21–22
  • For the one-dimensional Riesz kernel, the proof establishes capacity exactly 2.For N = 2, opposite scalar query-key pairs make the attention matrix approach the identity as τ grows.Source: G PROOFS · p. 22
  • Any univariate capacity attainable by ϕ transfers to its additive kernel Φ; for the Riesz kernel, the proof gives capacity at least 2D.The transfer embeds scalar keys and queries into each coordinate of the D-dimensional inputs.Source: G PROOFS · p. 22
  • For fixed D, additive kernels built from a quasi-linear univariate kernel can compute all query outputs in quasi-linear time.The proof computes D inner sums and then combines their results at cost O(MDC).Source: G PROOFS · p. 22
  • In the associative-recall proof, normalization over all keys yields A_{m,l}(q,kσ)=A_{m,σ(l)}(q,k), establishing the reindexing relation for activations under permuted keys.Source: G PROOFS · p. 23

Related papers

Loading 2609.35349v1…