Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

2026-06-26

AI Summary

Traditional Transformers struggle with long sequences due to the quadratic complexity of their attention mechanism, while existing linear attention methods suffer from limited expressivity owing to fixed or weakly learnable kernels. This work proposes Flexformer, a flexible linear Transformer that treats the spectral frequencies in random Fourier features as trainable parameters, enabling data-driven learning of attention kernels in both stationary and non-stationary variants. As the first fully learnable linear attention kernel, its non-stationary form offers enhanced representational capacity and supports efficient knowledge distillation from pretrained Transformers as well as cross-domain transfer. Experiments demonstrate that Flexformer consistently outperforms baseline models on language modeling and sequence classification tasks, achieving both high efficiency and competitive performance on long sequences while effectively approximating softmax attention behavior.

Abstract

Transformer models rely on attention mechanism to capture long-range dependencies but suffer from quadratic complexity, limiting their scalability to long sequences. Kernel-based linear attention reduces this complexity but typically relies on fixed or weakly learnable kernels, restricting expressiveness and performance. In this work, we propose Flexformer, a flexible linear Transformer that learns attention kernels in a fully data-driven manner. Flexformer builds on random Fourier feature-based linear attention and treats spectral frequencies as trainable parameters, enabling the model to learn a broad family of attention kernels. We develop both stationary and nonstationary variants, with the latter offering strictly greater expressiveness. Extensive experiments on language modeling and sequence classification demonstrate that Flexformer consistently outperforms baselines. Moreover, Flexformer can be effectively distilled from pretrained Transformers to recover softmax attention and exhibits strong kernel transferability across domains, achieving both high efficiency and competitive performance on long-sequence tasks.

Problem

• linear attention
• attention kernel
• Transformer
• long-range dependencies
• quadratic complexity

Innovation

• learnable attention kernel
• linear attention
• random Fourier features
• nonstationary kernel
• attention distillation