Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

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.

• The paper introduces Flexformer, a learnable attention kernel method that generalizes softmax attention using trainable Fourier features.
• The method supports both stationary and nonstationary kernel variants, capturing arbitrary input dependencies while maintaining linear time and space complexity.
• Empirical results on benchmarks like LRA and WikiText-103 show that Flexformer reduces memory usage and speeds up processing, often surpassing traditional quadratic attention.

Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

Motivation and Background

Transformer architectures have become central across NLP, vision, and sequential modeling, primarily due to the power of the self-attention mechanism. However, the vanilla self-attention operation suffers from quadratic time and space complexities with respect to the input sequence length $N$, limiting their scalability for long-context tasks. To address this bottleneck, multiple efficient attention variants have been proposed. Sparse attention schemes restrict context windows, while kernel-based linear attention methods seek to express the attention kernel as inner products in transformed feature spaces, reducing computation to linear complexity in $N$. However, standard kernel-based approaches, such as the Performer and RFA, typically leverage fixed, non-adaptive kernel mappings (e.g., Gaussian random Fourier features) and thus lack adaptability to diverse data distributions. Prior attempts at learnable kernels used hand-crafted families, often with limited expressivity and theoretical grounding.

Flexformer introduces a principled and highly flexible framework for learnable attention kernels, leveraging the spectral characterization of positive definite kernels via Fourier features and treating their frequencies as fully trainable. This approach generalizes previous random feature schemes, encompasses softmax attention as a strict subset, and supports both stationary and nonstationary kernel variants.

Methodology

Spectral-Kernel-Based Linear Attention

The kernelized linear attention paradigm reduces quadratic complexity by expressing the similarity function $k(\mathbf{x},\mathbf{y})$ as an inner product $$k(\mathbf{x},\mathbf{y}) = \phi(\mathbf{x})^\top\phi(\mathbf{y})$$ with a mapping into a feature space. For softmax attention, this kernel is $$k_{\mathrm{SM}}(\mathbf{x},\mathbf{y}) = \exp(\mathbf{x}^\top\mathbf{y}/\sqrt{d})$$, but the exact feature mapping is infinite-dimensional. Random Fourier features approximate stationary kernels by drawing frequencies from the spectral measure induced by the kernel, as per Bochner’s theorem.

Flexformer departs from previous work by learning the spectral measure directly. Rather than sampling $\{\omega_i\}$ from a fixed distribution (e.g., Gaussian), it treats the frequencies as trainable parameters. This allows the data to shape the spectral profile of the kernel, vastly broadening its expressivity.

Stationary and Nonstationary Extensions

• Stationary kernel ($\text{Flexformer}_\text{s}$): Frequencies $\{\omega_i\}$ are learnable, giving rise to attention kernels that generalize Gaussian RFF-based kernels without assuming a particular bandwidth or structure.
• Nonstationary kernel ($\text{Flexformer}_\text{n}$): Joint learnable frequency pairs $(\omega_{1i}, \omega_{2i})$ enable the kernel mapping to capture arbitrary input-pair dependencies, as grounded in the Yaglom extension of Bochner’s theorem.

Both variants include a trainable scaling factor in the norm-based exponent term inherited from the softmax kernel.

Complexity and Implementation Considerations

Flexformer maintains linear time and space complexity $N$0 and $N$1 respectively, where $N$2 is the kernel-induced feature dimension (typically $N$3). Despite the additional learnable parameters required for the frequencies (less than 1% of total model parameters in typical configurations), efficiency and scalability are preserved for long sequences.

Figure 1: Softmax attention mechanism, which exhibits quadratic complexity in sequence length.

Figure 2: Forward pass time and GPU memory usage scale linearly with sequence length for linear methods including Flexformer.

Empirical Evaluation

Long Range Sequence Classification

Flexformer is evaluated on the Long Range Arena (LRA) benchmark suite, which comprises tasks such as ListOps, document retrieval, byte-level text classification, image classification from pixel sequences, and Pathfinder. Across these diverse domains, both stationary and nonstationary Flexformer variants consistently outperform all kernel-based linear attention baselines, and in several settings even surpass (or closely match) the vanilla softmax Transformer in accuracy.

Figure 3: Performance, memory consumption, and training speed of $N$4 on LRA tasks demonstrate substantial speed and memory gains versus quadratic attention, with superior test accuracy compared to other linear kernels.

Autoregressive Language Modeling

Experiments on WikiText-103—using Transformer architectures with both small (41M) and large (247M) parameter counts—reveal that Flexformer matches or exceeds the perplexity of vanilla softmax attention when sufficiently scaled. The nonstationary variant is especially effective, achieving the lowest perplexity in the large model regime.

Figure 4: Performance comparison of $N$5 and other linear attention architectures on language modeling, indicating strong generalization.

Softmax Recovery and Distillation

Flexformer is capable of faithfully recovering softmax attention when distilled from large pretrained models (e.g., RoBERTa) using attention weight distillation. This enables efficient and accurate linear-inference replacements for pretrained quadratic attention models.

Figure 5: Distillation loss curves on MNLI illustrating that Flexformer achieves lower validation cross-entropy than competitive baselines during softmax attention recovery.

Transferability

Flexformer exhibits strong kernel transferability. When distilled on one GLUE task and finetuned on another, attention kernels learned by Flexformer maintain high accuracy across domains, outperforming prior learnable kernel baselines like Hedgehog—especially on challenging tasks such as CoLA.

Theoretical and Practical Implications

Flexformer’s design is rooted in spectral kernel theory, supplying a principled and expressive family of learnable kernels for attention. It strictly subsumes (and can closely approximate) softmax attention while mitigating the computational bottlenecks of quadratic models. On practical long-sequence modeling benchmarks, Flexformer delivers robust improvements in efficiency (up to $N$6 speedup and $N$7 memory reduction) without compromising accuracy.

Theoretically, the approach provides a unifying view, suggesting that optimal attention is not constrained to softmax-like kernels, and can benefit from data-adaptive spectral measures. However, the assumption that positive definite kernels are universally optimal remains to be rigorously justified.

Limitations and Future Directions

The paper acknowledges that the optimality of positive definite attention kernels is yet unsettled, and that the suite of tasks evaluated—while diverse—primarily cover textual and moderate-scale sequence modalities. Extending Flexformer to extremely long protein, genomic, or high-resolution vision data could further test its generality. Future work could also involve non-Fourier kernel parameterizations or exploration of indefinite kernels.

Conclusion

Flexformer represents a substantial advancement in linear transformer architectures by fusing rigorous spectral kernel methods with full data-driven learnability. Its ability to operate at linear computation and memory cost, match (and sometimes surpass) the performance of quadratic softmax attention, recover softmax weights via distillation, and transfer kernels across domains offers a compelling solution for scalable, performant attention in long-context sequence models.

(2606.27748)