Summary (Overview)

• FlashMorph proposes a novel formulation of hybrid layer selection for Transformer-to-hybrid conversion as a budget-constrained subset optimization problem, moving beyond heuristic fixed placement rules and isolated layerwise scoring.
• The method constructs a morphable model by equipping each full-attention layer with a trained linear-attention branch, then jointly optimizes layerwise gates on synthetic long-context retrieval data with a linearization regularization.
• FlashMorph achieves strong long-context retrieval performance (e.g., perfect NIAH-Single-1 accuracy at 256K on 1.7B backbone) while using only 20M layer-selection tokens—orders of magnitude fewer than prior methods (e.g., 20B for KL-LS, 50B for PostNAS).
• The method demonstrates substantial efficiency gains: 2.81× prefill speedup at 256K tokens and 2.07× decode speedup at 512K tokens compared to the full-attention baseline, while using significantly less GPU memory.
• Extensive experiments across Qwen3-series models (0.6B to 30B-A3B) and multiple linear-attention variants (Lightning Attention, GLA, GDN) show FlashMorph's effectiveness, efficiency, and scalability across diverse settings.

Introduction and Theoretical Foundation

The Transformer architecture has become the dominant backbone of modern LLMs, but its reliance on softmax attention introduces a fundamental efficiency bottleneck: quadratic growth in attention computation and linear growth in the KV cache with sequence length. Linear attention and state-space models reduce this cost but are generally less effective on long-context and recall-sensitive tasks. Hybrid attention models retain full attention in a subset of layers while replacing remaining layers with efficient linear sequence mixers, achieving a favorable trade-off between quality and efficiency.

Transformer-to-hybrid Conversion starts from a pretrained Transformer-based LLM, retains a small subset of full-attention layers, and replaces the remaining layers with linear attention through parameter transfer, distillation, and finetuning. The critical challenge is layer selection: under a fixed budget $K$ of retained full-attention layers from an $L$-layer model, identifying the optimal hybrid configuration requires evaluating $\binom{L}{K}$ possible subsets, which is computationally intractable.

Existing methods rely on heuristic approximations:

• Uniform interleaving imposes fixed attention placement rules, ignoring heterogeneous functional roles of layers.
• Search-based methods (e.g., PostNAS) explore placements through auxiliary architecture search but introduce substantial overhead.
• Layerwise methods (e.g., KL-LS, HALO) estimate marginal utility of each layer by perturbing one layer at a time, implicitly treating layer importance as isolated and overlooking interdependent layer effects.

The central question motivating this work: Can hybrid layer selection be formulated as a budget-constrained joint optimization problem that globally accounts for interdependent layer effects, rather than relying on heuristic approximations?

Methodology

Problem Formulation

Given a pretrained Transformer model with $L$ attention layers, let $I_{\text{full}} \subseteq [L] = \{1, \ldots, L\}$ denote the retained full-attention layers. Under a fixed budget $K = |I_{\text{full}}|$, the ideal layer selection objective is:

$$I_{\text{full}}^\star = \arg\max_{I_{\text{full}} \subseteq [L], |I_{\text{full}}| = K} \text{Score}\left(M(I_{\text{full}})\right)$$

This is inherently a subset optimization problem that requires evaluating $\binom{L}{K}$ possible subsets.

FlashMorph Pipeline

1. Morphable Layers Construction (Stage 1): Each full-attention layer is equipped with a trained linear-attention branch. The linear-attention branches are trained to imitate layerwise representations of the frozen full-attention teacher using hidden-state alignment:

$$\mathcal{L}_{\text{hidden}} = \frac{1}{L} \sum_{l=1}^{L} \left\| H_{\text{lin}}^{(l)} - H_{\text{full}}^{(l)} \right\|_2^2$$

This yields a morphable model enabling arbitrary full/linear layer configurations.

2. Layer Selection via Joint Optimization (Stage 2): For each layer $l$, a scalar gate $\alpha^{(l)} \in [0, 1]$ interpolates between the full-attention and linear-attention branches:

$$H_{\text{mix}}^{(l)} = \alpha^{(l)} H_{\text{full}}^{(l)} + (1 - \alpha^{(l)}) H_{\text{lin}}^{(l)}$$

Both the full-attention backbone and linear-attention branches are frozen; only the gates $\alpha = \{\alpha^{(l)}\}_{l=1}^L$ are optimized. The optimization objective is:

$$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{align}} + \lambda \mathcal{L}_{\text{reg}}$$

where:

