# Morphing into Hybrid Attention Models

> FlashMorph converts Transformers to hybrids via budget-constrained joint optimization of layerwise gates, achieving perfect 256K retrieval with 20M tokens and 2.81x prefill speedup.

- **Source:** [arXiv](https://arxiv.org/abs/2606.30562)
- **Published:** 2026-07-04
- **Permalink:** https://picx.dev/p/H3DGRI

## Summary

## 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

| Model | LS Tokens | NIAH-Single-1 (256K) | NIAH-Single-2 (256K) | NIAH-Single-3 (256K) |
|-------|-----------|---------------------|---------------------|---------------------|
| **1.7B backbone** | | | | |
| Uniform | N/A | 100 | 19.2 | 27.8 |
| PostNAS* | 50B | 99.2 | 73.8 | 57.6 |
| KL-LS | 20B | 94.4 | 34.6 | 3.8 |
| HALO | 234M | 100 | 95.2 | 52.8 |
| **FlashMorph** | **20M** | **100** | **88.2** | **73.2** |

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

| Model | LS Tokens | Avg. Commonsense | Avg. Recall |
|-------|-----------|-----------------|-------------|
| **1.7B + Lightning** | | | |
| Uniform | N/A | 62.9 | 58.1 |
| PostNAS* | 50B | 62.5 | 65.7 |
| HALO | 234M | 62.2 | 68.1 |
| **FlashMorph** | **20M** | **62.1** | **68.1** |

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:**

| Method | Tokens | FLOPs | GPU Hours |
|--------|--------|-------|-----------|
| PostNAS* | 50B | 8.0e20 | 2561.3 |
| KL-LS | 20B | 2.5e20 | 1071.8 |
| HALO | 234M | 6.5e17 | 15.4 |
| **FlashMorph** | **20M** | **2.5e17** | **2.1** |

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

---

_Markdown view of https://picx.dev/p/H3DGRI, served by PicX — AI-generated visual whiteboard summaries of research papers._