Quasi Linear Kernel Attention with Infinite Capacity

Nicolaj Rux; Johannes Hertrich; Sebastian Neumayer

cs.LGSep 28, 2026

Quasi Linear Kernel Attention with Infinite Capacity

Authors: Nicolaj Rux, Johannes Hertrich, Sebastian Neumayer

Organizations: Faculty of Mathematics, Chemnitz University of Technology, 09126 Chemnitz, Germany · Institute of Computer Science, University of Göttingen, 37073 Göttingen, Germany

Abstract

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.

Figures & tables

#	Step	Work	Memory
1	Input $\bm{s}\in\mathbb{R}^{M}$ , $\bm{t}\in\mathbb{R}^{N}$ and $\bm{v}\in\mathbb{R}^{N\times C}$	-	$(C{+}1)N{+}M$
2	Output $z_{m}\coloneqq\sum_{n=1}^{N}\|s_{m}-t_{n}\|v_{n}\in\mathbb{R}^{C}\quad m=1,\ldots,M$	-	$MC$
3	$\sigma\coloneqq\texttt{argsort}(\bm{t})$ s.t. $t_{\sigma(1)}\leq\ldots\leq t_{\sigma(N)}$	$N\log_{2}N$	$N$
4	$p_{m}\coloneqq\max\bigl(\{n\colon t_{\sigma(n)}\leq s_{m}\}\cup\{0\}\bigr)$ for $m=1,\ldots,M$	$M\log_{2}N$	$M$
5	$\textrm{pref\_t}_{n}\coloneqq\sum_{j=1}^{n}t_{\sigma(j)}v_{\sigma(j)}\in\mathbb{R}^{C}$ for $n=0,\ldots,N$	$(N{+}1)C$	$(N{+}1)C$
6	$\textrm{pref\_v}_{n}\coloneqq\sum_{j=1}^{n}v_{\sigma(j)}\in\mathbb{R}^{C}$ for $n=0,\ldots,N$	$(N{+}1)C$	$(N{+}1)C$

Algorithm 1 Weighted absolute value sum

Name	Formula	dim	stationary	additive	capacity	complexity
softmax	$\mathrm{e}^{q^{\top}k}$	$\infty$	no	no	$\infty$	quadratic
gauss	$\mathrm{e}^{-\frac{1}{2}\\|q-k\\|_{2}^{2}}$	$\infty$	yes	no	$\infty$	quadratic
laplace	$\mathrm{e}^{-\\|q-k\\|_{1}}$	$\infty$	yes	no	$\infty$	quasi linear
riesz	$\\|q\\|_{2}+\\|k\\|_{2}-\\|q-k\\|_{2}+\varepsilon$	$\infty$	partially	no	$-$	quadratic
add_riesz	$\|s\|+\|t\|-\|s-t\|+\varepsilon$	$\infty$	partially	yes	$\geq 2D$	quasi linear
add_laplace	$\mathrm{e}^{-\|s-t\|}$	$\infty$	yes	yes	$\infty$	quasi linear

Table 1: For additive kernels, we state the univariate

\phi(s,t)

. dpfp is described in Appendix A and the remaining kernels are defined for

q,k\in\mathbb{R}^{D}

. All kernels except tri are nonnegative. riesz and add_riesz use

\varepsilon=10^{-3}

and are partially stationary, i.e., stationary up to terms depending only on

q

or only on

k

. The last column is the best known kernel sum complexity in

N{+}M

$N$	GLOBAL	GLOBAL	FUSED	KEOPS	MEM_EFF32	MEM_EFF16	FLASH16	CUDNN16
	add_laplace fp32	add_riesz fp32			softmax
$128$	$0.442$	$0.455$	$0.305$	$0.381$	$\underline{0.029}$	$\mathbf{0.021}$	$\mathbf{0.021}$	$0.023$
$256$	$0.575$	$0.539$	$0.353$	$0.425$	$\underline{0.058}$	$0.031$	$\mathbf{0.022}$	$0.030$
$512$	$0.824$	$0.694$	$0.478$	$0.688$	$\underline{0.144}$	$0.058$	$\mathbf{0.041}$	$0.043$
$1024$	$1.305$	$1.014$	$0.719$	$1.648$	$\underline{0.431}$	$0.152$	$\mathbf{0.093}$	$0.119$
$2048$	$2.260$	$1.690$	$\underline{1.278}$	$4.754$	$1.645$	$0.537$	$0.266$	$\mathbf{0.265}$

