PaperlayerView exact revision on arXiv

Source-linked AI summary

Morphing into Hybrid Attention Models

Disen Lan, Jianbin Zheng, Yuxi Ren, Xin Xia, Xuanda Wang, Xuefeng Xiao, Xipeng Qiu, Yu Cheng

arXiv:2606.30562v1cs.CL

TL;DR

Selecting which layers retain full attention is difficult because layer effects are interdependent under a fixed hybrid-attention budget. FlashMorph jointly optimizes layer gates to choose the retained layers, improving the quality–efficiency trade-off while substantially reducing selection cost.

  • Problem

    Transformer-to-hybrid conversion must allocate a limited full-attention budget across layers, but existing methods rely on fixed placement or isolated layerwise scoring.

  • Method

    FlashMorph jointly optimizes frozen layerwise gates in a morphable model with linearization regularization, then discretizes them under a preset full-attention budget.

  • Results

    Across Qwen3 backbones and multiple attention variants, FlashMorph preserves strong retrieval and recall performance, maintains competitive reasoning, and substantially reduces layer-selection cost.

  • Takeaways & Limitations

    FlashMorph provides an effective, efficient, and scalable approach for selecting hybrid architectures from pretrained Transformers.

  • Takeaways & Limitations

    Exhaustively evaluating all layer subsets is computationally infeasible for modern large language models.

Abstract

from arXiv · show

Hybrid attention models improve long-context efficiency by retaining only a subset of full-attention layers and replacing the remaining layers with linear attention. However, the effectiveness of Transformer-to-hybrid conversion critically depends on which layers preserve full attention. Existing hybrid layer selection methods typically rely on heuristic strategies such as fixed placement patterns or layerwise scoring, implicitly treating layer importance as isolated and overlooking the interdependent layer effect under a global hybrid configuration. In this work, we formulate hybrid layer selection as a budget-constrained subset optimization problem. We further propose FlashMorph (Fast LAyer Selection for Hybrid MORPHing), an effective, efficient and scalable layer selection method for Transformer-to-hybrid conversion. FlashMorph first constructs a morphable model by equipping each full-attention layer with a converted linear-attention branch. It then freezes all model weights and jointly optimizes layerwise gates on synthetic long-context retrieval data, with a linearization regularization that encourages the model to rely on linear attention for efficiency. The learned gates are discretized under a preset full-attention budget to instantiate the hybrid architecture, followed by standard logits distillation and long-context finetuning. Extensive experiments show that FlashMorph discovers more effective hybrid configurations, preserves strong long-context recall and general benchmark performance while substantially reducing layer selection cost compared with existing layer selection methods, demonstrating its effectiveness, efficiency, and scalability.

1 Introduction

Transformer-to-hybrid conversion reduces the cost of full attention, but selecting retained layers is a combinatorial problem whose quality depends on interdependent layer effects. FlashMorph addresses this through budget-constrained joint gate optimization in a morphable model, improving the quality–efficiency trade-off while reducing selection cost.

  • Motivation: Hybrid attention retains full attention in selected layers and replaces the rest with efficient linear sequence mixers to improve the quality–efficiency trade-off.This approach avoids the prohibitive expense of training high-quality hybrid LLMs from scratch by converting pretrained Transformers.
  • Problem: Selecting K full-attention layers from an L-layer model requires evaluating all subsets, making the problem intractable and motivating heuristic approximations.Uniform interleaving ignores pretrained-model information and heterogeneous layer roles, while other methods search layer placements through auxiliary architecture search or supernet training.
  • Problem formulation: Hybrid layer selection is formulated as budget-constrained joint subset optimization that accounts for collective, interdependent effects rather than isolated layer scores or fixed placement rules.The objective balances quality and efficiency while modeling complementarity and redundancy among retained or converted layers.
  • FlashMorph: FlashMorph equips every pretrained full-attention layer with a converted linear-attention branch and jointly optimizes frozen layerwise gates on synthetic retrieval data.The gates estimate which layers need full attention under a global hybrid configuration, after which the selected architecture can be instantiated.
  • Experiments: FlashMorph improves Transformer-to-hybrid conversion across Qwen3-series models and multiple linear-attention variants while substantially reducing layer selection cost.Experiments cover long-context retrieval, commonsense reasoning, and recall-intensive tasks, demonstrating effectiveness, efficiency, and scalability.

2 Preliminaries

