멀티 헤드 어텐션

Multi-head attention, MHA; 멀티 헤드 어텐션
쿼리·키·값을 여러 개의 낮은 차원 부분 공간으로 나눠 어텐션을 병렬로 따로 계산한 뒤 이어 붙이는 트랜스포머의 핵심 연산

멀티 헤드 어텐션(Multi-head attention, MHA)은 어텐션을 한 번에 크게 계산하지 않고, 모델 차원을 $h$ 개의 "헤드"로 쪼개 각 헤드가 자기 몫의 쿼리·키·값으로 어텐션을 따로 계산한 뒤 결과를 이어 붙여 다시 투영하는 구조다. 2017년 〈Attention Is All You Need〉에서 트랜스포머와 함께 제안되었고, 인코더의 셀프 어텐션, 디코더의 마스크드 셀프 어텐션, 인코더-디코더(교차) 어텐션이 모두 이 형태다.^[1]

헤드를 나누는 이유는 한 번의 어텐션이 가중 평균 하나만 만든다는 데 있다. 한 위치가 문법상 주어와 앞 문장의 지시 대상을 동시에 봐야 할 때, 단일 헤드는 둘을 섞은 평균밖에 못 낸다. 헤드를 여럿 두면 서로 다른 표현 부분 공간에서 서로 다른 위치를 동시에 볼 수 있다. 원 논문의 표현으로는 "단일 헤드에서는 평균이 이를 막는다".^[1] 이후 추론 비용을 줄이려고 키·값 헤드를 줄인 그룹 쿼리 어텐션(GQA, MQA)과 키·값을 압축한 멀티 헤드 잠재 어텐션(MLA)이 나왔는데, 둘 다 이 문서의 MHA를 기준으로 비교된다.

수식

스케일드 점곱 어텐션(scaled dot-product attention)은

$Attention (Q, K, V) = softmax (\frac{Q K^{⊤}}{\sqrt{d_{k}}}) V$

이다. 멀티 헤드 어텐션은 이를 헤드마다 다른 투영으로 $h$ 번 계산한다.

$MultiHead (Q, K, V) = Concat ({h e a d}_{1}, \dots, {h e a d}_{h}) W^{O}, {h e a d}_{i} = Attention (Q W_{i}^{Q}, K W_{i}^{K}, V W_{i}^{V})$

$W_{i}^{Q}, W_{i}^{K} \in ℝ^{d_{m o d e l} \times d_{k}}$ , $W_{i}^{V} \in ℝ^{d_{m o d e l} \times d_{v}}$ , $W^{O} \in ℝ^{h d_{v} \times d_{m o d e l}}$ 이다. 원 논문은 $d_{m o d e l} = 512$ , $h = 8$ , $d_{k} = d_{v} = d_{m o d e l} / h = 64$ 를 썼다. 헤드마다 차원을 $1 / h$ 로 줄였으므로 전체 계산량은 차원을 그대로 쓴 단일 헤드 어텐션과 비슷하다.^[1]

셀프 어텐션에서는 $Q = K = V = X$ (이전 층 출력)이고, 교차 어텐션에서는 쿼리만 디코더에서, 키·값은 인코더 출력에서 온다.

텐서 모양으로 본 계산

배치 $B$ , 길이 $T$ , $D = h \cdot d_{k}$ 일 때 한 층의 흐름은 다음과 같다. 실제 구현에서 "헤드"는 별도 행렬이 아니라 한 행렬의 출력 차원을 $h$ 조각으로 나눠 읽는 것일 뿐이다.

단계	연산	모양
입력	$X$	(B, T, D)
투영	$X W^{Q}, X W^{K}, X W^{V}$ (각 D×D, 보통 하나의 D×3D로 합침)	(B, T, D) ×3
헤드 분리	view + transpose	(B, h, T, d_k)
점수	$Q K^{⊤} / \sqrt{d_{k}}$ , 필요하면 인과 마스크	(B, h, T, T)
가중치	마지막 축 소프트맥스	(B, h, T, T)
가중합	가중치 × V	(B, h, T, d_k)
헤드 결합	transpose + reshape(Concat)	(B, T, D)
출력 투영	$W^{O}$	(B, T, D)

매개변수와 계산량

투영 행렬 네 개가 각각 $D \times D$ 이므로 한 층의 매개변수는 편향을 빼고 $4 D^{2}$ 개다. 헤드 수 $h$ 는 매개변수 수에 영향을 주지 않는다. 계산량은 투영에 약 $8 T D^{2}$ FLOPs(곱셈·덧셈을 2로 셈), 점수와 가중합에 약 $4 T^{2} D$ FLOPs가 든다. 시퀀스가 $D$ 보다 길어지면 뒤의 $T^{2}$ 항이 지배하며, 이 부분의 메모리 이동을 줄인 것이 플래시 어텐션이다.

구현

