Sparse Attention as Compact Kernel Regression

Saul Santos; Nuno Gonçalves; Daniel C. McNamee; Marcos Treviso; André F. T Martins

cs.LGJan 30, 2026

Sparse Attention as Compact Kernel Regression

Authors: Saul Santos, Nuno Gonçalves, Daniel C. McNamee, Marcos Treviso, André F. T Martins

Organizations: Instituto Superior Técnico & Instituto de Telecomunicações, Universidade de Lisboa, Portugal · Carnegie Mellon University, Pittsburgh, USA · Champalimaud Research, Lisbon, Portugal · TransPerfect, Lisbon, Portugal

Abstract

Recent work has revealed a link between self-attention mechanisms in transformers and test-time kernel regression via the Nadaraya-Watson estimator, with standard softmax attention corresponding to a Gaussian kernel. However, a kernel-theoretic understanding of sparse attention mechanisms is currently missing. In this paper, we establish a formal correspondence between sparse attention and compact (bounded support) kernels. We show that normalized ReLU and sparsemax attention arise from Epanechnikov kernel regression under fixed and adaptive normalizations, respectively. More generally, we demonstrate that widely used kernels in nonparametric density estimation -- including Epanechnikov, biweight, and triweight -- correspond to $α$ -entmax attention with $α= 1 + \frac{1}{n}$ for $n \in \mathbb{N}$ , while the softmax/Gaussian relationship emerges in the limit $n \to \infty$ . This unified perspective explains how sparsity naturally emerges from kernel design and provides principled alternatives to heuristic top- $k$ attention and other associative memory mechanisms. Experiments with a kernel-regression-based variant of transformers -- Memory Mosaics -- show that kernel-based sparse attention achieves competitive performance on language modeling, in-context learning, and length generalization tasks, offering a principled framework for designing attention mechanisms.

Figures & tables

Panel of Figure 1 — Figure 1: Comparison between classical kernel functions and their corresponding attention activations. The left plot shows normalized one-dimensional kernel functions commonly used in kernel regression evaluated on an input $\bm{u}\in\mathbb{R}^{n}$ . The right panel shows different attention transformations applied to a 3D score vector $\bm{u}=[0,u_{1},u_{2}]$ . Each transformation allows an interpretation with corresponding kernels shown on the left: softmax ( $\alpha=1$ ) corresponds to the Gaussian kernel; normalized ReLU, ReLUmax, and sparsemax ( $\alpha=2)$ are derived from Epanechnikov kernel variants ( $r=1$ ); The 4/3- and 1.5-entmax transformations correspond to the biweight ( $r=2$ ) and triweight ( $r=3$ ) kernels, respectively, while 1- and 2-ReLUmax follow analogous patterns. top- $k$ uniform mirrors the top- $k$ uniform kernel, while top- $k$ softmax can be viewed as a truncated Gaussian kernel.

Figure 2: Validation loss of Memory Mosaic models with different kernels. Comparison on the BabiStories dataset for varying model depths. The horizontal axis shows the number of training iterations.

Method	1k		2.5k		5k		10k		20k		40k
Method	Acc $\uparrow$	TVD $\downarrow$	Acc $\uparrow$	TVD $\downarrow$	Acc $\uparrow$	TVD $\downarrow$	Acc $\uparrow$	TVD $\downarrow$	Acc $\uparrow$	TVD $\downarrow$	Acc $\uparrow$	TVD $\downarrow$
RetNet	51.3	67.4	50.5	63.9	81.5	33.3	88.4	25.7	90.8	22.5	92.6	18.9
H3	51.5	65.1	58.0	57.0	70.9	47.4	76.9	41.4	80.7	36.3	87.7	25.4
Gla	51.8	65.3	51.1	62.2	82.1	31.0	90.2	24.5	91.2	21.2	92.9	19.0
S4	52.1	62.8	51.7	62.2	54.5	59.0	63.3	52.9	72.1	45.8	82.5	33.8
Mamba	69.4	51.0	79.4	41.0	88.7	30.6	90.5	28.4	91.3	27.0	91.5	27.0