Full attention precisely models pairwise token interactions but incurs quadratic sequence-length computation and growing KV-cache memory, whereas linear attention offers linear processing with constant-size state caching. Hybrid attention retains full attention in a fixed subset of layers to balance retrieval effectiveness against computational and memory efficiency.

  • Full Attention: Full attention compares each query with all previous keys, enabling precise matching and long-range dependency modeling.Under causal masking, each position attends to preceding tokens through a weighted sum of their values.
  • Full Attention: O(T^2) computation and KV-cache memory that grows linearly with context length make full attention expensive for long sequences.These costs arise because every query attends to all previous keys during autoregressive inference.
  • Linear Attention: Linear attention replaces the softmax kernel with recurrently accumulated feature-map statistics, represented through a state-transition or decay matrix.Many variants maintain key-value statistics as a recurrent memory state after absorbing the feature map into queries and keys.
  • Linear Attention: O(T) sequence processing and constant-size state caching allow linear attention to avoid storing all previous keys and values.This efficiency comes from the recurrent formulation used during autoregressive decoding.
  • Hybrid Attention: Hybrid attention retains full attention in Ifull, a subset of layers, while replacing the remainder with linear attention to balance effectiveness and efficiency.Under a fixed full-attention budget K = |Ifull|, retaining more full-attention layers preserves retrieval and precise matching but increases computation and memory cost.

3 Morphing into Hybrid Attention Models

FlashMorph frames Transformer-to-hybrid conversion as a budget-constrained subset optimization problem and jointly learns layerwise full-attention preferences in a frozen morphable model. It then discretizes the learned gates and applies distillation and long-context finetuning to recover hybrid-model quality.

  • 3.1 Problem Formulation: Hybrid selection chooses K retained full-attention layers from L candidates, but exhaustive evaluation of the possible subsets is computationally infeasible.Existing approximations use fixed interleaving or isolated layerwise estimation, overlooking inter-layer dependencies.
  • 3.2 Morphable Model Construction: FlashMorph constructs a morphable model by pairing frozen full-attention layers with trained linear-attention replacements obtained through layerwise hidden-state alignment.The resulting model supports arbitrary full/linear layer configurations for subsequent selection.
  • 3.3 Optimization-Based Layer Selection: FlashMorph jointly optimizes continuous gates for all layers under a global hybrid configuration while freezing the backbone and linear-attention branches.Gates interpolate between attention types, with larger values indicating stronger reliance on full attention and only gate values being trainable.
  • 3.3 Optimization-Based Layer Selection: The gate objective aligns morphed and full-attention hidden states while linearization regularization penalizes reliance on full attention, with λ set to 0.1 by default.After optimization, the K layers with the largest gate values are retained as full attention, followed by logits distillation and long-context finetuning.
  • 3.3 Optimization-Based Layer Selection: Synthetic long-context retrieval data provides targeted supervision for learning which layers must retain full attention while capturing inter-layer dependency, redundancy, and complementarity.Random passkeys are inserted at varying depths and recovered at sequence ends; the data optimizes gates rather than model weights.

4 Experiments

Experiments show that FlashMorph selects effective hybrid architectures that preserve long-context retrieval and general task performance while substantially improving inference and layer-selection efficiency. Its benefits remain robust across hybrid ratios and supervision settings.

  • Needle-in-a-Haystack: FlashMorph achieves strong NIAH retrieval across 0.6B and 1.7B backbones, outperforming conventional Transformer and YaRN-based alternatives as context length increases.Hybrid conversion improves extended-context retrieval, with FlashMorph reaching near-perfect accuracy on NIAH-Single-1 for the 0.6B backbone.
  • Commonsense Reasoning and Recall-intensive Tasks: FlashMorph preserves competitive zero-shot commonsense reasoning and achieves the highest recall average on the 0.6B backbone across Lightning Attention, GLA, and GDN.On the 1.7B backbone, it remains close to the strongest baselines across all three attention backbones.
  • Prefilling: 2.81× speedup at 256K tokens: FlashMorph accelerates prefilling and uses less GPU memory, enabling 512K-token prefilling while Qwen3-1.7B runs out of memory.At 128K tokens, prefilling speedup is 2.24×; FlashMorph is comparable to Qwen3 at short sequence lengths.
  • Decoding: 2.07× speedup at 512K tokens: FlashMorph has flatter decoding latency and memory growth, remains executable at 1M decoding length, and outscales Qwen3-1.7B.At 256K tokens, decoding speedup is 1.56×, with prefilling fixed at 1K tokens.
  • Robustness and supervision: FlashMorph remains effective under 6:1, 3:1, and 1:1 linear:full ratios, while synthetic passkey-based supervision further improves performance over language-modeling supervision.The ratio study varies only layer selection while keeping post-selection distillation and finetuning unchanged.

5 Related Work

