Morphing into Hybrid Attention Models

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.

한 줄 요약

FlashMorph는 Transformer-to-hybrid 변환에서 효과적이고 효율적인 레이어 선택을 위한 글로벌 최적화 기반 방법이다.

핵심 기여도

Hybrid layer selection을 예산 제약 하의 subset optimization 문제로 정식화.
FlashMorph는 레이어별 게이트를 합성 retrieval 데이터로 공동 최적화하며, linearization regularization을 도입.
Qwen3-0.6B 및 Qwen3-1.7B 모델에서 기존 방법 대비 6:1 hybrid ratio에서 RULER 성능 +2.8~3.1% 개선.
레이어 선택 비용을 현저히 줄이며, 3:1 hybrid ratio에서 57.2 → 59.0 (GLA), 61.6 → 64.7 (GDN) 성능 향상.

핵심 아이디어

기존의 hybrid attention 모델은 full-attention 레이어를 고정된 패턴이나 레이어별 점수로 선택하는 방식을 사용했으나, 이는 레이어 간 상호작용을 무시하는 단점이 있었다. FlashMorph는 이 문제를 해결하기 위해, 레이어 선택을 **예산 제약 하의 subset optimization 문제**로 정식화한다. 이는 레이어 간 **상호보완성**과 **중복성**을 고려한 **글로벌 최적화**를 가능하게 한다.

모델 내 각 full-attention 레이어에 linear-attention branch를 추가한 morphable model을 구성하고, 학습된 가중치를 고정한 상태에서 **레이어별 게이트를 합성 retrieval 데이터로 공동 최적화**한다. 이 과정에서 **linearization regularization**을 도입하여 모델이 linear attention에 의존하도록 유도한다. 최적화된 게이트는 사전 정의된 full-attention 예산에 따라 이산화되어 hybrid 아키텍처를 생성한다. 이후 **logits distillation**과 **long-context fine-tuning**을 수행하여 최종 모델을 완성한다.

기술적 접근법

**Morphable model**: 각 full-attention 레이어에 linear-attention branch를 추가.
**Frozen weights**: 학습된 가중치는 고정하고, 레이어별 게이트만 최적화.
**Synthetic retrieval data**: passkey answer-token alignment 기반의 합성 데이터 사용.
**Linearization regularization**: linear attention에 의존하도록 유도.
**Discretization**: 사전 정의된 full-attention 예산에 따라 게이트 이산화.
**Distillation & fine-tuning**: logits distillation과 long-context fine-tuning을 통해 최종 모델 학습.

주요 결과

Qwen3-1.7B 기반 모델에서 6:1 hybrid ratio에서 FlashMorph는 GLA와 GDN backbone 모두에서 기존 방법 대비 RULER 성능 +2.8~3.1% 개선.
3:1 hybrid ratio에서 GLA backbone에서 RULER 성능 57.2 → 59.0 (+1.8%), GDN backbone에서 61.6 → 64.7 (+3.1%) 향상.
합성 retrieval 데이터를 사용한 FlashMorph (FlashMorph w/ syn)는 GLA에서 59.0, GDN에서 64.7의 RULER 성능 달성.
기존 PostNAS, KL-LS, HALO 대비 레이어 선택 비용을 현저히 줄임.

의의 및 한계

FlashMorph는 Transformer-to-hybrid 변환에서 레이어 선택의 효율성과 효과성을 동시에 달성하는 **글로벌 최적화 기반 접근법**으로, 기존 휴리스틱 방법의 한계를 극복한다. 특히, 레이어 간 상호작용을 고려한 **joint optimization**을 통해 더 나은 hybrid 구성이 가능하며, 합성 retrieval 데이터를 사용한 **supervision 방식**이 long-context 성능 향상에 기여한다.

그러나, FlashMorph는 합성 데이터에 의존하며, 실제 데이터와의 차이가 있을 수 있다. 또한, 특정 backbone (예: GLA, GDN)에 최적화된 성능을 보이지만, 다른 linear attention 메커니즘과의 호환성은 추가 실험 필요. 레이어 선택 후의 fine-tuning 과정도 성능에 영향을 줄 수 있으므로, 전체 파이프라인 최적화가 중요하다.

실용적 활용

FlashMorph는 대규모 언어 모델의 long-context 처리 효율성을 향상시키는 데 활용 가능하다. 특히, **자연어 처리**, **문서 검색**, **대화 시스템** 등에서 긴 입력을 처리해야 하는 상황에 적합하다. 또한, 기존 Transformer 모델을 low-cost로 hybrid 모델로 변환할 수 있어, **모델 배포 및 서비스 최적화**에 유용하다.