Attention Mechanism

어텐션 메커니즘은 각 출력을 입력 특성의 가중 평균으로 구성하는 신경망 연산입니다. 가중치는 처리 중인 요소와 각 입력이 얼마나 관련 있는지를 바탕으로 계산됩니다. 쿼리 벡터와 키 벡터를 비교해 관련성을 점수화하고, 그 점수로 해당 값 벡터에 가중치를 부여합니다. 어텐션은 2014년에 도입되어 기계 번역에 사용되었습니다. 이를 통해 순환 신경망(RNN) 디코더는 문장 전체를 하나의 고정 길이 벡터에 압축하는 대신 모든 원본 단어를 되돌아볼 수 있었습니다. 2017년 트랜스포머는 어텐션을 기반으로 전체 아키텍처를 구성했습니다. 어텐션은 현재 대규모 언어 모델, 비전 트랜스포머, 다수의 객체 감지 모델에서 핵심 연산으로 사용됩니다.

어텐션 작동 방식#

어텐션은 입력에 대해 학습된 세 개의 선형 투영을 사용합니다. 쿼리(Q)는 현재 토큰이 무엇을 찾는지 나타내고, 각 키(K)는 토큰이 무엇을 제공하는지 설명하며, 각 값(V)은 전달할 콘텐츠를 담습니다. 모델은 쿼리를 모든 키와 비교하고 비교 점수를 가중치로 변환한 뒤, 값의 가중합을 반환합니다. 번역 모델이 "The black cat"을 "noir"로 번역할 때, 해당 단계의 쿼리는 "black"의 키에 가장 높은 점수를 부여합니다. 프랑스어에서는 형용사가 명사 뒤에 오는데도 그렇습니다.

멀티헤드 어텐션은 각기 다른 저차원 투영을 사용해 이러한 연산을 여러 개 병렬로 실행한 다음 결과를 연결합니다. 그러면 여러 헤드가 같은 레이어 안에서 서로 다른 관계에 특화될 수 있습니다.

스케일드 닷프로덕트 어텐션#

트랜스포머의 어텐션은 Attention(Q, K, V) = softmax(QKᵀ / √d_k) V으로 표기합니다. 각 쿼리는 모든 키와의 내적을 계산해 비교하며, 점수가 벡터 크기에 따라 커져 소프트맥스를 포화시키지 않도록 키 차원 d_k의 제곱근으로 나눕니다. 소프트맥스는 각 쿼리의 점수를 합이 1인 가중치로 변환하며, 출력은 값 벡터의 가중합입니다.

import torch
import torch.nn.functional as F

torch.manual_seed(0)

# 각각 8차원인 4개 토큰 시퀀스의 쿼리, 키, 값입니다.
q, k, v = torch.randn(3, 4, 8).unbind(0)

# 모든 쿼리를 모든 키와 비교해 점수를 계산하고, sqrt(d_k)로 스케일링한 뒤 각 행을 소프트맥스로 정규화합니다.
weights = (q @ k.T / k.shape[-1] ** 0.5).softmax(dim=-1)

# 각 출력 토큰은 값 벡터의 가중 평균입니다.
output = weights @ v

print("Attention weights of token 0:", [round(w, 3) for w in weights[0].tolist()])
print("Matches PyTorch SDPA:", torch.allclose(output, F.scaled_dot_product_attention(q, k, v), atol=1e-6))

[0.121, 0.741, 0.106, 0.032]과 True을 출력합니다. 토큰 0은 토큰 1의 값 벡터에서 출력의 74%를 가져오며, 네 가중치의 합은 1입니다. 수동 계산 결과는 지원되는 GPU에서 FlashAttention과 같은 메모리 효율적인 커널을 호출하는 PyTorch의 융합 scaled_dot_product_attention 결과와 일치합니다.

어텐션의 유형#

• 셀프 어텐션: 쿼리, 키, 값이 모두 같은 시퀀스에서 나오므로 각 토큰은 다른 모든 토큰으로부터 문맥을 모읍니다. 트랜스포머 인코더와 ViT 내부에서 사용되는 어텐션입니다.
• 크로스 어텐션: 쿼리는 한 시퀀스에서, 키와 값은 다른 시퀀스에서 가져옵니다. 번역 디코더는 원본 문장에 어텐션을 적용하고, 텍스트-이미지 디퓨전 모델은 프롬프트에 어텐션을 적용하며, DETR 디코더의 객체 쿼리는 이미지 특성에 어텐션을 적용합니다.
• 인과적 어텐션: 마스크가 각 토큰이 뒤에 오는 토큰에 어텐션을 적용하지 못하도록 차단하므로 언어 모델은 왼쪽에서 오른쪽으로 텍스트를 생성할 수 있으며, 이전 토큰의 키와 값은 KV 캐시에 저장할 수 있습니다.
• 효율적인 어텐션: 희소 어텐션, 슬라이딩 윈도우 어텐션, 선형 어텐션은 모든 토큰 쌍을 비교할 때 발생하는 이차 비용을 줄입니다. 변형 가능 어텐션은 각 기준 위치 주변에서 학습된 몇 개의 지점만 샘플링합니다.
• CNN의 채널 어텐션: Squeeze-and-Excitation 네트워크는 특성 맵의 전역 평균을 사용해 채널의 가중치를 재조정합니다. 이는 계산 비용이 적은 어텐션 방식으로, ILSVRC 2017 분류 대회 우승 모델의 기반이 되었습니다.

어텐션과 컨볼루션#

합성곱은 모든 위치에 같은 작은 필터를 적용하므로 각 출력은 고정된 로컬 윈도우만 참조하고 가중치는 이미지에 따라 달라지지 않습니다. 어텐션은 콘텐츠로부터 가중치를 계산하며 한 레이어 안에서 임의의 두 위치를 연결할 수 있습니다. 대신 토큰 수의 제곱에 비례하는 비용이 발생합니다. 스트라이드 32의 640×640 이미지는 토큰 400개로 구성된 20×20 그리드를 만들어 쿼리-키 쌍 160,000개를 생성합니다. 같은 이미지에서 스트라이드 8 그리드는 토큰 6,400개로 구성되며 쌍은 약 4,100만 개에 이릅니다. YOLO26은 이에 맞춰 어텐션 블록을 배치합니다. 백본의 마지막에 있는 C2PSA 블록과 넥에 있는 어텐션 적용 C3k2 블록 하나는 모두 20×20 스트라이드 32 특성 맵에서 멀티헤드 셀프 어텐션을 실행하며, 해상도가 더 높은 모든 레이어는 합성곱을 유지합니다.