Related work addresses full attention’s quadratic cost through linear RNN-style mixers and hybrid architectures, while Transformer-to-hybrid conversion uses weight transfer, distillation, and finetuning. Existing layer-selection strategies use fixed interleaving or layerwise estimates, whereas FlashMorph targets inter-layer effects and more efficient selection.

  • Efficient Attention Architectures: Linear RNN architectures reduce full attention’s quadratic computation and memory demands for efficient decoding and scalable long-context processing.Fully replacing attention can create a fixed-state memory bottleneck that limits recall-intensive operations such as associative retrieval.
  • Hybrid Attention Models: Hybrid attention models retain a small subset of full-attention layers for global information access while using linear RNN mixers to improve inference efficiency.Most existing hybrid architectures are designed and pretrained from scratch, often with fixed allocation patterns such as uniform interleaving.
  • Transformer-to-Hybrid Conversion: Transformer-to-hybrid conversion replaces full-attention layers with linear attention through weight transfer, followed by distillation and continued finetuning to recover quality.This setting permits only a limited number of layers to retain full attention.
  • Hybrid Layer Selection: Existing layer-selection methods rely on uniform interleaving or layerwise importance estimation, while FlashMorph uses jointly optimized layerwise gates and linearization regularization to account for inter-layer dependencies, redundancy, and complementarity.This avoids repeated layer-by-layer replacement and evaluation, making selection substantially more efficient.

6 Conclusion

FlashMorph is an effective, efficient, and scalable method for converting pretrained Transformers into hybrid attention models. It selects layers through budget-constrained joint optimization that accounts for inter-layer dependencies, redundancy, and complementarity.

  • Method: FlashMorph converts pretrained Transformers into hybrid attention models through effective, efficient, and scalable layer selection.The method is presented as the paper’s central approach.
  • Optimization: FlashMorph formulates hybrid layer selection as a budget-constrained joint optimization problem rather than using fixed placement rules or isolated layerwise scoring.This formulation accounts for inter-layer dependencies, redundancy, and complementarity.
  • Optimization: FlashMorph constructs morphable attention layers and optimizes lightweight layerwise gates for hybrid model conversion.The gates are optimized within the joint layer-selection framework.

Appendix · A Model and Training Configuration

The appendix specifies FlashMorph’s model and training configuration, including data budgeting, optimization settings, and a four-stage Transformer-to-hybrid conversion pipeline.

  • A Model and Training Configuration: Table 5 reports the complete model and training configuration.The passage introduces the appendix’s consolidated configuration table.
  • A Model and Training Configuration: All training data are randomly sampled from the DCLM corpus [33].The corpus is used throughout the training setup described here.
  • A Model and Training Configuration: Each stage uses only a specified-size token subset to control data budgets across model variants and training stages.Token-subset sizing provides the stated mechanism for budget control.
  • A Model and Training Configuration: All training stages use AdamW [39] with beta values of (0.9, 0.95).The optimizer and beta settings are fixed throughout training.
  • A Model and Training Configuration: Weight decay is set to 0.0 throughout all training stages.The configuration keeps weight decay constant across stages.
  • A Model and Training Configuration: Following HALO [11], FlashMorph’s conversion pipeline consists of hidden-state alignment, layer selection, distillation, and logits distillation.The passage lists four stages, with the final stage described as “lo…” in the supplied text; the exact fourth-stage wording is therefore incomplete.

B Implementation Details

FlashMorph selects layers using a synthetic long-context retrieval dataset built from DCLM text with randomly inserted passkeys. The dataset varies passkey placement and context length to support layer-selection training across long contexts.

  • Synthetic Retrieval Data: The synthetic retrieval dataset uses DCLM text as background context and inserts ten randomly generated passkey sequences at different depths.Each passkey contains 32 words sampled from a fixed alphabet.
  • Synthetic Retrieval Data: Context lengths are randomly sampled from 50 intervals spanning 1K to 16K tokens.This sampling exposes layer selection to varied long-context retrieval conditions.
  • Synthetic Retrieval Data: Passkey insertion depth is discretized into 1,000 levels within each synthetic retrieval example.

C More Experiment Results

FlashMorph is further evaluated on Qwen3-8B and Qwen3-30B-A3B under HypeNet, with comparisons against uniform interleaving and HALO. The reported layer-selection results identify the top-25% layers retained as full-attention layers under a fixed hybrid-ratio budget.

  • C More Experiment Results: FlashMorph is evaluated on Qwen3-8B and Qwen3-30B-A3B under the HypeNet setting, compared with uniform interleaving and HALO.PostNAS and KL-LS are excluded at these scales because reproducing their layer-selection procedures would be prohibitively expensive, especially for larger models.
  • C More Experiment Results: Table 8 sorts layer indices from most important to least important and boxes the top-25% prefix retained as full-attention layers.The boxed prefix corresponds to the fixed hybrid ratio budget.
  • C More Experiment Results: Red indices in baseline rows mark selected layers not shared with FlashMorph, while asterisks denote selected-layer results taken from prior work.These annotations clarify how the layer selections are compared across methods.

D Complete Layer Importance Ranking

Table 8 reports complete layer-importance rankings for each backbone and linear-attention variant, with the highest-ranked layers retained as full-attention layers under a fixed hybrid budget.

  • Table 8 provides complete layer-importance rankings for each backbone and linear-attention variant.
  • Layers are ordered from most to least important, and the top-ranked layers are retained as full-attention layers under the fixed hybrid budget.
Loading 2606.30562v1…