MHLA: 토큰 수준 다중 헤드를 통한 선형 어텐션의 표현력 회복

초록

트랜스포머 아키텍처가 많은 분야를 지배하고 있지만, 그 자체 어텐션의 2차 복잡도는 대규모 응용 분야에서의 사용을 방해한다. 선형 어텐션은 효율적인 대안을 제공하지만, 이를 직접 적용하면 성능 저하가 흔히 발생하며, 기존 해결책들은 일반적으로 추가 모듈(예: 깊이별 분리 가능 합성곱)을 통해 계산 오버헤드를 다시 도입하여 본래 목적을 훼손한다. 본 연구에서는 이러한 방법들의 핵심 실패 원인, 즉 모델이 표현의 다양성을 상실하는 '전역 콘텍스트 붕괴'를 규명한다. 이를 해결하기 위해 우리는 토큰 차원을 따라 분할된 헤드 내에서 어텐션을 계산함으로써 이러한 다양성을 보존하는 다중 헤드 선형 어텐션(MHLA)을 제안한다. 우리는 MHLA가 선형 복잡도를 유지하면서 소프트맥스 어텐션의 표현력을 대부분 회복함을 증명하고, 여러 분야에서 그 효과를 입증한다. 동일한 시간 복잡도 하에서 ImageNet 분류에서 3.6%, 자연어 처리에서 6.3%, 이미지 생성에서 12.6%, 비디오 생성에서 41%의 성능 향상을 달성하였다.

English

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.

MHLA: 토큰 수준 다중 헤드를 통한 선형 어텐션의 표현력 회복

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

초록

Support