Flexformer Introduces Flexible Linear Transformers with Learnable Attention Kernels.

Traditional Transformer models, while effective at capturing long-range dependencies, face a significant limitation due to their attention mechanism's quadratic computational complexity. This complexity restricts their application to very long sequences of data. Kernel-based linear attention offers a solution by reducing this complexity, but often relies on fixed or only weakly learnable kernels, which can limit the model's expressive power and overall performance. Researchers have introduced Flexformer, a novel flexible linear Transformer designed to address these issues. Flexformer learns its attention kernels entirely from data by treating spectral frequencies as trainable parameters. This innovative approach allows the model to adapt and learn a much broader family of attention kernels, significantly enhancing its flexibility and expressiveness. The Flexformer architecture includes both stationary and nonstationary variants, with the latter providing even greater expressive capabilities. Extensive testing on language modeling and sequence classification tasks demonstrates that Flexformer consistently outperforms existing baselines. Furthermore, it can be effectively distilled from pre-trained Transformers to recover softmax attention behavior and exhibits strong kernel transferability across different domains, achieving both high efficiency and competitive performance on tasks involving long sequences.