Преобразование в гибридные модели внимания

Аннотация

Гибридные модели внимания повышают эффективность работы с длинными контекстами, сохраняя лишь часть слоёв полного внимания и заменяя остальные слои линейным вниманием. Однако эффективность преобразования «трансформер → гибрид» критически зависит от того, какие именно слои сохраняют полное внимание. Существующие методы выбора гибридных слоёв обычно опираются на эвристические стратегии, такие как фиксированные шаблоны размещения или поуровневые оценки, неявно трактуя важность слоя изолированно и игнорируя взаимозависимый эффект слоёв в рамках глобальной гибридной конфигурации. В данной работе мы формулируем задачу выбора гибридных слоёв как задачу оптимизации подмножества с ограничением на бюджет. Далее мы предлагаем 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.