Learnable low rank polynomial sketch improves transformer attention efficiency
Learnable Low-Rank Polynomial Sketch for Effective Linear Attention
Machine LearningArtificial IntelligenceComputation and Language
Summary
The way Transformers pay attention to long text often takes a lot of time because they look at every word with every other word. The authors find this slow method can be sped up by changing how attention is calculated using something called polynomial kernels. They introduce a new way to approximate these kernels that is more accurate and efficient, allowing Transformers to handle longer texts faster without losing quality. Their tests show this new method works better than previous similar approaches.
What this means in practice
- •For machine learning engineers: Implement faster Transformer models for processing long sequences in NLP tasks with reduced computational cost.
- •For ai infrastructure teams: Optimize AI systems to handle long context inputs more efficiently using improved linear attention methods.
Authors
Zongyue Qin, Haoran Deng, Jason Cong, Yizhou Sun
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 \textbf{L}earnable \textbf{Lo}w-Rank \textbf{P}olynomial \textbf{S}ketch (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.