演变为混合注意力模型

摘要

混合注意力模型通过仅保留部分全注意力层而将剩余层替换为线性注意力，从而提升长上下文处理效率。然而，Transformer到混合架构的转换效果关键取决于哪些层保留全注意力。现有的混合层选择方法通常依赖启发式策略（如固定放置模式或逐层评分），隐含地将层重要性视为独立因素，忽略了全局混合配置下层的相互依赖效应。本研究将混合层选择形式化为预算约束的子集优化问题，并进一步提出FlashMorph（面向混合架构转换的快速层选择方法），这是一种高效、可伸缩的Transformer到混合架构转换层选择方法。FlashMorph首先通过为每个全注意力层配备转换后的线性注意力分支来构建可变形模型，随后冻结所有模型权重，在合成长上下文检索数据上联合优化逐层门控，并引入线性化正则化以鼓励模型依赖线性注意力提升效率。学习得到的门控在预设全注意力预算下进行离散化以实例化混合架构，随后进行标准logits蒸馏与长上下文微调。大量实验表明，与现有层选择方法相比，FlashMorph能够发现更有效的混合配置，在显著降低层选择成本的同时保持强大的长上下文召回能力与通用基准性能，验证了其有效性、效率与可扩展性。

English

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.