One-sentence Summary

Researchers from Fudan University, ByteDance Seed, and The Chinese University of Hong Kong propose FlashMorph, which morphs Transformers into hybrid attention models by formulating layer selection as a budget-constrained subset optimization using a morphable model with linearization regularization to capture global layer interdependencies, outperforming heuristic strategies at reduced selection cost.

Key Contributions

• Hybrid layer selection is formulated as a budget-constrained subset optimization problem to capture inter-layer dependencies, rather than relying on isolated layerwise scoring.
• FlashMorph constructs a morphable model with linear-attention branches on each full-attention layer, then jointly optimizes layerwise gates under a full-attention budget using linearization regularization on synthetic long-context retrieval data.
• Evaluations on Qwen3 backbones and multiple linear-attention variants show that FlashMorph maintains long-context retrieval and reasoning performance while significantly reducing layer selection cost compared to heuristic and layerwise scoring methods.

Introduction

The Transformer's softmax attention scales quadratically with sequence length, making long-context inference expensive due to growing KV caches. Hybrid models that keep full attention in only a subset of layers while using linear attention elsewhere can mitigate this, but converting a pretrained Transformer into an efficient hybrid requires choosing exactly which layers retain full attention. Prior selection methods rely on heuristics like uniform placement, costly search, or isolated layer scoring, which ignore how layers interact when jointly linearized and often incur large selection overhead. The authors propose FlashMorph, a method that treats hybrid layer selection as a budget-constrained joint optimization. It constructs a morphable model where each attention layer can be continuously varied between full and linear modes via learnable gates, then aligns the gates on synthetic retrieval data to capture inter-layer complementarity and redundancy before discretizing them into a final hybrid architecture, thereby improving the quality-efficiency trade-off while drastically reducing selection cost.

Method

Theauthors introduce FlashMorph, an effective and scalable layer selection method for converting standard Transformer models into hybrid attention architectures. The overall framework formulates hybrid layer selection as a budget-constrained subset optimization problem, moving beyond fixed allocation patterns and isolated layerwise scoring.

The pipeline begins with morphable layer construction. To equip each full-attention layer with a trained linear-attention replacement, the authors distill an all-linear model from the original pretrained Transformer. During this stage, the parameters of the original full-attention model are frozen, and only the linear-attention branches are trained. They optimize these branches using a layerwise hidden-state alignment loss: $\mathcal{L}_{\text{hidden}} = \frac{1}{L} \sum_{l=1}^{L} \left\| \mathbf{H}_{\text{lin}}^{(l)} - \mathbf{H}_{\text{full}}^{(l)} \right\|_{2}^{2}$Lhidden​=L1​∑l=1L​​Hlin(l)​−Hfull(l)​​22​ This yields a trained all-linear student model, providing a linear-attention replacement for every layer in the original architecture.

Next, the authors perform layer selection via joint optimization and linearization regularization. To avoid the intractability of exhaustively evaluating all possible hybrid configurations, they relax the discrete subset-selection problem into a continuous optimization problem over layerwise gates. For each layer $l$l, a scalar gate $\alpha^{(l)}$α(l) interpolates between the full-attention branch and the linear-attention branch: $\mathbf{H}_{\text{mix}}^{(l)} = \alpha^{(l)} \mathbf{H}_{\text{full}}^{(l)} + (1 - \alpha^{(l)}) \mathbf{H}_{\text{lin}}^{(l)}$Hmix(l)​=α(l)Hfull(l)​+(1−α(l))Hlin(l)​ The gate values are optimized while keeping both the full-attention backbone and the trained linear-attention branches frozen. To preserve the behavior of the original model, the authors compute an alignment loss at the answer-token positions: $\mathcal{L}_{\text{align}} = \frac{1}{L |\mathcal{T}(x)|} \sum_{l=1}^{L} \sum_{t \in \mathcal{T}(x)} \left\| \mathbf{H}_{\text{mix}, t}^{(l)} - \mathbf{H}_{\text{full}, t}^{(l)} \right\|_{2}^{2}$Lalign​=L∣T(x)∣1​∑l=1L​∑t∈T(x)​​Hmix,t(l)​−Hfull,t(l)​​22​ To encourage efficiency, a linearization regularization term is introduced to penalize reliance on full attention: $\mathcal{L}_{\text{reg}} = \sum_{l=1}^{L} \alpha^{(l)}$Lreg​=∑l=1L​α(l) The final optimization objective combines these terms: $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{align}} + \lambda \mathcal{L}_{\text{reg}}$Ltotal​=Lalign​+λLreg​ This optimization is performed on synthetic long-context retrieval data, where randomly generated passkeys are inserted at different depths of a long-context document. This provides a targeted signal for measuring whether replacing full attention with linear attention disrupts long-range information access.

After optimization, the learned gate values indicate the relative necessity of preserving full attention in each layer. Given a target budget $K$K of full-attention layers, the authors discretize the hybrid layers by selecting the layers with the largest gate values: $\mathcal{I}_{\text{full}}^{\text{Hybrid}} = \operatorname{TopK} \left( \{ \alpha^{(l)} \}_{l=1}^{L}, K \right)$IfullHybrid​=TopK({α(l)}l=1L​,K) The remaining layers are instantiated with their trained linear-attention replacements.