Table 2: Forward-pass of kernel attention ( 1 ) with runtime in milliseconds, mean over 10 runs after 5 warm-up iterations. Relative standard deviations are below

18\%

for

N\leq 512

, below

4\%

N=1024

and below

3\%

for

N\geq 2048

. Bold marks the fastest method overall; underline marks the fastest among the five fp32 methods. Shape:

B=4

H=12

D=64

C=64

M=N

		MTEB $\uparrow$	LoCo (nDCG@10) $\uparrow$
Kernel	scale $\tau$	$\leq 512$	$2048$	$4096$	$8192$
softmax (teacher)	$2.828$	$64.41$	$0.873$	$0.882$	$0.883$
gauss	$2.828$	$64.41$	$0.875$	$0.884$	$0.886$
laplace	$6.0$	$64.38$	$0.885$	$0.891$	$0.887$
riesz	$1.0$	$61.69$	$0.805$	$0.839$	$0.858$
add_riesz	$1.0$	$61.72$	$0.769$	$0.813$	$0.832$

Table 3: Results after full two-stage distillation for MTEB and LoCo. Best performance among quasi linear kernels is underlined.

		MTEB $\uparrow$	LoCo (nDCG@10) $\uparrow$
Kernel	scale $\tau$	$\leq 512$	$2048$	$4096$	$8192$
softmax (teacher)	$2.828$	$64.41$	$0.873$	$0.882$	$0.883$
gauss	$2.828$	$64.41$	$0.875$	$0.884$	$0.886$
laplace	$6.0$	$64.38$	$0.885$	$0.891$	$0.887$
riesz	$1.0$	$61.69$	$0.805$	$0.839$	$0.858$
add_riesz	$1.0$	$61.72$	$0.769$	$0.813$	$0.832$

Table 3: Results after full two-stage distillation for MTEB and LoCo. Best performance among quasi linear kernels is underlined.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Paper	$\varphi$	$J$	$f_{j}$
Katharopoulos et al. (2020)	$\varphi_{\mathrm{elu}}$	$1$	$f_{1}(x)=\mathrm{elu}(x)+1$
Peng et al. (2021)	$\varphi_{\mathrm{tri}}$	$2$	$f_{1}(x)=\sin(x),f_{2}(x)=\cos(x)$
Choromanski et al. (2021)	$\varphi_{\mathrm{ReLU}}$	$1$	$f_{1}(x)=\mathrm{ReLU}(x)$

Appendix

Table 4: Examples of FFMs of the form ( 19 ).

	add_riesz fp32			softmax
$N$	GLOBAL	FUSED	KEOPS	MEM_EFF32	MEM_EFF16	FLASH16	CUDNN16
$128$	$1.066$	$0.913$	$1.598$	$\underline{0.098}$	$0.071$	$0.059$	$\mathbf{0.057}$
$256$	$1.329$	$1.147$	$2.424$	$\underline{0.233}$	$0.104$	$\mathbf{0.067}$	$0.069$
$512$	$1.832$	$1.620$	$4.087$	$\underline{0.609}$	$0.231$	$\mathbf{0.132}$	$0.143$
$1024$	$2.950$	$2.664$	$11.96$	$\underline{1.923}$	$0.634$	$\mathbf{0.324}$	$0.363$
$2048$	$5.283$	$\underline{4.927}$	$40.33$	$7.360$	$2.262$	$\mathbf{0.978}$	$1.050$

Appendix

Table 5: Forward+backward of kernel attention ( 1 ) with runtime in milliseconds, mean over 10 runs after 5 warm-up iterations, with forward and backward timed together in a single region. Relative standard deviations are below

10\%

for

N\leq 512

, below

1.2\%

N=1024

and below

0.7\%

for

N\geq 2048

. Bold marks the fastest method overall; underline marks the fastest among the four fp32 methods. Shape:

B=4

H=12

D=64

C=64

M=N

1	Input $\bm{s}\in\mathbb{R}^{M}$ , $\bm{t}\in\mathbb{R}^{N}$ and $\bm{v}\in\mathbb{R}^{N\times C}$
2	Output $(z_{m})_{m=1}^{M}=\mathrm{CausalSorting}(\bm{s},\bm{t},\bm{v})$ given by $z_{m}\coloneqq\sum_{n=1}^{m}\phi(s_{m},t_{n})v_{n}\in\mathbb{R}^{C}$
3	Denote by $\mathrm{Sorting}(\bm{s},\bm{t},\bm{v})$ the application of Algorithm 1
4	If $\min(M,N)=1$ :
5	compute $(z_{m})_{m=1}^{M}$ naively
6	Else:
7	compute $(z_{m})_{m=1}^{M/2}=\mathrm{CausalSorting}((s_{m})_{m=1}^{M/2},(t_{n})_{n=1}^{N/2},(v_{n})_{n=1}^{N/2})$