아래는 PyTorch로 쓴 최소 구현이다. nn.MultiheadAttention과 가중치를 맞춰 넣으면 출력이 부동소수점 오차(10^-7 수준) 안에서 일치한다(PyTorch 2.5에서 확인).

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.h, self.d_k = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model)   # W_Q, W_K, W_V 를 한 행렬로
        self.out = nn.Linear(d_model, d_model)       # W_O

    def forward(self, x, causal=False):
        B, T, D = x.shape
        q, k, v = self.qkv(x).split(D, dim=-1)
        # (B, T, D) -> (B, h, T, d_k): 헤드는 차원을 나눈 것일 뿐이다
        q, k, v = (t.view(B, T, self.h, self.d_k).transpose(1, 2) for t in (q, k, v))
        scores = q @ k.transpose(-2, -1) / self.d_k ** 0.5          # (B, h, T, T)
        if causal:
            mask = torch.ones(T, T, dtype=torch.bool, device=x.device).triu(1)
            scores = scores.masked_fill(mask, float("-inf"))
        y = scores.softmax(dim=-1) @ v                               # (B, h, T, d_k)
        y = y.transpose(1, 2).reshape(B, T, D)                       # Concat(head_1..head_h)
        return self.out(y)

실무에서는 점수 행렬을 직접 만들지 않고 torch.nn.functional.scaled_dot_product_attention(q, k, v, is_causal=True)를 부른다. PyTorch 2.0부터 있는 함수로, 입력 조건에 따라 플래시 어텐션이나 메모리 효율 커널을 골라 쓴다.^[2] nn.MultiheadAttention은 기본 입력 모양이 (T, B, D)이므로 (B, T, D)로 쓰려면 batch_first=True를 줘야 한다.

헤드는 무엇을 배우는가

설계 의도는 헤드마다 다른 관계를 보는 것이었고, 실제로 그런 헤드가 관찰된다. 다만 모든 헤드가 똑같이 쓸모 있지는 않다.

역할 분화: 클라크(Kevin Clark) 등은 BERT의 헤드들이 구분 토큰([SEP])에 몰리거나, 고정된 상대 위치(바로 앞·뒤 토큰)를 보거나, 문장 전체에 넓게 퍼지는 패턴을 보이며, 같은 층 헤드끼리는 비슷한 경향이 있다고 보고했다. 동사의 직접 목적어, 명사의 한정사, 상호 참조 대상을 높은 정확도로 가리키는 헤드도 찾았다.^[3]
가지치기 가능성: 보이타(Elena Voita) 등은 영어-러시아어 번역 모델 인코더 헤드 48개 중 38개를 잘라도 BLEU가 0.15만 떨어지며, 끝까지 남는 것은 언어학적으로 해석 가능한 특화 헤드라고 보고했다.^[4] 미셸(Paul Michel) 등도 학습이 끝난 뒤에는 상당수 헤드를 지워도 성능이 크게 변하지 않고, 어떤 층은 헤드 하나로 줄여도 된다는 것을 보였다. 동시에 여러 헤드가 학습 과정 자체에는 도움이 된다는 정황도 제시했다.^[5] 이것이 헤드 단위 가지치기의 근거가 된다.
헤드 수의 최적점: 원 논문의 절제 실험에서 계산량을 고정하고 헤드 수를 바꾸면 단일 헤드는 최적 설정보다 BLEU가 0.9 낮았지만, 헤드가 너무 많아도 품질이 떨어졌다.^[1] 보자나팔리(Srinadh Bhojanapalli) 등은 헤드 수를 늘리면 헤드당 차원 $d_{k} = D / h$ 가 줄어 점수 행렬의 계수(rank)가 제한되는 "저계수 병목"이 생긴다고 분석했다.^[6]
회로 단위 기능: 여러 층의 헤드가 조합되어 하나의 알고리즘을 이루기도 한다. 앞 토큰 정보를 옮기는 헤드와 그것을 이용해 패턴을 복사하는 헤드의 조합인 인덕션 헤드가 대표적이며, 기계적 해석 가능성 연구의 주요 대상이다.
어텐션 싱크: 대형 언어 모델의 여러 헤드는 의미와 상관없이 시퀀스 첫 토큰에 큰 어텐션을 준다. 쉬안(Guangxuan Xiao) 등은 이 현상을 어텐션 싱크(attention sink)라 부르고, 첫 토큰 몇 개의 KV만 남겨 두면 슬라이딩 윈도우 추론이 무너지지 않는다는 것을 보였다(StreamingLLM).^[7] 소프트맥스 가중치의 합이 항상 1이어서, 볼 것이 없는 헤드도 어딘가에 가중치를 쏟아야 하기 때문으로 해석된다.

추론 비용과 변형

