MultiheadAttention — Attention Is All You Need의 구현체

MultiheadAttention — Attention Is All You Need의 구현체

멀티헤드 어텐션은 서로 다른 표현 부분공간(representation subspaces)에서 정보에 '함께 주목'할 수 있게 하는 층이에요. PyTorch의 torch.nn.MultiheadAttention(embed_dim, num_heads) 는 그 원 논문의 구조를 참조 구현으로 담고 있어요.

수식으로 보는 핵심

멀티헤드 어텐션은 다음과 같이 정의돼요.

MultiHead(Q,K,V) = Concat(head_1,...,head_h)W^O
head_i = Attention(QW^Q_i, KW^K_i, VW^V_i)

즉 입력을 여러 헤드로 나눠 각자 어텐션을 계산한 뒤 다시 결합하는 구조예요. 최신 아키텍처로 확장하기 전 기본 원리를 익히는 기준(reference) 구현으로 적합해요.

성능 최적화

nn.MultiheadAttention 은 가능하면 최적화된 scaled_dot_product_attention() 구현을 사용해요. 추론 시 need_weights=Falsetorch.inference_mode 조합이면 fastpath를 타서 훨씬 빠르고, Nested Tensor(배치마다 길이가 다른 입력)도 지원해요.

마스킹과 주의점

  • attn_mask: 특정 위치에 주목하지 않도록 막는 2D/3D 마스크.
  • key_padding_mask: 패딩을 무시하게 하는 마스크.
  • 주의: scaled_dot_product_attention 과 마스크 의미가 반대라 전환 시 반전이 필요해요.

더 알아보기