Transformação em Modelos de Atenção Híbridos

Resumo

Modelos de atenção híbrida melhoram a eficiência em contextos longos ao reter apenas um subconjunto de camadas de atenção completa e substituir as demais camadas por atenção linear. No entanto, a eficácia da conversão de Transformer para híbrido depende criticamente de quais camadas preservam a atenção completa. Os métodos existentes de seleção de camadas híbridas geralmente dependem de estratégias heurísticas, como padrões de posicionamento fixo ou pontuação por camada, tratando implicitamente a importância da camada como isolada e ignorando o efeito de interdependência das camadas sob uma configuração híbrida global. Neste trabalho, formulamos a seleção de camadas híbridas como um problema de otimização de subconjunto com restrição de orçamento. Propomos ainda o FlashMorph (Fast LAyer Selection for Hybrid MORPHing), um método de seleção de camadas eficaz, eficiente e escalável para conversão de Transformer para híbrido. O FlashMorph primeiro constrói um modelo maleável equipando cada camada de atenção completa com um ramo de atenção linear convertido. Em seguida, congela todos os pesos do modelo e otimiza conjuntamente portas por camada em dados sintéticos de recuperação de contexto longo, com uma regularização de linearização que incentiva o modelo a depender da atenção linear para eficiência. As portas aprendidas são discretizadas sob um orçamento de atenção completa predefinido para instanciar a arquitetura híbrida, seguido por destilação de logits padrão e ajuste fino em contexto longo. Extensos experimentos mostram que o FlashMorph descobre configurações híbridas mais eficazes, preserva um forte recall de contexto longo e desempenho geral em benchmarks, enquanto reduz substancialmente o custo de seleção de camadas em comparação com métodos existentes, demonstrando sua eficácia, eficiência e escalabilidade.

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.