Table 1: Performance comparison across varying training set sizes in the Regbench in-context learning benchmark. Accuracy (

\uparrow

) measures the correct prediction rate, and TVD (

\downarrow

) measures divergence from the target distribution. We bold the best performing models.

	MQMTAR $(L=4)$							Reverse $(L=4)$		Sort $(L=2)$
Kernel	$1\times$	$2\times$	$4\times$	$8\times$	$16\times$	$32\times$	$64\times$	$1\times$	$1.5\times$	$1\times$	$2\times$	$4\times$
Transformer	1.00 ${}^{(\underline{1.00})}$	1.00 ${}^{(\underline{1.00})}$	0.99 ${}^{(0.99})$	0.09 ${}^{(0.15})$	0.01 (0.01)	0.00 (0.00)	0.00 (0.00)	1.00 ${}^{(\underline{1.00})}$	0.52 (0.62)	1.00 ${}^{(\underline{1.00})}$	0.00 (0.00)	0.00 (0.00)
Gaussian	0.99 ${}^{(\underline{1.00})}$	0.96 ${}^{(\underline{1.00})}$	0.93 (0.99)	0.85 (0.97)	0.64 (0.81)	0.14 (0.20)	0.00 (0.00)	1.00 ${}^{(\underline{1.00})}$	0.14 (0.38)	1.00 ${}^{(\underline{1.00})}$	0.06 (0.17)	0.00 (0.00)
Gaussian kNN, $k=32$	0.99 ${}^{(\underline{1.00})}$	0.97 ${}^{(\underline{1.00})}$	0.93 (0.99)	0.84 (0.96)	0.68 (0.85)	0.27 (0.36)	0.05 (0.05)	0.99 ${}^{(\underline{1.00})}$	0.00 (0.00)	1.00 ${}^{(\underline{1.00})}$	0.80 (0.91)	0.01 (0.02)
Uniform kNN, $k=32$	0.04 (0.08)	0.00 (0.00)	0.00 (0.00)	0.00 (0.00)	0.00 (0.00)	0.00 (0.00)	0.00 (0.00)	0.42 (0.66)	0.00 (0.00)	1.00 ${}^{(\underline{1.00})}$	0.00 (0.00)	0.00 (0.00)
Epanechnikov (normReLU)	1.00 ${}^{(\underline{1.00})}$	1.00 ${}^{(\underline{1.00})}$	0.89 ${}^{(\underline{1.00})}$	0.66 ${}^{(\underline{1.00})}$	0.44 ${}^{(\underline{1.00})}$	0.33 ${}^{(\underline{0.98})}$	0.31 ${}^{(\underline{0.94})}$	1.00 ${}^{(\underline{1.00})}$	0.19 (0.56)	1.00 ${}^{(\underline{1.00})}$	0.05 (0.14)	0.00 (0.00)

Table 2: Exact match accuracy on representative synthetic tasks. In-distribution results (

n=64

) and out-of-distribution performance at increasing sequence lengths are reported. Values show the mean across three seeds, with the maximum across seeds indicated in superscript. Best average performance is in bold and overall maximum performance is underlined .

L

is the number of layers.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Hyperparameter	1 layer	8 layers	12 layers
Optimization
Batch size (total)	512	512	512
Learning rate	$5{\times}10^{-3}$	$5{\times}10^{-3}$ ∗	$1{\times}10^{-3}$
Min learning rate	$1{\times}10^{-4}$	$1{\times}10^{-}4$ ∗	$5{\times}10^{-5}$
Max iterations	80k	80k	80k
Weight decay	0.1	0.1	0.1

Appendix

Table 3: Main training and model hyperparameters across model depths. Unless otherwise specified, hyperparameters are shared across all experiments. ∗ means that for 2-ReLUmax we used a learning rate of