자기회귀 생성에서는 이전 토큰의 키와 값을 KV 캐시에 보관한다. MHA는 토큰마다 층당 키와 값 $2 n_{h} d_{h}$ 개( $n_{h}$ 는 헤드 수, $d_{h}$ 는 헤드 차원)를 저장해야 하고, 생성 단계마다 이 캐시 전체를 메모리에서 읽어야 한다. 섀지어(Noam Shazeer)는 이 메모리 대역폭이 디코딩 속도의 병목임을 지적하고 키·값 헤드를 하나만 두는 멀티 쿼리 어텐션(MQA)을 제안했다.^[8] 이후 변형은 대부분 "키·값을 얼마나, 어떻게 줄이느냐"의 문제다.

방식	쿼리 헤드	키·값	토큰·층당 KV 캐시(원소 수)	품질 경향	자세히
MHA	$n_{h}$	헤드마다 따로	$2 n_{h} d_{h}$	기준	이 문서
MQA	$n_{h}$	1세트를 모든 헤드가 공유	$2 d_{h}$	떨어질 수 있음	그룹 쿼리 어텐션
GQA	$n_{h}$	$n_{g}$ 세트, 그룹별 공유	$2 n_{g} d_{h}$	MHA에 근접	그룹 쿼리 어텐션
MLA	$n_{h}$	저차원 잠재 벡터 하나로 압축	$d_{c} + d_{h}^{R}$	딥시크는 MHA 이상이라고 보고	멀티 헤드 잠재 어텐션

헤드 구조 자체를 고친 변형도 있다. 큐원(Qwen) 팀은 2025년 스케일드 점곱 어텐션 출력에 헤드별 시그모이드 게이트를 곱하는 단순한 수정이 30여 가지 변형 비교에서 일관되게 성능과 학습 안정성을 높이고, 어텐션 싱크를 줄인다고 보고했다.^[9] 이 게이트 어텐션은 2025년 9월 공개된 Qwen3-Next에 선형 어텐션 층과 섞여 들어갔다.^[10]

같이 보기

각주

↑ ^1.0 ^1.1 ^1.2 ^1.3 Ashish Vaswani 외, 〈Attention Is All You Need〉, NeurIPS 2017. arXiv:1706.03762
↑ PyTorch 문서, torch.nn.functional.scaled_dot_product_attention
↑ Kevin Clark 외, 〈What Does BERT Look At? An Analysis of BERT's Attention〉, BlackboxNLP 2019. arXiv:1906.04341
↑ Elena Voita 외, 〈Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned〉, ACL 2019. arXiv:1905.09418
↑ Paul Michel, Omer Levy, Graham Neubig, 〈Are Sixteen Heads Really Better than One?〉, NeurIPS 2019. arXiv:1905.10650
↑ Srinadh Bhojanapalli 외, 〈Low-Rank Bottleneck in Multi-head Attention Models〉, 2020. arXiv:2002.07028
↑ Guangxuan Xiao 외, 〈Efficient Streaming Language Models with Attention Sinks〉, ICLR 2024. arXiv:2309.17453
↑ Noam Shazeer, 〈Fast Transformer Decoding: One Write-Head is All You Need〉, 2019. arXiv:1911.02150
↑ Zihan Qiu 외, 〈Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free〉, 2025. arXiv:2505.06708
↑ Qwen/Qwen3-Next-80B-A3B-Instruct 모델 카드, 허깅 페이스.

[vaswani-1] 1.0 ^1.1 ^1.2 ^1.3 Ashish Vaswani 외, 〈Attention Is All You Need〉, NeurIPS 2017. arXiv:1706.03762

[sdpa-2] PyTorch 문서, torch.nn.functional.scaled_dot_product_attention

[clark-3] Kevin Clark 외, 〈What Does BERT Look At? An Analysis of BERT's Attention〉, BlackboxNLP 2019. arXiv:1906.04341

[voita-4] Elena Voita 외, 〈Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned〉, ACL 2019. arXiv:1905.09418

[michel-5] Paul Michel, Omer Levy, Graham Neubig, 〈Are Sixteen Heads Really Better than One?〉, NeurIPS 2019. arXiv:1905.10650

[lowrank-6] Srinadh Bhojanapalli 외, 〈Low-Rank Bottleneck in Multi-head Attention Models〉, 2020. arXiv:2002.07028

[sink-7] Guangxuan Xiao 외, 〈Efficient Streaming Language Models with Attention Sinks〉, ICLR 2024. arXiv:2309.17453

[mqa-8] Noam Shazeer, 〈Fast Transformer Decoding: One Write-Head is All You Need〉, 2019. arXiv:1911.02150

[gated-9] Zihan Qiu 외, 〈Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free〉, 2025. arXiv:2505.06708

[qwen3next-10] Qwen/Qwen3-Next-80B-A3B-Instruct 모델 카드, 허깅 페이스.

[1]

[2]

[3]

[4]

[5]

[6]

[7]

[8]

[9]

[10]