Morphing into Hybrid Attention Models

📄 arXiv: 2606.30562v1 📥 PDF

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

分类: cs.CL

发布日期: 2026-06-29


💡 一句话要点

提出FlashMorph以优化Transformer到混合注意力模型的转换

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 混合注意力模型 层选择 Transformer 长上下文 优化算法 线性注意力 深度学习

📋 核心要点

  1. 现有混合注意力模型的层选择方法通常依赖启发式策略,未能充分考虑层间的相互依赖性。
  2. 本文提出FlashMorph,通过构建可变形模型并优化层级门来实现高效的层选择,解决了现有方法的不足。
  3. 实验结果显示,FlashMorph在长上下文回忆和基准性能上表现优异,同时显著降低了层选择成本。

📝 摘要(中文)

混合注意力模型通过保留部分全注意力层并用线性注意力替代其余层来提高长上下文的效率。然而,Transformer到混合模型的转换效果依赖于哪些层保留全注意力。现有的混合层选择方法通常依赖启发式策略,忽视了层间的相互依赖性。本文将混合层选择形式化为预算约束的子集优化问题,并提出FlashMorph(快速层选择用于混合变换),一种有效、快速且可扩展的层选择方法。FlashMorph首先构建一个可变形模型,为每个全注意力层配备一个线性注意力分支,然后在合成长上下文检索数据上共同优化层级门,最终通过标准的logits蒸馏和长上下文微调来实例化混合架构。实验表明,FlashMorph能够发现更有效的混合配置,保持强大的长上下文回忆能力,同时显著降低层选择成本。

🔬 方法详解

问题定义:本文旨在解决Transformer到混合注意力模型转换中层选择的有效性问题。现有方法通常采用固定模式或逐层评分,未能考虑层间的相互影响,导致选择效果不佳。

核心思路:本文将混合层选择视为预算约束的子集优化问题,提出FlashMorph,通过构建可变形模型和优化层级门来实现高效的层选择,确保模型在长上下文任务中的表现。

技术框架:FlashMorph的整体架构包括构建可变形模型、冻结模型权重、在合成长上下文数据上优化层级门、离散化门以实例化混合架构,以及进行logits蒸馏和长上下文微调。

关键创新:FlashMorph的主要创新在于其层选择方法的系统性和优化过程的高效性,区别于现有方法的启发式选择,能够更好地适应长上下文任务。

关键设计:在设计中,FlashMorph使用线性化正则化来鼓励模型在效率上依赖线性注意力,同时设定全注意力预算以离散化学习到的门,确保最终模型的有效性和可扩展性。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明,FlashMorph在长上下文回忆能力上优于现有方法,且在基准测试中表现出色,层选择成本显著降低,展示了其有效性和可扩展性。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的长文本理解、机器翻译和对话系统等。FlashMorph的高效层选择方法能够在保持性能的同时显著降低计算成本,具有广泛的实际价值和未来影响。

📄 摘要(原文)

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.