ハイブリッドアテンションモデルへの変容

要旨

ハイブリッド・アテンション・モデルは、フルアテンション層の一部のみを保持し、残りの層を線形アテンションに置き換えることで、長文脈効率を向上させる。しかし、Transformerからハイブリッドへの変換の効果は、どの層にフルアテンションを維持するかに決定的に依存する。既存のハイブリッド層選択手法は、典型的には固定配置パターンや層単位のスコアリングといったヒューリスティックな戦略に依存しており、層の重要度を暗黙的に独立したものとして扱い、グローバルなハイブリッド構成下での層間の相互依存効果を見過ごしている。本研究では、ハイブリッド層選択を予算制約付き部分集合最適化問題として定式化する。さらに、Transformerからハイブリッドへの変換のための、効果的かつ効率的でスケーラブルな層選択手法であるFlashMorph（Fast LAyer Selection for Hybrid MORPHing）を提案する。FlashMorphはまず、各フルアテンション層に変換された線形アテンション分岐を装備することで、モーフィング可能なモデルを構築する。次に、すべてのモデル重みを凍結し、合成された長文脈検索データ上で層単位のゲートを共同最適化し、モデルが効率性のために線形アテンションに依存するよう促す線形化正則化を施す。学習されたゲートは、事前設定されたフルアテンション予算の下で離散化され、ハイブリッド・アーキテクチャを具現化した後、標準的なロジット蒸留と長文脈ファインチューニングが行われる。広範な実験により、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.