증류 후 대체: 효율적인 작업 특화 하이브리드 어텐션 모델 구축

Distill-then-Replace: Efficient Task-Specific Hybrid Attention Model Construction

트랜스포머 아키텍처는 완전 어텐션을 통해 최첨단 정확도를 제공하지만, 시퀀스 길이에 대한 2차 복잡성으로 인해 시간 및 메모리 제약으로 인해 실제 배포에 어려움이 있습니다. 선형 어텐션 메커니즘은 선형 또는 거의 선형 스케일링을 제공하지만, 종종 성능 저하를 초래합니다. 완전 어텐션과 선형 어텐션 레이어를 통합하는 하이브리드 모델은 효율성과 표현력 사이의 균형을 제공할 수 있지만, 다음과 같은 두 가지 주요 과제에 직면합니다. 첫째, 이러한 하이브리드 모델을 처음부터 훈련하는 것은 계산 비용이 많이 들고, 둘째, 어텐션 유형의 최적 배치를 수동으로 설계하는 것은 매우 어렵습니다. 우리는 이러한 문제를 해결하기 위해 먼저 사전 훈련된 완전 어텐션 모듈에서 선형 어텐션 모듈로 블록 단위 로컬 증류를 통해 가중치를 전송하고, 둘째, 목표 작업에서 검증 성능을 모니터링하면서 완전 어텐션 블록을 선형 어텐션 블록으로 반복적으로 대체하는 탐욕적인 레이어 대체 전략을 도입합니다. 이를 통해 비용이 많이 드는 재훈련이나 신경망 아키텍처 검색 없이 단일 효율적인 단계로 작업 특화 하이브리드 모델을 얻을 수 있으며, 다양한 다운스트림 작업에 적용 가능한 사전 훈련된 완전 어텐션 백본에 적용할 수 있습니다.

Transformer architectures deliver state-of-the-art accuracy via dense full-attention, but their quadratic time and memory complexity with respect to sequence length limits practical deployment. Linear attention mechanisms offer linear or near-linear scaling yet often incur performance degradation. Hybrid models that integrate full and linear attention layers promise a balance between efficiency and expressiveness, but face two major challenges: training such hybrid models from scratch is computationally expensive, and manually designing the optimal placement of attention types is highly nontrivial. We address both issues by first transferring weights from the pretrained full-attention modules to its linear attention counterparts through blockwise local distillation, and second, introducing a greedy layer replacement strategy that iteratively substitutes full attention blocks with linear ones while monitoring validation performance on the target task. This yields a task-specific hybrid model in a single efficient pass, without costly re-training or neural architecture search, and can be applied to any pretrained full-attention backbone for diverse downstream tasks.