헤드당 세 개의 행렬, 또는 하나의 큰 행렬

헤드별 관점

각 헤드는 고유한 쿼리 투영, 키 투영, 값 투영이 필요합니다. 헤드 h에 대해:

Q_h = X · W_Q^h    여기서 W_Q^h의 형태는 [d_model, head_dim]입니다.
K_h = X · W_K^h    여기서 W_K^h의 shape는 [d_model, head_dim]입니다
V_h = X · W_V^h    여기서 W_V^h의 shape는 [d_model, head_dim]입니다

X는 입력 shape [batch, seq_len, d_model]을 가집니다. 프로젝션 후 Q_h, K_h, V_h는 각각 shape [batch, seq_len, head_dim]을 가집니다.

융합된 뷰

각 헤드별 행렬은 메모리에서 서로 나란히 배치됩니다. shape가 [d_model, d_model]인 단일 융합 행렬 W_Q가 모든 헤드를 한 번에 생성합니다:

Q_fused = X · W_Q              # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused)  # [batch, n_head, seq_len, head_dim]

융합 행렬 곱은 12개의 BLAS 호출 대신 하나만 사용합니다. CUDA 텐서 코어는 이 크기의 행렬 곱에서 최고 처리량에 도달합니다. 헤드별 행렬 곱은 하드웨어를 충분히 활용하지 못합니다.

매개변수 수

세 개의 융합된 행렬 W_Q, W_K, W_V, 각각 d_model × d_model. 출력 프로젝션 W_O도 d_model × d_model. ANDREA-120M의 경우:

레이어당 어텐션 params = 4 × 768² = 2,359,296 ≈ 2.36M
12개 레이어 전체 params      = 12 × 2.36M ≈ 28.3M

ANDREA-120M의 총 매개변수 중 대략 4분의 1이 어텐션 프로젝션에 있습니다. 나머지 4분의 3은 MLP 서브레이어와 임베딩에 있습니다.

프로젝션 명명하기

열두 개의 벡터가 하나가 되다

각 헤드가 계산한 후

각 헤드는 [batch, seq_len, head_dim] 모양의 출력 텐서를 생성합니다. 12개의 헤드가 12개의 이러한 텐서를 생성합니다. 특징 차원沿着의 연결은 이를 다시 쌓아줍니다:

concat_output = concat(head_1, head_2, ..., head_12)
shape         = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768]    # for ANDREA-120M

Concat은 split을 역으로 수행합니다. 총 feature dimension은 d_model로 되돌아갑니다. 차원에서 정보 손실은 없으며; 차이는 각 chunk가 담고 있는 내용에 있습니다: head 1의 chunk는 head 1의 학습된 attention 패턴을 반영합니다.

출력 프로젝션 W_O

단순 Concatenation만으로는 head들이 고립된 상태로 남습니다: head 4의 출력이 head 7의 출력 옆에 있지만, 서로에 대해 인식하지 못합니다. 출력 프로젝션 W_O (모양 [d_model, d_model])가 이를 섞어줍니다:

attention_output = concat_output · W_O
모양            = [batch, seq_len, d_model]

W_O 이후, 모든 출력 차원은 12개의 헤드 모두의 학습된 선형 조합을 담고 있습니다. 단일 행렬 곱을 통해 헤드 간 정보가 자유롭게 흐릅니다.

왜 헤드가 전문화되는가

아키텍처에서 head 4가 동사 시제를 배우거나 head 9가 짝짓기된 구두점을 배우도록 강제하는 것은 없습니다. 전문화는 그래디언트 압력에서 나타납니다: 훈련 중에 중복으로 기여하는 헤드들은 독특하게 기여하는 헤드들보다 더 작은 그래디언트를 받습니다. 수천 단계에 걸쳐 각 헤드는 총 손실을 가장 효과적으로 줄이는 틈새로 자리 잡습니다.

경험적으로, 훈련된 트랜스포머는 다음과 같은 헤드들을 보여줍니다: 위치 패턴(헤드가 이전 토큰을 봄), 구문 패턴(헤드가 일치하는 닫는 괄호를 봄), 의미 패턴(헤드가 가장 최근 명명된 엔티티를 봄). 라벨 없이 이 전문화가 훈련됩니다. 훈련 신호만 W_O를 통해 전파되어 헤드들을 분류합니다.

왜 12개의 헤드인가, 하나의 더 넓은 헤드가 아닌가

CUDA가 헤드를 저장하는 방법

단일 텐서, 재구성됨

ANDREA의 훈련 엔진 microgpt_cuda.cu는 12개의 헤드를 위한 12개의 별도 버퍼를 할당하지 않습니다. 하나의 융합된 텐서를 할당하고 헤드 차원을 스트라이드 패턴으로 처리합니다:

// Q = X · W_Q (한 번의 행렬 곱셈, 헤드 간 융합됨) 후
// Q의 형태: [batch, seq_len, d_model]
// [batch, seq_len, n_head, head_dim]으로 재구성
// [batch, n_head, seq_len, head_dim]으로 전치
// 이제 각 헤드의 내부 두 차원이 메모리에서 연속적임

전치는 n_head를 seq_len 앞에 이동시킵니다. 왜냐하면 다음 연산(Q_h · K_h^T)이 각 헤드의 seq_len × head_dim 슬라이스가 메모리에서 연속적이어야 하기 때문입니다. CUDA 행렬 곱은 연속 텐서에서 더 빠르게 실행됩니다.

하나의 커널, 여러 헤드

단일 어텐션 CUDA 커널은 모든 헤드에서 병렬로 실행됩니다. 각 스레드 블록은 하나의 (배치, 헤드) 쌍을 처리합니다; 블록 내부의 스레드들은 seq_len × head_dim 타일에서 협력합니다. 커널은 여러 헤드를 처리한다는 사실을 전혀 알지 못합니다; 런치 그리드가 병렬성을 처리합니다.

하드웨어를 반영한 구성

ANDREA-120M의 n_head=12, head_dim=64 선택은 RTX 4090 텐서 코어와 맞습니다. 텐서 코어는 16의 배수인 matmul 타일을 선호합니다. head_dim=64 = 4 × 16은 타일 모양과 정확히 일치합니다. head_dim=32 (ANDREA-12M)도 일치하지만 타일을 충분히 활용하지 않습니다. head_dim=72는 일치하지 않아 폴백 커널을 강제합니다.

최종 그림

다섯 단계. 세 개의 행렬 곱셈이 입력(Q, K, V 프로젝션)에 접근. 하나의 행렬 곱셈이 연결된 헤드(W_O)에 접근. 하나의 어텐션 커널이 모든 헤드를 병렬로 처리. ANDREA-120M은 레이어당 한 번씩 다섯 단계를 실행하며, 12개의 레이어 깊이로, 모든 순전파에서 실행됩니다.

파이프라인 읽기