MHLA: 토큰 레벨 멀티 헤드를 활용한 선형 어텐션의 표현력 복원

MHLA: Restoring Expressivity of Linear Attention via Token-Level Multi-Head

트랜스포머 아키텍처가 많은 분야에서 지배적이지만, 그 제곱 복잡도의 셀프 어텐션은 대규모 애플리케이션에서의 활용을 제한합니다. 선형 어텐션은 효율적인 대안을 제공하지만, 직접 적용하면 성능이 저하되는 경향이 있으며, 기존의 해결책들은 종종 추가적인 모듈(예: 깊이 분리 컨볼루션)을 통해 계산 오버헤드를 다시 도입하여 원래의 목적을 무효화합니다. 본 연구에서는 이러한 방법들의 핵심적인 문제점인 '글로벌 컨텍스트 붕괴', 즉 모델이 표현적 다양성을 잃는 현상을 밝혀냈습니다. 이를 해결하기 위해, 우리는 토큰 차원을 따라 분할된 헤드 내에서 어텐션을 계산하여 이러한 다양성을 유지하는 '멀티 헤드 선형 어텐션 (MHLA)'을 제안합니다. 우리는 MHLA가 선형 복잡도를 유지하면서 소프트맥스 어텐션의 상당한 표현력을 회복한다는 것을 증명했으며, 다양한 분야에서 그 효과를 검증했습니다. 결과적으로, ImageNet 분류에서 3.6%, NLP에서 6.3%, 이미지 생성에서 12.6%, 그리고 비디오 생성에서 41%의 성능 향상을 달성했으며, 이는 동일한 시간 복잡도 하에서 이루어진 결과입니다.

While the Transformer architecture dominates many fields, its quadratic self-attention complexity hinders its use in large-scale applications. Linear attention offers an efficient alternative, but its direct application often degrades performance, with existing fixes typically re-introducing computational overhead through extra modules (e.g., depthwise separable convolution) that defeat the original purpose. In this work, we identify a key failure mode in these methods: global context collapse, where the model loses representational diversity. To address this, we propose Multi-Head Linear Attention (MHLA), which preserves this diversity by computing attention within divided heads along the token dimension. We prove that MHLA maintains linear complexity while recovering much of the expressive power of softmax attention, and verify its effectiveness across multiple domains, achieving a 3.6\% improvement on ImageNet classification, a 6.3\% gain on NLP, a 12.6\% improvement on image generation, and a 41\% enhancement on video generation under the same time complexity.