🤖 machine learning

Flexformer: Flexible Linear Transformer with Learnable Attention Kernel

Flexformer is a flexible linear Transformer that enhances expressiveness and scalability for long sequences by learning attention kernels in a fully data-driven manner through trainable spectral frequencies, consistently outperforming baselines while maintaining efficiency and transferability.

Original authors: Haoran Zhang, Feng Zhou

Published 2026-06-29

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Haoran Zhang, Feng Zhou

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to organize a massive party with thousands of guests. In the world of AI, these guests are "tokens" (pieces of text or data), and the goal is to figure out who needs to talk to whom to understand the whole story.

The Problem: The "Handshake" Bottleneck

Traditional AI models (called Transformers) use a method called Softmax Attention. Think of this as a rule where every single guest at the party must shake hands with every other guest to see who is relevant.

The Good: It's very accurate. Everyone gets a perfect understanding of the group.
The Bad: If you have 1,000 guests, that's 1,000,000 handshakes. If you have 10,000 guests, it's 100,000,000 handshakes. The work grows quadratically. It's like trying to organize a party for a whole city; the computer runs out of memory and time very quickly.

The Old Fix: The "Rigid Script"

To fix this, researchers tried Linear Attention. Instead of everyone shaking hands, they use a shortcut. They assign everyone a "tag" (a feature map) and only compare tags.

The Catch: Most of these shortcuts use a fixed script. They assume there is only one correct way to tag people (usually based on a specific mathematical formula called the "softmax kernel").
The Flaw: Just because a script works for one type of party doesn't mean it works for all. Sometimes the "fixed tag" misses important connections, making the AI less smart.

The New Solution: Flexformer

The authors of this paper propose Flexformer. Think of Flexformer as a smart, customizable tagging system that learns the best tags while the party is happening.

Here is how it works, using simple analogies:

1. The "Radio Tuner" Analogy

Imagine the way the AI connects people is like tuning a radio.

Old Linear Attention: The radio is stuck on one specific frequency. It can only hear one station.
Flexformer: The radio has a tunable dial. Instead of being stuck on one frequency, Flexformer treats the "frequencies" (the mathematical settings that determine how people connect) as learnable knobs.
How it learns: The AI turns these knobs up and down, listening to the data, to find the perfect "station" that captures the most important relationships. It doesn't guess; it learns exactly what works best for the specific text it is reading.

2. The "Flexible Rubber Band"

The paper creates two versions of this system:

Stationary Version (Flexformer_s): Imagine a rubber band that stretches the same way no matter where you pull it. It's flexible, but the rules of stretching are consistent.
Non-Stationary Version (Flexformer_n): This is like a super-flexible, shape-shifting rubber band. It can stretch, twist, and change its rules depending on exactly where you are in the sequence. The paper claims this version is even more powerful because it can adapt to complex patterns that the "consistent" version might miss.

Why is this a big deal?

The paper proves three main things:

It's Fast and Lean: Just like the old linear methods, Flexformer keeps the "handshake" count low. It scales linearly (1,000 guests = 1,000 handshakes, not 1,000,000). This means it can handle very long documents (like entire books or long video transcripts) without crashing the computer.
It's Smarter: Because it learns its own "tags" instead of using a fixed script, it actually understands the data better than the old linear methods. In tests (like reading comprehension and predicting the next word in a sentence), Flexformer beat all the other fast methods and even matched or beat the slow, heavy "gold standard" models.
It Can "Mimic" the Old Way: If you already have a slow, perfect AI model and want to make it fast, you can "distill" (teach) Flexformer to copy the behavior of the slow model. Flexformer can learn to act exactly like the slow, perfect model but run at lightning speed. Furthermore, if you teach it on one topic (like news articles), it can transfer that knowledge to a different topic (like scientific papers) better than other fast models can.

Summary

Flexformer is a new way for AI to read long texts. Instead of forcing the AI to use a rigid, pre-written rulebook for connecting ideas, it gives the AI a set of adjustable dials. The AI learns to turn these dials to find the perfect way to connect ideas, resulting in a system that is fast enough for long documents but smart enough to understand complex details.

Technical Summary: Flexformer

Problem Statement

Transformer models rely on the dot-product attention mechanism to capture long-range dependencies, but this mechanism incurs quadratic time and space complexity ( $O(N^2)$ ) relative to the sequence length $N$ . This bottleneck severely limits scalability for long sequences. While kernel-based linear attention methods reduce complexity to linear ( $O(N)$ ) by approximating the attention mechanism via feature maps, existing approaches face a trade-off between efficiency and expressiveness:

Fixed Kernels: Methods like Performer and RFA rely on random Fourier features to approximate the softmax kernel. However, they treat the kernel as fixed and non-learnable, assuming softmax is optimal, which may not hold in all scenarios.
Weakly Learnable Kernels: Recent works (e.g., Hedgehog, Polaformer) introduce learnable mappings combined with specific functions (exponentials, power functions) to mimic the low-entropy property of softmax. However, these methods do not guarantee that the learned kernel family is sufficiently expressive to encompass a broad range of attention patterns, potentially compromising performance.

