Learnable Low-Rank Polynomial Sketch for Effective Linear Attention

Zongyue Qin; Haoran Deng; Jason Cong; Yizhou Sun

cs.LGOct 9, 2026

Learnable Low-Rank Polynomial Sketch for Effective Linear Attention

Authors: Zongyue Qin, Haoran Deng, Jason Cong, Yizhou Sun

Organizations: Department of Computer Science, University of California Los Angeles, CA, USA

Abstract

Softmax attention in Transformers suffers from quadratic complexity in sequence length, making it impractical for long-context applications. Linear attention alleviates this issue by replacing the exponential kernel with alternative functions that enable linear-time computation. Among existing linear attention approaches, recent studies have shown that polynomial kernels are particularly effective, as they exhibit sparse and spiky behavior similar to softmax attention, which emphasizes large dot products while suppressing irrelevant interactions. However, the exact computation of high-degree polynomial kernels is infeasible for high-dimensional representations. As a result, prior work relies on approximate polynomial kernels. This introduces a non-negligible approximation error. In this paper, we show that, from the perspective of polynomial kernel approximation, existing linear attention methods are still suboptimal. We propose Learnable Low-Rank Polynomial Sketch (LLoPS), a principled and flexible framework for approximating polynomial kernels with linear attention. Our method learns a low-rank polynomial sketch that provably achieves a strictly smaller approximation error than existing approaches. Experiment results show that LLoPS achieves the highest performance across extensive benchmarks, comparing with various linear attention baselines.

Figures & tables

Figure 1: Illustration of LLoPS with an example where $p=2,d=2,m=3$ . (a) Given the optimal linear projection matrix $\bm{W}$ , each column of $\bm{W}$ can be written as a 2 by 2 matrix, and be approximated by its rank-1 CP decomposition; (b) The rank-1 CP decomposition can be stacked into two weight matrices $\bm{\Theta_{1}},\bm{\Theta_{2}}$ ; (c) In this way, the linear projection $\bm{W}^{\top}\text{vec}(\bm{xx}^{\top})$ can be approximated as $(\bm{\Theta_{1}}^{\top}\bm{x})*(\bm{\Theta_{2}}^{\top}\bm{x})$ where $*$ is an element-wise product.

360M Models
Method	Wino	PIQA	LAM	Hella	ARC-E	ARC-C	FDA	SWDE	SQuAD	Avg
Softmax Attn	51.78	63.71	29.32	33.60	42.85	24.32	57.89	56.08	27.88	43.05
Mamba	50.36	63.76	27.89	33.88	42.42	24.66	5.90	17.37	24.83	32.34
BASED	51.30	63.55	26.64	33.23	41.88	24.23	17.42	27.18	27.41	34.76
BASED+CT	50.91	65.40	27.73	36.85	43.10	25.51	8.35	30.87	26.27	35.00
BASED+ LLoPS	53.04	66.70	29.23	37.07	43.94	25.51	7.53	31.41	31.94	36.26

Table 1: Zero-shot accuracy (%) across diverse NLP benchmarks for pretrained models at two scales. Bold numbers indicate the best performance among linear-complexity methods within each model size. BASED+CT is BASED under the same continual training setting as LLoPS.

Method	Image	Text	Retrieval	ListOps	Pathfinder	Avg
Softmax Attn	64.36	82.96	83.23	48.49	70.45	69.90
Nyström	58.66	66.19	83.90	45.46	69.71	64.78
Skyformer	58.54	65.25	79.55	46.77	72.73	64.57
Hedgehog	52.68	66.59	82.24	41.03	74.16	63.34
PolaFormer	47.74	73.06	80.50	38.86	70.53	62.14
Flatten	51.48	66.13	78.72	38.51	70.27	61.02

Table 2: Performance comparison across Long Range Arena (LRA) benchmarks. Bold numbers indicate the best performance in linear attention methods, underlined numbers indicate the second-best, and italicized numbers indicate the third-best.

$p$	Method	Layer 0	Layer 8	Layer 15	Avg
$2$	PSF	0.186	0.247	0.445	0.293
	NN	0.035	0.070	0.058	0.054
	LLoPS	0.008	0.032	0.025	0.022
$3$	PSF	0.740	0.738	0.745	0.741
	NN	0.059	0.114	0.143	0.105
	LLoPS	0.018	0.040	0.058	0.038

Table 3: Relative mean absolute error of different sketches approximating polynomial kernels with degree

p=2,3

Figure 2: Time cost of different attention kernels at different input length with $p=2$ .

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

d	128	192
BASED	91.9	98.1
LLoPS	93.2	99.0

Appendix

Table 4: Comparion of Based and LLoPS on synthetic associative recall task. For LLoPS,

m=256

p=4

. For both methods

d=128

$p$	2	4	6
Acc	88.7	93.2	86.6

Appendix

Table 5: Accuracy of synthetic association recall task under different polynomial degree

p

with

d=128

Panel of Figure 3 — Figure 3: Results of how accuracy and forward time changes with sketch dimension $m\in\{64,128,192,256\}$ and $d=128$ , $p=4$ .

Explore similar work

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
Flexformer: Flexible Linear Transformer with Learnable Attention Kernel
Jun 26, 2026Haoran Zhang, Feng ZhouFourier Feature Embeddings Long-Context Language Modeling
Center for Applied Statistics and School of Statistics, Renmin University of China Beijing, China
Exact Linear Attention
May 13, 2026Weinuo OuEfficient Transformer Inference Self-Attention
Wuyi 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

Explore similar work

Quasi Linear Kernel Attention with Infinite Capacity

Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

Exact Linear Attention