Morphing into Hybrid Attention Models

2026-06-29

AI Summary

This work addresses the challenge of efficiently selecting which layers to retain full attention in Transformers when converting them into hybrid attention models, aiming to balance long-context performance with computational efficiency. The authors propose FlashMorph, a method that formulates layer selection as a budget-constrained subset optimization problem. It constructs a deformable architecture featuring parallel full and linear attention branches per layer. After freezing the backbone weights, the approach jointly optimizes inter-layer gating on synthetic long-context data and introduces a linearization regularizer to encourage sparse usage of full attention. The resulting configuration is then discretized into an efficient hybrid model. FlashMorph is the first to enable global inter-layer dependency modeling and joint gating learning, surpassing heuristic strategies by significantly reducing search costs while achieving superior long-context recall, general performance, and inference efficiency compared to baselines.

Abstract

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.

Problem

• hybrid attention
• layer selection
• Transformer
• long-context efficiency
• subset optimization

Innovation

• hybrid attention
• layer selection
• linear attention
• budget-constrained optimization
• FlashMorph