트랜스포머 멀티헤드 어텐션 텐서 모양: 쿼리·키·값과 512차원을 8개 헤드로 나누는 원리

핵심 메시지

• 어텐션은 입력과 출력의 텐서 모양이 같다. 4개 토큰이 각각 512차원으로 들어와 4×512 블록이 되고, 어텐션을 통과해도 다시 4×512로 나온다. 변화는 중간 계산의 '모양'에 숨어 있다.
• 각 토큰은 학습된 행렬을 곱해 쿼리(내가 찾는 것)·키(내가 광고하는 것)·값(넘겨줄 정보)의 세 관점으로 쪼개진다. 4×512에 512×64 행렬을 곱하면 4×64가 되고 512는 상쇄된다.
• 쿼리와 키의 내적으로 토큰 간 관련도를 재면 4×64 × 64×4가 4×4 점수 격자가 된다. 관련도를 재는 64차원 공간은 최종 점수에서 사라진다.
• 512 = 8 × 64. 모델 폭을 8개 헤드로 나눠 각 헤드가 64차원에서 독립적으로 전체 파이프라인을 돈다. 헤드끼리는 서로의 점수를 보지 못하고 마지막 WO 행렬에서만 합쳐진다.
• 8개 헤드는 반복문으로 도는 게 아니라 큰 행렬 하나가 한 번에 투영하고, reshape·transpose로 헤드 축을 배치 축 옆에 붙여 단일 배치 행렬곱으로 모두 처리한다. 병렬성은 사실상 공짜다.

쉽게 이해하기

영상은 어텐션의 작은 역설로 시작한다. 'The cat sat down' 같은 4개 토큰이 각각 512차원 벡터로 들어오면 입력은 4×512 블록이고, 어텐션을 통과한 결과도 토큰 수와 폭이 그대로인 4×512다. 모양이 전혀 바뀌지 않는다면 어텐션은 대체 무엇을 하는 것일까? 이야기는 그 사이 중간 텐서의 모양에 담겨 있다.

각 토큰은 먼저 학습된 행렬을 곱해 쿼리·키·값 세 가지 관점으로 쪼개진다. 여기서 모든 모양을 지배하는 규칙이 등장한다. a×b 행렬과 b×c 행렬을 곱하면 안쪽 차원 b가 상쇄되어 a×c가 된다. 그래서 4×512 토큰에 512×64 행렬을 곱하면 4×64가 되고, 64는 쿼리·키·값이 살아가는 공간의 폭이다.

다음은 비교 단계다. 한 토큰이 다른 토큰에 얼마나 주목할지는 그 쿼리와 상대의 키의 내적으로 잰다. 모양으로는 4×64 × 64×4이고, 안쪽 차원 64가 맞아떨어져 상쇄되면 4×4 점수 격자가 남는다. 즉 어텐션은 토큰 위에서 이뤄지지만 관련도는 64차원 공간에서 재며, 그 공간은 최종 점수에 드러나지 않는다.

왜 512를 통째로 쓰지 않고 64차원씩 8개로 나눌까? 512 = 8 × 64이기 때문이다. 모델의 폭을 8개 헤드로 파티션하면 각 헤드가 64차원의 자기 공간에서 독립적으로 쿼리·키·값과 점수, 가중 혼합을 계산한다. 어떤 헤드도 다른 헤드의 점수를 보지 않고, 마지막에 하나의 행렬 WO에서만 다시 꿰매어 섞인다.

구현은 반복문일 것 같지만 아니다. 큰 행렬 하나가 8개 헤드를 한 번에 투영하고, reshape가 512폭 출력을 8×64로 재해석하며, transpose가 헤드 축을 배치 축 옆으로 옮긴다. 행렬곱은 언제나 마지막 두 차원만 건드리므로, 헤드는 배치와 나란히 서서 단일 배치 행렬곱으로 모든 시퀀스의 모든 헤드가 동시에 계산된다. 그 512×512 행렬은 하나의 텐서지만 8개의 서로 다른 블록으로 이뤄져 있어, 한 헤드는 동사의 주어를 다른 헤드는 가까운 구두점을 추적할 수 있다.

주요 인사이트

• 어텐션이 '모양을 바꾸지 않는다'는 사실 자체가 핵심 직관이다. 실제 연산은 눈에 보이는 입출력이 아니라 그 사이 중간 텐서 모양의 변화에 담겨 있다.
• 헤드는 하나의 큰 행렬을 reshape로 8개처럼 재해석한 '뷰(view)'다. 물리적으로 8개의 별도 텐서가 아니지만, 그 안의 숫자 블록은 진짜로 서로 다르다.
• 멀티헤드의 존재 이유는 다양성이다. 블록이 모두 같으면 8개 헤드가 똑같은 것을 계산하는 무의미한 복사가 되므로, 서로 달라야 각 헤드가 다른 관계를 포착할 수 있다.
• 행렬곱이 항상 마지막 두 차원만 다룬다는 성질 덕분에, 헤드 축을 배치 축 옆에 두면 모든 시퀀스의 모든 헤드를 단 한 번의 배치 행렬곱으로 계산할 수 있다.

자주 묻는 질문

관련 AI 소식