Flexformer: Learnable-Kernel & Compression Models
- Flexformer is a term for two distinct architectures: one that learns flexible Fourier-feature attention kernels for linear Transformers and another that compresses video-text context via an autoencoder.
- The learnable-kernel variant uses trainable spectral frequencies to approximate softmax attention, ensuring linear time and memory complexity while enhancing performance on language tasks.
- The FlexFormer in LoViC employs a self-attention based autoencoder with replicated query tokens and I-RoPE to efficiently compress multimodal context for long video generation.
Searching arXiv for papers titled or mentioning “Flexformer” to ground the article and check for multiple usages. Flexformer denotes two distinct architectures in recent arXiv literature. In the 2026 paper "Flexformer: Flexible Linear Transformer with Learnable Attention Kernel" (Zhang et al., 26 Jun 2026), it is a kernel-based linear Transformer that replaces fixed random-feature kernels with trainable spectral frequencies, thereby learning the attention kernel in a data-driven manner while preserving linear-time, linear-memory scaling in sequence length. In the 2025 paper "LoViC: Efficient Long Video Generation with Context Compression" (Jiang et al., 17 Jul 2025), "FlexFormer" names a Q-Former-inspired autoencoder that jointly compresses video and text context into a small set of position-aware latent tokens for long-horizon diffusion-based video generation. The shared label reflects an emphasis on flexibility, but the two models address different bottlenecks: efficient long-sequence attention in one case, and scalable multimodal context compression in the other.
1. Nomenclature and scope
The term appears in two separate technical lineages.
| Usage | Core mechanism | Primary setting |
|---|---|---|
| Flexformer (Zhang et al., 26 Jun 2026) | Learnable Fourier-feature linear attention | Language modeling and sequence classification |
| FlexFormer (Jiang et al., 17 Jul 2025) | Self-attention autoencoder with replicated query tokens and I-RoPE | Long video generation with joint video-text context |
The 2026 Flexformer paper addresses the quadratic cost of softmax self-attention. Standard self-attention computes
which requires an attention matrix for sequence length , yielding time complexity and memory complexity (Zhang et al., 26 Jun 2026). Its starting point is the kernelized reformulation of attention, in which a similarity kernel is written as an inner product of features, , so that computation can be reordered into linear complexity in .
The 2025 FlexFormer paper in LoViC addresses a different scaling problem: diffusion transformers for long video incur prohibitive cost if all previous spatiotemporal tokens are passed directly into self-attention. There, FlexFormer is introduced as an expressive autoencoder that jointly compresses video and text into unified latent representations, producing context tokens that are concatenated with current-segment tokens inside the DiT (Jiang et al., 17 Jul 2025).
2. Flexformer as a learnable-kernel linear Transformer
The central claim of the 2026 architecture is that prior kernel-based linear attention methods either fix the kernel or permit only weakly learnable, hand-designed families. Performer and Random Feature Attention approximate the softmax kernel with random Fourier features drawn from a fixed distribution, while methods such as Hedgehog and Polaformer introduce learnable parameters but retain restricted functional forms (Zhang et al., 26 Jun 2026). Flexformer removes the assumption that softmax is always the correct similarity function and instead learns the spectral representation of the kernel itself.
The construction begins from the random Fourier feature view of stationary kernels. By Bochner’s theorem, a bounded, continuous, positive-definite stationary kernel admits a spectral representation, and Monte Carlo sampling yields a feature map of the form
so that (Zhang et al., 26 Jun 2026). Traditional random-feature linear attention fixes the frequencies 0 by sampling from a prescribed distribution. Flexformer instead treats them as trainable parameters.
This change is interpreted spectrally. For stationary kernels, learning frequencies corresponds to learning the spectral density 1; for nonstationary kernels, the analogous object is a density over frequency pairs 2 via Yaglom’s representation (Zhang et al., 26 Jun 2026). The attention kernel is therefore no longer hard-coded as a softmax approximation. The paper characterizes the resulting family as broad enough to approximate any bounded, continuous, positive-definite kernel of the appropriate class, stationary or nonstationary depending on the variant (Zhang et al., 26 Jun 2026).
Flexformer also retains the norm-dependent multiplicative factors used in softmax-kernel decompositions. Writing softmax similarity as
3
the kernel can be decomposed into norm terms and a Gaussian term in 4, motivating the feature map
5
for the stationary case (Zhang et al., 26 Jun 2026). This preserves compatibility with standard linear-attention normalization while allowing the spectral part to be learned.
3. Stationary and nonstationary formulations
The 2026 paper defines two variants: 6 for stationary kernels and 7 for nonstationary kernels (Zhang et al., 26 Jun 2026).
In the stationary variant, only the frequencies 8 are learned. Attention is computed in standard kernel-linear form,
9
with 0 defined through trainable Fourier frequencies (Zhang et al., 26 Jun 2026). This preserves the usual reordering trick: one precomputes 1 and 2, then evaluates each query against those aggregates instead of materializing the full similarity matrix.
The nonstationary variant invokes Yaglom’s theorem, according to which a bounded, continuous, positive-definite kernel has a spectral representation over 3. To preserve positive definiteness, the paper uses a symmetrized spectral density and derives a real-valued random-feature expansion from sampled frequency pairs (Zhang et al., 26 Jun 2026). The corresponding feature map is more complex than in the stationary case and is accompanied by a learnable scaling parameter 4 in the norm factor,
5
The nonstationary model learns all paired frequencies 6 together with the scaling, and the paper states that it is strictly more expressive because stationary kernels arise as a special case whereas the converse does not hold (Zhang et al., 26 Jun 2026).
Computationally, both variants preserve linear scaling in sequence length. With feature dimension 7, the reported complexity is time 8 and space 9 (Zhang et al., 26 Jun 2026). In practice the paper sets 0, so 1. The surrounding architecture remains that of a standard Transformer—multi-head attention, residual connections, feed-forward blocks, layer norms, and positional encodings—while only the attention block is altered (Zhang et al., 26 Jun 2026).
The additional parameter cost is described as modest. With 2, stationary Flexformer adds about 3 parameters per head and nonstationary Flexformer about 4 per head; for 5, the total overhead is reported as under 6 of model parameters, with language-modeling overheads of 7–8 depending on variant (Zhang et al., 26 Jun 2026).
4. Training, distillation, and empirical behavior
From-scratch training follows ordinary Transformer practice. The paper states that Transformer weights are initialized as usual, spectral frequencies are typically initialized from a Gaussian so that the initial kernel approximates softmax, and then frequencies are fine-tuned jointly with all other parameters under the task loss; no special regularizers on frequencies are required, though Adam and gradient clipping are used (Zhang et al., 26 Jun 2026).
A prominent use case is attention-weight distillation from a pretrained softmax Transformer. The reported procedure is to take a pretrained and task-finetuned model such as RoBERTa-125M, replace softmax attention with Flexformer attention, freeze all original model parameters, and train only the kernel parameters using an attention-weight distillation loss defined as the cross-entropy between the softmax teacher attention distribution and the Flexformer attention distribution (Zhang et al., 26 Jun 2026). In the GLUE experiments, the paper reports training only the kernel parameters for 9 steps with batch size 0 and learning rate 1, followed optionally by whole-model finetuning (Zhang et al., 26 Jun 2026).
The empirical picture in the paper is organized around long-sequence classification, language modeling, and recovery of pretrained softmax behavior. On Long Range Arena, 2 is reported as the second-best overall method and 3 as the best average performer with 4 average accuracy, described as about 5 relative improvement over the best prior linear method; the nonstationary model is best on Retrieval, Image, and Pathfinder and highly competitive on ListOps and Text (Zhang et al., 26 Jun 2026). The paper also reports that runtime and memory scale linearly with sequence length, with slight overhead relative to other linear-attention methods because of the more complex feature map.
On WikiText-103, the small model results are reported as follows: Transformer achieves dev/test perplexity 6, 7 achieves 8, and 9 achieves 0; among prior linear baselines, Polaformer is reported around 1 and Hedgehog around 2 (Zhang et al., 26 Jun 2026). For the big model, Transformer achieves 3, 4 achieves 5, and 6 achieves 7, which the paper describes as significantly better than Transformer and all baselines (Zhang et al., 26 Jun 2026). This suggests that the learned-kernel formulation benefits more strongly from increased model capacity.
On GLUE, the paper reports that directly replacing softmax attention with unlearnable linear attention such as Performer or RFA causes large performance drops, especially on CoLA and STS-B, whereas distilled Hedgehog and 8 almost fully recover softmax-level performance (Zhang et al., 26 Jun 2026). It further reports that 9 attains a slightly better average than Hedgehog and is best on most individual tasks, including CoLA, where it surpasses the original softmax RoBERTa. Distillation loss curves on MNLI and QNLI are reported to converge faster and lower for Flexformer than for Hedgehog, and cross-domain kernel-transfer experiments show smaller performance degradation for Flexformer, particularly on CoLA (Zhang et al., 26 Jun 2026).
The paper also gives a concrete efficiency figure for document retrieval with sequence length 0 and head dimension 1: 2 is reported as about 3 faster in training than a softmax Transformer and as using about 4 less peak GPU memory (Zhang et al., 26 Jun 2026).
5. FlexFormer in LoViC: multimodal context compression for long video generation
In LoViC, FlexFormer is not an attention-kernel method but a context compressor. The paper identifies it as the module that makes segment-wise long-video generation computationally feasible by compressing arbitrary-length history into a small number of tokens (Jiang et al., 17 Jul 2025). The overall pipeline is: raw frames are encoded by a VAE into latent video tokens, text prompts are encoded into text tokens, previous segments’ video and text tokens are passed through the FlexFormer encoder to produce compressed context tokens, and these context tokens are concatenated with the current segment’s tokens in the DiT (Jiang et al., 17 Jul 2025).
Architecturally, this FlexFormer is described as a Q-Former-inspired autoencoder with an encoder and decoder, both transformer-based (Jiang et al., 17 Jul 2025). For a segment, the model receives 5 flattened video tokens and 6 text tokens. It introduces a single learnable query token 7 which is replicated at runtime into a sequence 8, where 9 is determined by input length and compression strategy (Jiang et al., 17 Jul 2025). The encoder forms the joint sequence
0
and processes it with self-attention, not cross-attention. After the encoder, only the query-token slice is retained as the compressed representation (Jiang et al., 17 Jul 2025).
The paper emphasizes three departures from vanilla Q-Former. First, instead of a fixed bank of learnable queries, there is only one learnable query vector that is copied 1 times, allowing the number of query tokens to vary with input length (Jiang et al., 17 Jul 2025). Second, self-attention replaces cross-attention, which enables the use of RoPE-based positional encodings in a symmetric setting. Third, the replicated query tokens are differentiated spatially and temporally through assigned 3D positions, making them pseudo video tokens rather than an undifferentiated latent set (Jiang et al., 17 Jul 2025).
Compression is adaptive. The paper states that the number of query copies is determined by the total number of input tokens and the specified compression strategy, and it evaluates uniform, linear, and logarithmic strategies (Jiang et al., 17 Jul 2025). In the uniform case, 2 for total input length 3 and average compression ratio 4. The reported experiments include global compression ratios such as 5, 6, and “7,” with the latter meaning that compression linearly changes from 8 for distant frames to 9 near the current segment, yielding an overall average around 0 (Jiang et al., 17 Jul 2025).
A key technical mechanism is Interpolated RoPE (I-RoPE). Naïve treatment of query tokens as 1D tokens under M-RoPE led to blurry reconstructions, which the authors attribute to the lack of an explicit spatial prior (Jiang et al., 17 Jul 2025). I-RoPE assigns each query token a spatiotemporal 3D position by interpolation over the video-token grid:
1
with positions determined by the interpolation scheme and compression strategy (Jiang et al., 17 Jul 2025). Uniform compression tiles query tokens evenly across the grid; linear and log strategies allocate more query density near frames adjacent to the current segment, emphasizing continuity-relevant regions of the temporal context (Jiang et al., 17 Jul 2025).
FlexFormer is trained as a deterministic autoencoder. The paper reports training it from scratch for 2K steps with batch size 3 using MSE loss on reconstructed VAE latents (Jiang et al., 17 Jul 2025). The downstream DiT is then trained with standard flow matching loss, where conditioning includes compressed context from FlexFormer and current text (Jiang et al., 17 Jul 2025).
6. Empirical properties, misconceptions, and limitations
The LoViC paper provides ablations that isolate FlexFormer’s design choices. On video reconstruction, Table 4 compares four variants and reports PSNR/SSIM/LPIPS as follows: M-RoPE + Multiple Query Tokens, 4; M-RoPE + Single Query Token, 5; I-RoPE + Single Query Token, 6; and I-RoPE + MLP Query, 7 (Jiang et al., 17 Jul 2025). The default configuration—single query token with variable copies and I-RoPE—therefore gives the best reported reconstruction metrics.
The compression-strategy ablations make a further point: lower compression is not uniformly better. Table 5 is summarized as showing that reducing the compression ratio from 8 to 9 does not always improve DiT metrics because of redundancy, while linear compression often yields better continuity metrics than uniform compression at similar capacity (Jiang et al., 17 Jul 2025). This directly counters the simplified assumption that allocating more context tokens necessarily improves downstream generation quality.
On long-video generation benchmarks, the LoViC paper reports that LoViC + FlexFormer with 0B parameters outperforms or matches CausVid, SkyReel-V2, FramePack (1B), and LTX-Video-v0.9.6 on PSNR, LPIPS, Subject Consistency, Background Consistency, and Motion Smoothness while using compressed context rather than full context sequences (Jiang et al., 17 Jul 2025). The paper also states that context compression substantially reduces memory and runtime and allows more frames under the same resource constraints.
Across the two literatures, several objective clarifications follow from the cited papers. First, Flexformer in the 2026 sense is not merely a faster approximation to a fixed softmax kernel; the paper explicitly frames softmax as a special case and learns spectral frequencies to access a broader positive-definite kernel family (Zhang et al., 26 Jun 2026). Second, FlexFormer in LoViC is not a variant of the linear-attention model despite the shared name; it is a self-attention autoencoder for joint video-text context compression (Jiang et al., 17 Jul 2025). Third, neither work claims that flexibility eliminates all trade-offs. The linear-attention paper notes dependence on positive-definiteness assumptions, hyperparameter sensitivity in the number of Fourier features 2, and limited evaluation beyond NLP and LRA (Zhang et al., 26 Jun 2026). The video paper notes a trade-off between compression and detail, possible difficulty on extremely long or high-resolution structured scenes, and slightly lower aesthetic and prompt-alignment metrics than the base LTX-Video in some multi-shot cases, reflecting a tendency to prioritize coherence with previous segments over strict text alignment (Jiang et al., 17 Jul 2025).
Taken together, the two usages establish "Flexformer" as a label associated with adaptable internal representations rather than a single canonical architecture. In one case, flexibility is spectral and kernel-theoretic; in the other, it is multimodal, positional, and compression-oriented. A plausible implication is that the shared nomenclature reflects a broader design tendency in efficient sequence modeling: replacing fixed structural assumptions—whether fixed attention kernels or fixed query-token budgets—with learned, task-adaptive mechanisms.