Appendix

Algorithm 2 Sorting algorithm for causal masked kernel sums

#	Step	Work	Memory
1	Input $\bm{s}\in\mathbb{R}^{M}$ , $\bm{t}\in\mathbb{R}^{N}$ and $\bm{v}\in\mathbb{R}^{N\times C}$	-	$(C{+}1)N{+}M$
2	Output $z_{m}\coloneqq\sum_{n=1}^{N}\exp(-\|s_{m}-t_{n}\|)v_{n}\in\mathbb{R}^{C}\quad m=1,\ldots,M$	-	$MC$
3	$\sigma\coloneqq\texttt{argsort}(\bm{t})$ s.t. $t_{\sigma(1)}\leq\ldots\leq t_{\sigma(N)}$	$N\log_{2}N$	$N$
4	$p_{m}\coloneqq\max\bigl(\{n\colon t_{\sigma(n)}\leq s_{m}\}\cup\{0\}\bigr)$ for $m=1,\ldots,M$	$M\log_{2}N$	$M$
5	$\delta_{n}\coloneqq t_{\sigma(n)}-t_{\sigma(n-1)}\geq 0$ for $n=2,\ldots,N$	$N$	$N$
6	$L_{n}\coloneqq e^{-\delta_{n}}L_{n-1}+v_{\sigma(n)}\in\mathbb{R}^{C}$ for $n=1,\ldots,N$ , $L_{0}\coloneqq\bm{0}$	$2NC$	$(N{+}1)C$

Appendix

Algorithm 3 Weighted Laplace sum

			distillation from softmax
Kernel	scale $\tau$	scratch $\uparrow$	copy $\uparrow$	+layerwise $\uparrow$	+end-to-end $\uparrow$
softmax (teacher)	$2.828$	$96.57$	—	—	—
gauss	$2.828$	$96.70$	$96.10$	$96.56$	$96.55$
laplace	$6.0$	$96.71$	$94.92$	$96.57$	$96.64$
riesz	$1.0$	$94.14$	$58.12$	$95.61$	$95.83$
add_riesz	$1.0$	$93.68$	$55.93$	$95.84$	$96.21$

Appendix

Table 6: CIFAR-10 test accuracy (%). Training from scratch as well as piecewise distillation.

Explore similar work

Flexformer: Flexible Linear Transformer with Learnable Attention Kernel
Jun 26, 2026Haoran Zhang, Feng ZhouLanguage Modeling
Center for Applied Statistics and School of Statistics, Renmin University of China Beijing, China
Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones
Jul 19, 2026Ayoub Ghriss, Sourav ChakrabortyKimi Delta Attention Linear Attention
Department of Computer Science, University of Colorado Boulder
Exact Linear Attention
May 13, 2026Weinuo OuKimi Delta Attention Transformer Attention
Wuyi University

Jun 26, 2026, cs.LGJ/K move · Enter open · S save

Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

Haoran Zhang, Feng Zhou

Center for Applied Statistics and School of Statistics, Renmin University of China Beijing, China

Language Modeling

Open paper AskPDF arXiv

Transformer models rely on attention mechanism to capture long-range dependencies but suffer from quadratic complexity, limiting their scalability to long sequences. Kernel-based linear attention reduces this complexity but typically relies on fixed or weakly learnable kernels, restricting expressiveness and performance. In this work, we propose Flexformer, a flexible linear Transformer that learns attention kernels in a fully data-driven manner. Flexformer builds on random Fourier feature-based linear attention and treats spectral frequencies as trainable parameters, enabling the model to learn a broad family of attention kernels. We develop both stationary and nonstationary variants, with the latter offering strictly greater expressiveness. Extensive experiments on language modeling and sequence classification demonstrate that Flexformer consistently outperforms baselines. Moreover, Flexformer can be effectively distilled from pretrained Transformers to recover softmax attention and exhibits strong kernel transferability across domains, achieving both high efficiency and competitive performance on long-sequence tasks.

Abstract

Explore similar work

Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones

Exact Linear Attention