Transformation en modèles d'attention hybrides

Résumé

Les modèles d'attention hybrides améliorent l'efficacité en contexte long en ne conservant qu'un sous-ensemble de couches d'attention complète et en remplaçant les couches restantes par de l'attention linéaire. Cependant, l'efficacité de la conversion de Transformer en hybride dépend crucialement du choix des couches qui préservent l'attention complète. Les méthodes existantes de sélection de couches hybrides reposent généralement sur des stratégies heuristiques, telles que des schémas de placement fixes ou un score par couche, traitant implicitement l'importance des couches de manière isolée et négligeant l'effet d'interdépendance des couches dans une configuration hybride globale. Dans ce travail, nous formulons la sélection de couches hybrides comme un problème d'optimisation de sous-ensemble sous contrainte budgétaire. Nous proposons en outre FlashMorph (Fast LAyer Selection for Hybrid MORPHing), une méthode de sélection de couches efficace, efficiente et scalable pour la conversion de Transformer en hybride. FlashMorph construit d'abord un modèle morphable en équipant chaque couche d'attention complète d'une branche d'attention linéaire convertie. Il gèle ensuite tous les poids du modèle et optimise conjointement des portes par couche sur des données synthétiques de récupération en contexte long, avec une régularisation de linéarisation qui encourage le modèle à s'appuyer sur l'attention linéaire pour l'efficacité. Les portes apprises sont discrétisées sous un budget d'attention complète prédéfini pour instancier l'architecture hybride, suivies d'une distillation standard des logits et d'un ajustement fin en contexte long. Des expériences approfondies montrent que FlashMorph découvre des configurations hybrides plus efficaces, préserve un fort rappel en contexte long et des performances générales sur les benchmarks tout en réduisant considérablement le coût de sélection des couches par rapport aux méthodes existantes, démontrant ainsi son efficacité, son efficience et sa scalabilité.

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.