Finally, the authors apply logits distillation and long-context finetuning to recover the quality of the selected hybrid attention model. The hybrid model is distilled from the original full-attention teacher using the Kullback-Leibler divergence: $\mathcal{L}_{\text{KD}} = D_{\text{KL}} (p_{\text{T}}(\mathbf{X}) \parallel p_{\text{H}}(\mathbf{X}))$LKD​=DKL​(pT​(X)∥pH​(X)) Subsequently, the hybrid model is finetuned on long-context sequences with the standard language modeling objective: $\mathcal{L}_{\text{FT}} = - \sum \log p_{\text{H}}(x_{t} \mid x_{<t})$LFT​=−∑logpH​(xt​∣x<t​) This completes the FlashMorph Transformer-to-hybrid conversion pipeline.

Experiment

FlashMorph employs joint optimization to select full-attention layers far more efficiently than prior selection methods, requiring orders of magnitude fewer selection tokens while achieving strong long-context retrieval and recall across diverse hybrid ratios and linear-attention backbones. It preserves competitive commonsense reasoning and significantly boosts recall-intensive task performance, especially when guided by retrieval-oriented synthetic supervision. The selected hybrid architectures deliver substantial inference speedups and reduced memory usage, enabling scalable long-sequence processing beyond full-attention baselines.

FlashMorph is the only hybrid layer selection method that is both non-isolated and optimization-based, and it achieves the lowest selection cost among search-based approaches. In contrast, other methods sacrifice either the ability to examine layers jointly (non-isolated) or the use of optimization, or they incur much higher token budgets for selection. FlashMorph uniquely satisfies both non-isolated selection and optimization-based selection, while requiring only 20M tokens, far fewer than KL-LS (20B) and PostNAS (50B). HALO has a low 234M token cost but is neither non-isolated nor optimization-based, while KL-LS and PostNAS are optimization-based but fail to be non-isolated or come with extreme selection costs.

FlashMorph delivers strong needle-in-a-haystack retrieval using just 20 million layer-selection tokens. Original models collapse beyond 64K, YaRN fails to restore long-context performance, and prior hybrid methods lose accuracy on harder multi-needle variants at long lengths. FlashMorph achieves near-perfect to perfect single-needle retrieval and substantially better multi-needle accuracy, with gains growing at the 1.7B scale. Original Qwen3 models reach perfect retrieval at 32K and 64K but drop to zero accuracy at 128K and 256K across all NIAH variants. Applying YaRN degrades performance to near zero at all context lengths, never recovering long-context retrieval ability. Uniform hybrid attention preserves high single-needle accuracy but falls to low scores on harder multi-needle settings at 256K (e.g., dropping to the low teens). KL-LS and HALO improve over uniform but still suffer steep declines on NIAH-Single-2 and Single-3 at 256K. FlashMorph achieves near-perfect NIAH-Single-1 on the 0.6B backbone and perfect NIAH-Single-1 on the 1.7B backbone across all contexts. On multi-needle tasks, FlashMorph provides clear gains over prior layer selection methods, especially at short and medium sequence lengths.

On a 0.6B backbone, FlashMorph preserves competitive zero-shot commonsense reasoning across attention backbones while delivering substantially stronger recall-intensive task performance than other efficient attention variants. Under Lightning attention, it achieves the highest recall average among linear methods, requires far fewer selection tokens than alternatives, and closes much of the gap to the original full-attention model. FlashMorph with Lightning attention reaches a recall-intensive average of 60.1, outperforming HALO (55.9) and uniform linear conversion (51.1) on the 0.6B scale. FlashMorph uses only 20M tokens for layer selection, compared to 20B for KL-LS and 234M for HALO, while attaining similar commonsense averages (52.7) to other efficient methods. On the challenging SQuAD task, FlashMorph improves accuracy to 41.7 versus 34.7 for HALO, substantially narrowing the gap to full-attention Qwen3 (44.1).

FlashMorph performs layer selection using only 20 million tokens and 2.1 GPU hours, dramatically undercutting prior methods. It reduces GPU hours by over 1200× compared to PostNAS, 510× compared to KL‑LS, and 7.3× compared to HALO, while requiring orders of magnitude fewer tokens. FlashMorph uses 20M tokens, making token cost 2500× lower than PostNAS (50B) and 1000× lower than KL‑LS (20B). GPU hours drop from 2561.3 (PostNAS) to 2.1 with FlashMorph, a 1219.7× reduction. FlashMorph achieves the lowest FLOPs (2.5e17), significantly below HALO’s 6.5e17 FLOPs.

The experiments evaluate FlashMorph, a hybrid layer selection method that jointly optimizes layer assignments using only 20 million tokens, on long-context needle-in-a-haystack retrieval, zero-shot reasoning, recall-intensive tasks, and computational efficiency. FlashMorph achieves near-perfect single-needle accuracy and substantial multi-needle gains where original models collapse beyond 64K and prior methods like YaRN, uniform hybrid attention, KL-LS, and HALO degrade sharply, while also preserving competitive commonsense reasoning and delivering clear improvements on recall benchmarks such as SQuAD. Critically, FlashMorph completes layer selection with over a 1200-fold reduction in GPU hours and orders-of-magnitude lower token cost compared to search-based alternatives, demonstrating a Pareto-optimal balance between selection cost and downstream performance.