1\times 10^{-3}

and minimum learning rate of

5\times 10^{-5}

due to training instability.

Dataset partition	# Stories	# Tokens (GPT-2 tokenizer)	Avg. # Characters per story
Train	2.2M	474,704,907	888
Valid	2.2k	4,749,107	889

Appendix

Table 4: BabiStories dataset statistics [ 49 ] .

Figure 3: Training and validation loss of Memory Mosaic models with different kernels. Comparison on the BabiStories dataset for varying model depths. The horizontal axis shows the number of training iterations.

Parameter	Range
weight decay	{ $10^{-1}$ , $10^{-2}$ }
depth	{2, 4, 8}
embedding size	{32, 128, 256}
number of heads	{1, 2, 4, 8}
learning rate	{ $5\times 10^{-4}$ }
maximum number of epochs	{200}

Appendix

Table 5: Hyperparameter space for the in-context learning experiment. Depth denotes the number of stacked blocks comprising persistent and contextual memories. Embedding size corresponds to the head dimension in transformers. We report results from the epoch achieving the lowest validation loss.

Panel of Figure 4 — Figure 4: Illustration of the MQMTAR, reverse, and sort tasks. In the MQMTAR task, Soft colors indicate different keys, values, and queries for visual clarity. In the reverse task, the model receives a sequence of integers and must output them in reverse order. In the sort task, the model receives a sequence of integers and must output the sequence in ascending order.

Parameter	MQMTAR	Reverse	Sorting
Training examples	$50\times 10^{6}$	$30\times 10^{6}$	$40\times 10^{6}$
# epochs	1	1	1
Vocab size	263	36	36
Weight decay	$10^{-1}$	$10^{-1}$	$10^{-1}$
Depth	4	4	2
Embedding size	512	512	512

Appendix

Table 6: Hyperparameters used for MQMTAR recall, reverse, and sorting tasks. Hyperparameters were fixed to keep models simple and comparable.

MQMTAR

(L=4)

Reverse

(L=4)

Sort

(L=2)

Kernel

1\times

2\times

4\times

8\times

16\times

32\times

64\times

1\times

1.5\times

1\times

2\times

4\times

Gaussian kNN,

k=16

0.99

{}^{(\underline{1.00})}

0.97

{}^{(\underline{1.00})}

0.93

{}^{(\underline{0.99})}

0.85

{}^{(\underline{0.96})}

0.65 (0.78)

0.25

{}^{(\underline{0.36})}

0.05

{}^{(\underline{0.09})}

0.86

{}^{(\underline{1.00})}

0.00

{}^{(\underline{0.00})}

1.00

{}^{(\underline{1.00})}

0.74

{}^{(0.89})

0.00 (0.00)

Gaussian kNN,

k=32

0.99

{}^{(\underline{1.00})}

0.97

{}^{(\underline{1.00})}

0.93

{}^{(\underline{0.99})}

0.84

{}^{(\underline{0.96})}

0.68

{}^{(\underline{0.85})}

0.27

{}^{(\underline{0.36})}

0.05 (0.05)

0.99

{}^{(\underline{1.00})}

0.00

{}^{(\underline{0.00})}

1.00

{}^{(\underline{1.00})}

0.80

{}^{(\underline{0.91})}

0.01

{}^{(\underline{0.02})}

Uniform kNN,

k=16

0.58

{}^{(\underline{0.71})}

0.01

{}^{(\underline{0.01})}

0.00

{}^{(\underline{0.00})}

0.00

{}^{(\underline{0.00})}

0.00

{}^{(\underline{0.00})}

0.00

{}^{(\underline{0.00})}

0.00

{}^{(\underline{0.00})}

0.00

{}^{(0.00})

0.00

{}^{(\underline{0.00})}

1.00

{}^{(\underline{1.00})}

0.41

{}^{(\underline{0.63})}

