Sparse Attention as Compact Kernel Regression

TL;DR

This work reframes self-attention in transformers as a kernel-regression problem, showing that standard softmax attention corresponds to a Gaussian kernel while sparse attention can be realized with compact-support kernels. It develops a unified theory in which Epanechnikov, biweight, and triweight kernels map to α-entmax variants (α = 1 + 1/r), with top- extit{k} methods interpreted as data-adaptive truncations, and introduces ReLUmax as a max-anchored Epanechnikov variant. The authors instantiate Memory Mosaics, a kernel-regression–based transformer with contextual and persistent memories, and demonstrate that adaptive, compact kernels yield competitive or superior performance on language modeling, in-context learning, and length generalization. This framework provides a principled approach to designing attention mechanisms with controllable sparsity and locality, bridging theory, architecture, and empirical evaluation.

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 (5)

• 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$); 4/3- and 1.5-entmax mirror biweight ($r=2$) and triweight ($r=3$) kernels; 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.
• 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.
• 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.
• 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.

Theorems & Definitions (2)

• proposition 1
• proof