[Paper] MHLA: Token-Level Multi-Head를 통한 Linear Attention의 표현력 복원

Source: arXiv - 2601.07832v1

개요

논문 **“MHLA: Restoring Expressivity of Linear Attention via Token‑Level Multi‑Head”**는 트랜스포머의 오래된 병목 현상인 소프트맥스 자체‑주의의 2차 비용을 해결합니다. 선형‑주의 변형은 O(N) 시간과 메모리를 약속하지만, 일반적으로 전역 컨텍스트가 평범하고 저‑랭크 표현으로 붕괴되어 정확도가 희생됩니다. MHLA (Multi‑Head Linear Attention)는 전체 주의의 표현력을 않고 선형‑시간 보장을 깨뜨리지 않으면서 다시 도입하여 비전, 언어 및 생성 작업에서 상당한 향상을 제공합니다.

주요 기여

• Token‑level multi‑head design: 토큰 시퀀스를 토큰 차원을 따라 여러 헤드로 분할하여(일반적인 특성 차원 분할 대신) 다양한 컨텍스트 신호를 보존합니다.
• Theoretical guarantee: MHLA가 선형 시간·공간 복잡성을 유지하면서 소프트맥스 어텐션의 표현 능력을 근사함을 증명합니다.
• Empirical validation across domains:
  • +3.6 % ImageNet 분류에서 top‑1 정확도 향상.
  • +6.3 % 벤치마크 NLP 작업(e.g., GLUE)에서 개선.
  • +12.6 % 이미지 생성 품질 향상(FID 감소).
  • +41 % 동일 실행 시간 하에서 비디오 생성 충실도 향상.
• Lightweight implementation: 추가적인 컨볼루션이나 순환 모듈이 없으며, 기존 Transformer 코드베이스에 한 줄만 수정하면 적용할 수 있습니다.

방법론

1. Linear attention 요약 – 표준 선형 어텐션은 softmax 커널을 특징 맵 ϕ(·) 로 재작성하여 어텐션을 일련의 행렬 곱셈으로 계산할 수 있게 하며, O(N) 복잡도를 제공합니다.
2. “global context collapse”(전역 컨텍스트 붕괴) 식별 – 모든 토큰이 동일한 ϕ‑임베딩을 공유하면 어텐션 출력이 각 위치마다 거의 동일해져 모델이 미세한 패턴을 구분하는 능력이 약화됩니다.
3. Token‑level multi‑head formulation –
   • 입력 토큰 시퀀스 X ∈ ℝ^{N×D} 를 H 개의 연속 토큰 그룹으로 나누며, 각 그룹의 크기는 ≈ N/H 입니다.
   • 각 헤드 h 에 대해 별도의 선형 어텐션 모듈이 해당 토큰 슬라이스를 사용해 자체 컨텍스트를 계산하고, 헤드‑특정 출력 Y_h 를 생성합니다.
   • 헤드들을 연결(또는 합산)하여 최종 표현을 형성합니다.
4. Complexity analysis – 각 헤드는 N/H 토큰을 처리하므로 전체 비용은 O(N·D) (선형) 로 유지됩니다. 이는 헤드별 연산이 독립적이며 합산되기 때문입니다.
5. Training details – 저자들은 기본 Transformer와 동일한 옵티마이저 설정을 유지하고, 어텐션 레이어만 교체합니다. 추가적인 정규화나 보조 손실은 필요하지 않습니다.

결과 및 발견

• 표현력 복원: 어텐션 맵 시각화 결과 MHLA가 토큰별로 뚜렷한 패턴을 유지하는 반면, vanilla linear attention은 맵이 균일해짐을 보여줍니다.
• 학습 안정성: 수렴 곡선이 softmax attention과 일치하여 토큰‑수준 분할이 최적화 난제를 일으키지 않음을 나타냅니다.
• 확장성: 최대 16 k 토큰 길이까지 실험한 결과, 실행 시간과 메모리는 선형적으로 유지되면서 정확도도 경쟁력을 유지합니다.

실용적인 시사점

• 대규모 배포 가능: 개발자들은 이제 엣지 디바이스, 장문 NLP 파이프라인, 혹은 고해상도 비디오 생성에서도 2차 메모리 한계에 부딪히지 않고 Transformer‑스타일 모델을 실행할 수 있습니다.
• 드롭인 교체: MHLA는 어텐션 레이어만 변경하기 때문에 기존 코드베이스(예: Hugging Face Transformers, PyTorch Lightning)를 최소한의 리팩터링으로 적용할 수 있습니다.
• 비용 효율적인 학습: 선형 복잡도는 GPU 메모리 부담을 줄여 더 큰 배치 크기나 더 긴 컨텍스트 윈도우를 가능하게 하며, 이는 빠른 반복 사이클과 낮은 클라우드 비용으로 이어집니다.
• 새로운 제품 기회: 실시간 비디오 합성, 대규모 추천 시스템, 그리고 온‑디바이스 언어 어시스턴트는 MHLA가 제공하는 속도‑정확도 트레이드오프의 혜택을 받을 수 있습니다.

Limitations & Future Work

• Head granularity trade‑off: 토큰‑헤드 수 H 를 선택하는 것은 하이퍼파라미터이며, 헤드가 너무 많으면 컨텍스트가 파편화되고, 너무 적으면 붕괴 현상이 발생합니다. 논문에서는 휴리스틱을 제공하지만 자동 튜닝 방법은 제시하지 않습니다.
• Benchmarks limited to vision and standard NLP: 결과는 인상적이지만, 초장 시퀀스(예: 100k‑토큰 문서)나 멀티모달 작업에 대한 평가가 아직 이루어지지 않았습니다.
• Theoretical bounds: 표현력 회복에 대한 증명은 특징 맵 ϕ의 특정 속성을 전제로 합니다. 다른 커널(예: 코사인 기반)로 분석을 확장하면 적용 범위를 넓힐 수 있습니다.
• Hardware‑specific optimizations: 현재 구현은 밀집 행렬 연산에 의존하고 있습니다. 향후 작업에서는 GPU/TPU에서 추가 속도를 끌어내기 위해 융합 커널이나 희소성 인식 커널을 탐색할 수 있습니다.

Bottom line: MHLA는 선형 확장성을 얻기 위해 소프트맥스 어텐션의 핵심 성능을 희생할 필요가 없음을 보여줍니다. 차세대 AI 시스템을 구축하는 엔지니어에게 더 크고, 더 빠르며, 메모리 효율이 높은 트랜스포머를 위한 실용적인 경로를 제공합니다.

저자

• Kewei Zhang
• Ye Huang
• Yufan Deng
• Jincheng Yu
• Junsong Chen
• Huan Ling
• Enze Xie
• Daquan Zhou

논문 정보

• arXiv ID: 2601.07832v1
• 카테고리: cs.CV, cs.AI
• 발행일: 2026년 1월 12일
• PDF: PDF 다운로드