Kernel attention methods with infinite capacity and faster computation
Quasi Linear Kernel Attention with Infinite Capacity
Machine Learning
Summary
Transformer models, which help computers understand language and sequences, often get slow as the information grows longer because they involve complex calculations. The authors look for ways to keep the attention mechanism powerful while making it run faster on long inputs. They introduce a way to measure how expressive different attention methods are and find some popular fast alternatives lose expressiveness. To fix this, they design new kernel-based attention methods that remain very expressive but run almost as fast as the simpler options. Their experiments show these new methods work efficiently for very long sequences.
What this means in practice
- •For machine learning engineers: Implement efficient attention for transformers to handle longer sequences with high expressivity and reduced computational cost.
- •For natural language processing teams: Process and analyze long text sequences more quickly without losing detail in attention-based models.
Authors
Nicolaj Rux, Johannes Hertrich, Sebastian Neumayer
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.