The core challenge is constructing a flexible, learnable attention kernel that maintains linear complexity while offering sufficient expressiveness to outperform both fixed kernels and weakly learnable baselines.

Methodology

The authors propose Flexformer, a linear Transformer framework that learns attention kernels in a fully data-driven manner using Random Fourier Features (RFF).

Theoretical Foundation

The method is grounded in spectral representation theory.

Stationary Kernels: Based on Bochner's theorem (Theorem 3.1), a stationary kernel $k(x-y)$ corresponds to a spectral density $p(\omega)$ . Standard RFF approximates this by sampling frequencies $\omega$ from a fixed distribution (e.g., Gaussian for softmax).
Non-Stationary Kernels: Based on the Yaglom theorem (Theorem 3.2), non-stationary kernels can be represented via a joint spectral density.

Flexformer Architecture

Instead of sampling frequencies from a fixed distribution, Flexformer treats the spectral frequencies as trainable parameters optimized end-to-end.

Flexformer $_s$ (Stationary): The frequencies $\{\omega_i\}$ are learnable parameters. This allows the model to learn an optimal spectral density $p(\omega)$ directly from data. The feature map is constructed as:
$\tilde{\phi}_n(x) = \exp\left(\frac{\|x\|^2}{2\sqrt{d}}\right) \phi_n(x)$
where $\phi_n(x)$ consists of cosine and sine terms of the learnable frequencies. This formulation strictly contains the unbiased estimate of the softmax kernel, guaranteeing the ability to recover softmax attention if optimal.
Flexformer $_n$ (Non-Stationary): To increase expressiveness, the authors extend the framework to non-stationary kernels using the Yaglom theorem. This involves learning pairs of frequencies $\{(\omega_{1i}, \omega_{2i})\}$ and a scaling parameter $\tau$ (replacing the fixed $2\sqrt{d}$ ). The feature map becomes:
$\tilde{\phi}_n(x) = \exp\left(\frac{\|x\|^2}{\exp(\tau)}\right) \phi_n(x)$
The non-stationary variant is theoretically strictly more expressive than the stationary one.

Complexity

Both variants maintain linear time complexity $O(Ndd')$ and linear space complexity $O(Nd + Nd' + dd')$ with respect to sequence length $N$ , where $d$ is the hidden dimension and $d'$ is the feature map dimension.

Key Contributions

Learnable Kernel Framework: Proposes Flexformer, which leverages learnable spectral frequencies within a random Fourier feature framework, enabling the model to learn a broad family of attention kernels data-drivenly.
Theoretical Expressiveness: Demonstrates that the Flexformer kernel family strictly contains the softmax kernel (for the stationary variant) and offers strictly greater expressiveness in the non-stationary variant, addressing the limitations of fixed or empirically designed kernels.
Distillation and Transferability: Shows that Flexformer can be effectively distilled from pretrained Transformers to recover softmax attention performance and exhibits strong cross-domain kernel transferability.

Experimental Results

The authors evaluated Flexformer on language modeling, sequence classification, and distillation tasks.

Long-Range Arena (LRA): Flexformer achieved the highest average accuracy across all tasks (ListOps, Text, Retrieval, Image, Pathfinder). Specifically, Flexformer $_n$ attained the best results on document retrieval, image classification, and Pathfinder, outperforming the best existing linear attention baseline by a relative 4.4% in average accuracy.
Autoregressive Language Modeling (WikiText-103):
- In small models (41M params), Flexformer variants outperformed all other linear attention baselines and approached the performance of the vanilla Transformer.
- In large models (247M params), Flexformer $_n$ surpassed the vanilla Transformer, suggesting that learnable kernels benefit significantly from increased model capacity.
Efficiency: On the document retrieval task, Flexformer achieved a 2.6 $\times$ training speedup and an 84% reduction in memory usage compared to the vanilla Transformer, while maintaining linear scaling in time and memory.
Distillation (GLUE Benchmark): When distilling a pretrained RoBERTa model, Flexformer $_n$ nearly fully recovered the softmax attention behavior, outperforming non-learnable linear methods (Performer, RFA) and even surpassing standard softmax on the CoLA dataset.
Kernel Transferability: In cross-domain experiments, kernels learned by Flexformer $_n$ on one GLUE dataset transferred more effectively to other datasets than those learned by Hedgehog, showing smaller performance drops on out-of-domain tasks.

Significance and Claims

The paper claims that Flexformer provides an effective and scalable alternative to softmax attention for long-sequence modeling. Its primary significance lies in resolving the tension between efficiency and expressiveness in linear attention:

It moves beyond the assumption that softmax is optimal by learning the kernel directly from data.
It offers a principled approach (via spectral representation theory) to constructing flexible kernels, rather than relying on empirical heuristics.
It demonstrates that linear attention can match or exceed the performance of quadratic attention while retaining computational efficiency, particularly when scaled or distilled.

Limitations: The authors note that Flexformer assumes the learned attention kernel is positive definite. Whether positive definite kernels are theoretically the optimal choice for attention remains an open question for future research.