• $\mathcal{L}_{\text{align}} = \frac{1}{L|\mathcal{T}(x)|} \sum_{l=1}^L \sum_{t \in \mathcal{T}(x)} \left\| H_{\text{mix},t}^{(l)} - H_{\text{full},t}^{(l)} \right\|_2^2$ (alignment at answer-token positions)
• $\mathcal{L}_{\text{reg}} = \sum_{l=1}^L \alpha^{(l)}$ (linearization regularization with $\lambda = 0.1$)

Synthetic retrieval data is used for optimization: randomly generated passkeys are inserted at different depths of a long-context document, and the model must recover the corresponding passkeys at the end of the sequence.

3. Discretization: After optimization, layers with the largest gate values are selected as full-attention:

$$I_{\text{full}}^{\text{Hybrid}} = \text{TopK}\left(\{\alpha_l\}_{l=1}^L, K\right)$$

4. Distillation and Long-context Finetuning (Stages 3-4): Standard logits distillation (KL divergence) and long-context finetuning recover the quality of the selected hybrid model.

Empirical Validation / Results

Experimental Setup

• Backbones: Qwen3-0.6B, Qwen3-1.7B, Qwen3-8B, Qwen3-30B-A3B
• Linear-attention variants: Lightning Attention, Gated Linear Attention (GLA), Gated DeltaNet (GDN)
• Baselines: Uniform interleaving, PostNAS, KL-LS, HALO
• Default hybrid ratio: 3:1 (linear:full)

Needle-in-a-Haystack (NIAH) Results

FlashMorph achieves perfect accuracy on NIAH-Single-1 across all context lengths (32K-256K) on the 1.7B backbone, and substantially improves performance on the more challenging NIAH-Single-2 and NIAH-Single-3 settings.

Commonsense Reasoning and Recall-intensive Tasks

FlashMorph preserves competitive commonsense reasoning while achieving the best or tied-best recall averages across attention backbones.

Efficiency Results

Inference Efficiency (1.7B backbone):

• Prefilling: 2.24× speedup at 128K, 2.81× at 256K
• Decoding: 1.56× speedup at 256K, 2.07× at 512K
• FlashMorph scales to 512K-token prefilling and 1M-token decoding on a single GPU, while Qwen3-1.7B encounters OOM

Layer Selection Cost:

FlashMorph reduces selection cost by 7.3× compared to HALO, 510.4× compared to KL-LS, and 1219.7× compared to PostNAS.

Analysis

Robustness across Hybrid Ratios: FlashMorph consistently achieves strong RULER performance across linear:full ratios of 6:1, 3:1, and 1:1, with the advantage most pronounced under sparse full-attention allocation.

Effect of Supervision: FlashMorph with language-modeling supervision already outperforms prior methods, while synthetic retrieval data supervision further improves RULER scores (from 57.2 to 59.0 on GLA, from 61.6 to 64.7 on GDN).

Theoretical and Practical Implications

Theoretical Contributions:

• Formulates hybrid layer selection as a budget-constrained subset optimization problem, providing a principled framework that accounts for inter-layer dependencies, redundancy, and complementarity.
• Introduces a continuous relaxation of the discrete selection problem through learnable gates, enabling joint optimization under a global hybrid configuration.
• Demonstrates that synthetic retrieval data provides a targeted selection signal for measuring whether replacing full attention with linear attention disrupts long-range information access.

Practical Implications:

• FlashMorph enables efficient Transformer-to-hybrid conversion with minimal selection overhead (20M tokens vs. 20B-50B for prior methods), making it practical for large-scale models.
• The method achieves strong long-context retrieval performance while maintaining general reasoning capabilities, offering a practical path to efficient LLM deployment.
• The substantial inference speedups (2-3×) and memory savings enable deployment on resource-constrained hardware and support longer context lengths.
• The approach is agnostic to the choice of linear-attention variant, making it broadly applicable across different efficient attention mechanisms.

Conclusion

FlashMorph presents an effective, efficient, and scalable layer selection method for converting pretrained Transformers into hybrid attention models. By formulating hybrid layer selection as a budget-constrained joint optimization problem and constructing morphable attention layers with lightweight learnable gates, FlashMorph accounts for inter-layer dependencies, redundancy, and complementarity under a global hybrid configuration. Extensive experiments across multiple model scales and linear-attention variants demonstrate that FlashMorph preserves strong long-context retrieval and recall-intensive performance, maintains competitive commonsense reasoning ability, and substantially reduces layer selection overhead compared to existing methods.

Future Directions:

• Extending FlashMorph to other architecture conversion scenarios beyond attention linearization
• Exploring adaptive hybrid ratios that vary with input characteristics
• Investigating the application of joint optimization-based selection to other architectural decisions in LLM deployment