0.00

{}^{(\underline{0.00})}

Uniform kNN,

k=32

0.04 (0.08)

0.00

{}^{(\underline{0.00})}

0.00

{}^{(.00})

0.00

{}^{(\underline{0.00})}

0.00

{}^{(\underline{0.00})}

0.00

{}^{(\underline{0.00})}

0.00

{}^{(\underline{0.00})}

0.42

{}^{(\underline{0.66})}

0.00

{}^{(\underline{0.00})}

1.00

{}^{(\underline{1.00})}

0.00

{}^{(0.00})

0.00

{}^{(\underline{0.00})}

Appendix

Table 7: Exact match accuracy on representative synthetic tasks. In-distribution results (

n=64

L

represents the number of layers.

	MQMTAR
Model	$1\times$	$2\times$	$4\times$	$8\times$	$16\times$	$32\times$	$64\times$
Transformer (ALiBi)	1.00 ${}^{(\underline{1.00})}$	1.00 ${}^{(\underline{1.00})}$	0.99 (0.99)	0.09 (0.15)	0.01 (0.01)	0.00 (0.00)	0.00 (0.00)
MM, Epanechnikov (ALiBi)	0.99 ${}^{(\underline{1.00})}$	0.96 ${}^{(\underline{1.00})}$	0.92 ${}^{(\underline{1.00})}$	0.85 ${}^{(\underline{1.00})}$	0.72 (0.94)	0.43 (0.71)	0.14 (0.29)
MM, Epanechnikov (no ALiBi)	1.00 ${}^{(\underline{1.00})}$	1.00 ${}^{(\underline{1.00})}$	1.00 ${}^{(\underline{1.00})}$	0.98 ${}^{(\underline{1.00})}$	0.86 ${}^{(\underline{0.98})}$	0.63 ${}^{(\underline{0.84})}$	0.31 ${}^{(\underline{0.51})}$

Appendix

Table 8: Exact match accuracy on the MQMTAR task. In-distribution results (

1\times

) and out-of-distribution performance at increasing sequence lengths are reported. Values show the mean across seeds, with the maximum indicated in superscript. Best average performance is in bold and overall maximum performance is underlined .

Figure 5: Prediction performance on the Simple English Wikipedia dataset using models trained on the BABISTORIES corpus. The plot reports the per-token average loss as a function of the token position within a 512-token input sequence.

Transformer Attention

Transformer attention lets models relate tokens by weighting information from other positions, enabling context-dependent sequence and multimodal processing.

See all 320 papers

Quasi Linear Kernel Attention with Infinite Capacity
Sep 28, 2026Nicolaj Rux, Johannes Hertrich, Sebastian NeumayerTransformer Expressivity Attention Mechanisms
Faculty of Mathematics, Chemnitz University of Technology, 09126 Chemnitz, Germany · Institute of Computer Science, University of Göttingen, 37073 Göttingen, Germany
Projection-Free Transformers via Gaussian Kernel Attention
May 4, 2026Debarshi Kundu, Archisman Ghosh, Swaroop Ghosh +1Transformer Self-Attention
The Pennsylvania State University, University Park, USA
Sparse Attention Is Matrix Approximation, Not Choosing from a Bag of Values
Oct 7, 2026Fang Wan, Xufeng Liu, Fan Li +1Sparse Attention Efficient Attention
Stony Brook University · Duke University

Sep 28, 2026J/K move · Enter open · S save

Quasi Linear Kernel Attention with Infinite Capacity

Nicolaj Rux, Johannes Hertrich, Sebastian Neumayer

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

Transformer Expressivity Attention Mechanisms

Open paper AskPDF arXiv

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.

Abstract

Transformer Attention

Explore similar work

Quasi Linear Kernel Attention with Infinite Capacity

Projection-Free Transformers via Gaussian Kernel Attention

Sparse Attention Is Matrix Approximation, Not Choosing from a Bag of Values