Transformación en Modelos de Atención Híbridos

Resumen

Los modelos de atención híbrida mejoran la eficiencia en contextos largos al retener solo un subconjunto de capas de atención completa y reemplazar las capas restantes con atención lineal. Sin embargo, la efectividad de la conversión de Transformer a híbrido depende críticamente de qué capas conservan la atención completa. Los métodos existentes de selección de capas híbridas suelen basarse en estrategias heurísticas, como patrones de colocación fijos o puntuación por capas, tratando implícitamente la importancia de cada capa como aislada y pasando por alto el efecto de interdependencia entre capas bajo una configuración híbrida global. En este trabajo, formulamos la selección de capas híbridas como un problema de optimización de subconjuntos con restricción de presupuesto. Además, proponemos FlashMorph (Selección Rápida de Capas para el MORFEO Híbrido), un método de selección de capas efectivo, eficiente y escalable para la conversión de Transformer a híbrido. FlashMorph primero construye un modelo maleable equipando cada capa de atención completa con una rama de atención lineal convertida. Luego, congela todos los pesos del modelo y optimiza conjuntamente compuertas por capas en datos sintéticos de recuperación de contexto largo, con una regularización de linealización que fomenta que el modelo dependa de la atención lineal para lograr eficiencia. Las compuertas aprendidas se discretizan bajo un presupuesto preestablecido de atención completa para instanciar la arquitectura híbrida, seguido de destilación de logits estándar y ajuste fino de contexto largo. Experimentos exhaustivos muestran que FlashMorph descubre configuraciones híbridas más efectivas, preserva un buen recuerdo de contexto largo y un rendimiento general de referencia, al tiempo que reduce sustancialmente el costo de selección de capas en comparación con los métodos existentes, demostrando su efectividad, eficiencia y escalabilidad.

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.