[논문 리뷰] Explicit Multi-head Attention for Inter-head Interaction in Large Language Models

Runyu Peng, Yunhua Zhou외 5명|arXiv (Cornell University)|2026. 01. 27.

Multimodal Machine Learning Applications인용 수 0

한 줄 요약

본 논문은 Multi-head Explicit Attention(MEA)를 도입하여 Head-level Linear Composition(HLC)과 Group Normalization을 통해 교차-head 상호작용을 명시적으로 모델링하고, 사전학습 수렴을 개선하며 지식/과학 과제에서 50% KV-캐시 메모리 감소를 최소한의 성능 손실로 가능하게 한다. 또한 여러 주의 변형을 하나로 통일하고 저랭크 재구성을 통한 KV-캐시 압축을 지원한다.

ABSTRACT

In large language models built upon the Transformer architecture, recent studies have shown that inter-head interaction can enhance attention performance. Motivated by this, we propose Multi-head Explicit Attention (MEA), a simple yet effective attention variant that explicitly models cross-head interaction. MEA consists of two key components: a Head-level Linear Composition (HLC) module that separately applies learnable linear combinations to the key and value vectors across heads, thereby enabling rich inter-head communication; and a head-level Group Normalization layer that aligns the statistical properties of the recombined heads. MEA shows strong robustness in pretraining, which allows the use of larger learning rates that lead to faster convergence, ultimately resulting in lower validation loss and improved performance across a range of tasks. Furthermore, we explore the parameter efficiency of MEA by reducing the number of attention heads and leveraging HLC to reconstruct them using low-rank "virtual heads". This enables a practical key-value cache compression strategy that reduces KV-cache memory usage by 50% with negligible performance loss on knowledge-intensive and scientific reasoning tasks, and only a 3.59% accuracy drop for Olympiad-level mathematical benchmarks.

연구 동기 및 목표

• 트랜스포머의 어텐션 성능 향상을 위해 헤드 간 커뮤니케이션의 활용을 모티브로 삼는다.
• Head-level Linear Composition를 도입하여 명시적 교차-head 상호작용을 가능하게 한다.
• GroupNorm을 통해 학습 안정성을 높이고 MEA를 DFA 및 THA와 같은 기존 변형과 연결지어 해석한다.
• 스케일링 법칙에 의해 더 큰 학습률로 학습 속도를 더 빠르게 수렴시키는 방법을 제시한다.
• 저랭크 재구성을 통한 KV-캐시 압축으로 성능 저하 없이 메모리를 감소시킨다.]
• method: [

제안 방법

• 헤드 간 정보를 섞기 위한 Head-level Linear Composition(HLC)를 정의한다.
• K와 V를 HLC로 혼합된 버전으로 대체하고 헤드 출력에 대해 GroupNorm을 적용하여 MEA를 개발한다.
• DFA와 THA가 MEA의 특수 사례로 어떻게 관계하는지를 하나의 통합 시각으로 제시한다.
• 학습률 선택을 효율적으로 하기 위해 스케일링 법칙을 활용하고,From-scratch 사전학습 비교를 수행한다.
• 메모리 절감을 위해 저랭크 근사치를 통한 KV-캐시 압축을 제안한다.

실험 결과

연구 질문

• RQ1 MEA가 표준 트랜스포머 및 다른 헤드 간 변형에 비해 최적화 및 최종 성능을 개선하는가?
• RQ2 GroupNorm이 MEA의 학습 안정성 및 표현 다양성에 어떤 영향을 미치는가?
• RQ3 MEA가 지식/과학 과제에서 큰 규모의 KV-캐시를 메모리 효율적으로 가능하게 하는가?
• RQ4 DFA 및 Talking-Heads 변형이 MEA와 통일된 이론적 관점에서 어떻게 연결되는가?
• RQ5 계속된 사전학습 이후 KV-캐시 압축이 복잡한 추론 벤치마크에 미치는 영향은 무엇인가?

주요 결과

데이터셋	PIQA	OBQA	WinoGrande	HellaSwag	ARC-e	ARC-c	평균
Transformer	71.93	21.00	56.04	40.62	59.51	26.19	45.88
+GroupNorm	71.38	21.00	56.12	40.59	59.13	25.77	45.67
+DFA	71.76	22.20	54.38	41.29	60.69	27.82	46.36
Ours	73.18	19.80	54.14	42.02	61.57	27.65	46.39

• GroupNorm과 함께 MEA가 평가된 변형들 중에서 가장 우수한 평균 하류 성능을 달성한다.
• MEA가 사전학습 중에 기반 모델들보다 더 큰 안정적 학습률과 더 빠른 수렴을 가능하게 한다.
• 지식 집중적 및 과학 과제에서 성능 손실이 거의 없이 KV-캐시 메모리 사용량을 50% 감소시킬 수 있으며, 완전 압축 하에서 올림피아드 수준의 수학 벤치마크에서 약 3.59%의 손실이 발생한다.
• DFA와 THA가 MEA 프레임워크 내에서 해석될 수 있으며, GroupNorm이 없는 DFA는 표준 주의로 수렴하는 경향이 있다.
• GroupNorm은 헤드 간 상호작용을 유지하고 최적화를 안정화시키며, 정규화가 결여된 변형들보다 MEA가 더 나은 성능을 보이게 한다.