어텐션 층이 하나 생기면 모델은 처음으로 앞 문맥을 볼 수 있다. 그 결과 “keep이 앞에 있고 지금 단어가 in이면 mind를 예측하라” 같은 규칙을 배운다. 이처럼 떨어져 있는 세 토큰을 잇는 규칙이 스킵 트라이그램(skip-trigram)이다.
이번 편의 수식은 이 규칙이 두 번의 표 조회로 만들어진다는 사실을 보이는 도구다. 먼저 어디를 볼지 찾고, 그곳에서 어떤 단어의 점수를 올릴지 찾는다. 이 두 단계만 놓치지 않으면 경로 전개와 고유값 부분을 건너뛰어 읽어도 전체 줄거리를 따라갈 수 있다.
그림 1. 1-레이어 어텐션 전용 트랜스포머. 임베딩된 토큰에 각 헤드의 출력이 더해지고, 곧바로 언임베딩되어 로짓이 된다.
경로 전개
모델의 답은 두 갈래 길에서 온다. 첫째는 지난 편에서 본 내장 자동완성이다. 둘째는 앞 문맥에서 정보를 가져오는 어텐션 헤드들이다. 각 길의 기여를 따로 적는 것이 경로 전개(path expansion)다.
수식으로는 다음과 같다. 처음 읽을 때는 “직접 경로 + 여러 헤드 경로”라는 덧셈 구조만 보면 충분하다.
$$T = \underbrace{W_U W_E}_{\text{직접 경로}} + \sum_h \underbrace{A^h \otimes (W_U W_{OV}^h W_E)}_{\text{헤드 } h \text{ 경로}}$$수식의 $\otimes$는 두 차원의 결합을 나타내는 표기다. 왼쪽(패턴)은 어느 위치에서 가져올지를, 오른쪽(행렬)은 가져온 내용이 어떻게 변할지를 맡는다. 2편의 QK/OV 분해가 수식에 그대로 새겨진 셈이다.
그림 2. 층마다 하나씩 곱해지던 항들을 펼치면 직접 경로와 헤드 경로들의 합이 된다. 어텐션 패턴은 토큰 쌍을 $W_{QK}$의 양쪽에 곱한 뒤, 자기회귀 마스킹이 적용된 소프트맥스를 취해 계산된다.
각 헤드 경로의 정체는 2편에서 본 두 장의 표다. 어디를 볼지 정하는 QK 표, 주목받은 토큰이 출력에 무엇을 올리는지 정하는 OV 표. 둘 다 토큰에서 토큰으로 가는 표라서, 모델을 돌려보지 않고 가중치만 읽어도 해석할 수 있다.
그림 3. 소스 토큰에서 로짓까지 이어지는 OV 회로(금색, $W_U W_O W_V W_E$)와, 소스·목적지 토큰 쌍의 어텐션 점수를 만드는 QK 회로(자주색, $W_E^\top W_Q^\top W_K W_E$). 각각 하나의 행렬로 요약된다.
스킵 트라이그램
표를 실제로 읽으면 어떤 규칙이 나올까. 구체 예로 보자. 문맥 앞쪽에 “keep"이 있었고, 지금 토큰이 “in"이다.
1단계, QK 표를 찾아본다. “in"의 쿼리는 앞쪽의 “keep"에 강하게 끌리도록 학습되어 있다.
(keep, in) 큰 값 ← "in"일 때는 "keep"을 봐라
2단계, OV 표를 찾아본다. “keep"이 주목받으면 “mind"의 점수가 오르도록 학습되어 있다.
(keep → mind) 큰 값 ← "keep"을 봤으면 "mind"를 올려라
두 표가 이어지며 규칙 하나가 완성된다. “keep … in → mind”. 앞쪽 어딘가의 [소스], 지금의 [목적지], 밀어 올릴 [출력]. 소스와 목적지 사이가 떨어져 있어도 되는 세 토큰 패턴이라서, 논문은 이를 스킵 트라이그램(skip-trigram)이라 부른다. 학습된 헤드에서 이런 규칙이 대량으로 발견되고, 이것이 1-레이어 모델의 능력 전부다. 자동완성 표 위에 스킵 트라이그램 표를 얹은 것.
그림 4. 전개된 QK/OV 행렬에서 값이 큰 항목들. 소스 토큰마다 주목하는 목적지 토큰과 확률이 오르는 출력 토큰을 읽으면 그대로 스킵 트라이그램 표가 된다. 넷째 줄은 LaTeX를, 다섯째 줄은 HTML 이스케이프를 배운 흔적이다.
모든 헤드가 내용을 보고 주목하는 것은 아니다. 주로 위치를 근거로, 이를테면 바로 앞 토큰 근처에 주목하는 헤드들도 발견된다. 지금은 조연처럼 보이지만, 다음 편에서 이전 토큰 헤드(previous token head)라는 이름으로 인덕션 헤드(induction head)의 파트너가 된다.
복사 헤드와 고유값
발견된 스킵 트라이그램의 가장 흔한 형태는 복사다. “Tom … T → Tom"처럼, 나왔던 토큰을 또 밀어주는 규칙이다. “문서에 이미 나온 단어는 또 나오기 쉽다"는 통계의 구현이고, 논문은 이것을 원시적인 인컨텍스트 러닝(in-context learning)이라 부른다. 역할 분담은 이제 익숙하다. OV 표는 주목받은 토큰 자신의 점수를 올리도록 세팅되고, QK 표는 그 토큰이 다음에 올 법한 자리에서만 뒤를 돌아본다. 복사하되, 자동완성이 허락하는 자리에만 복사하는 것이다.
재미있는 특수 사례도 있다. 토크나이저는 보통 공백을 단어 앞에 붙여 자르기 때문에, 드문 단어가 공백 없이 등장하면 두 토큰으로 쪼개진다(” Ralph"는 한 토큰, “Ralph"는 “R”+“alph”). 일부 헤드는 이 경우를 부분적으로 전담해서, 조각 토큰 “R"을 보면 앞쪽의 " Ralph"를 찾아 “alph"를 예측한다. 다음 편에서 볼 인덕션 헤드를 1-레이어가 흉내 내는 아주 특수한 사례인 셈이다.
그림 5. 공백 없이 쪼개진 단어의 복사를 다루는 항목들. 완전한 단어(” Ralph”)를 소스로, 조각 토큰(“R”)을 목적지로 삼아 나머지 조각(“alph”)을 출력한다.
복사 헤드인지는 어떻게 한눈에 알아볼까. 직관은 표의 대각선이다. “Tom을 보면 Tom을 올린다"는 규칙은 OV 표의 대각선 칸이 크다는 뜻이다. 다만 5만 x 5만 표의 대각선을 일일이 보는 대신, 논문은 고유값이라는 요약 통계 하나로 이것을 잰다. 고유값은 행렬이 특정 방향의 벡터를 늘리는 배율이다. $Mv = \lambda v$에서 $\lambda$가 양수면 그 방향 $v$가 뒤집히지 않고 커져서 나온다. 그러니 OV 표의 고유값이 양수 쪽으로 쏠려 있다는 것은, 토큰(들의 조합)을 넣으면 바로 그 토큰들의 점수가 커진다는 뜻이다. 표가 전반적으로 복사라는 얘기다. 이 잣대로 재면 분석한 12-헤드 모델에서 10개 헤드가 뚜렷한 복사 성향을 보였다. 겉보기에 거대한 이 행렬들이 실은 랭크 $d_{head}$(64나 128)짜리 저랭크 행렬이라 계산도 가능하다. 다만 단서가 붙는다. 복사 행렬은 반드시 양의 고유값을 갖지만 그 역은 성립하지 않아서, 양의 고유값은 “평균적으로 복사한다"는 강한 증거일 뿐 결정적 증명은 아니다. 이 요약 통계는 이후 논문들에서도 계속 쓰인다.
그림 6. 고유값 중 양수 비율로 헤드들을 세운 히스토그램. 12개 중 10개가 오른쪽 끝, 즉 복사 쪽 극단에 몰려 있다.
분해 구조가 낳는 버그
같은 헤드가 “keep … in → mind"와 “keep … at → bay"를 둘 다 배웠다고 하자. 헤드의 규칙은 언제나 QK 표와 OV 표, 두 장의 곱이다. QK 표에는 (keep, in)과 (keep, at) 칸이 크고, OV 표에는 (keep → mind)와 (keep → bay) 칸이 크다. 그런데 곱 구조에는 “in일 때는 mind만, at일 때는 bay만"이라고 조합을 묶어 줄 칸이 없다. 그래서 교차 조합인 “keep … in → bay"에도 확률이 샌다. 세 변수 함수 $f(a, b, c)$를 $f_1(a, b) \cdot f_2(a, c)$ 꼴로만 쓸 수 있는 셈이라, 세 토큰이 함께 만드는 상호작용을 자유롭게 담지 못하는 것이다. 모델의 오류가 무작위 잡음이 아니라 구조의 논리적 귀결이라는 것. 역설계가 주는 예측력의 좋은 예다.
그림 7. 제한된 표현력이 밖에서 보면 이상한 버그를 만든다. 올바른 스킵 트라이그램들을 배운 헤드는 소스와 출력이 교차된 조합(강조 표시)의 확률도 함께 올릴 수밖에 없다.
다 이해한 것인가
논문은 이 질문에 신중하게 답한다. 어떤 의미에서는 그렇다. 경로 전개(path expansion)로 모든 파라미터가 맥락 속에서 의미를 갖게 되었고, 알고리즘적 미스터리는 남지 않았다. 하지만 이 이해는 거대한 선형회귀의 가중치를 이해한다고, 또는 커다란 데이터베이스에 질의할 줄 안다고 말할 때의 이해다. 어휘가 5만 개면 전개된 OV 행렬 하나에 약 25억 개의 항목이 있다. 요약하는 방법을 더 만들지 않는 한, 사람이 머리에 담을 수 없는 양이다. 논문의 표현을 빌리면, 1-레이어 모델이 압축된 중국어 방이었음을 밝혀냈더니 이제 거대한 카드 더미가 남은 상황이다.
기술적 유보도 몇 가지 붙는다. 상관된 변수 위의 선형 모델이 늘 그렇듯 두 헤드가 서로의 역할을 대신할 수 있어 가중치 0이 “안 한다"는 뜻이 아닐 수 있고, 서로 다른 쿼리의 QK 점수를 비교할 표준적 정규화 방법도 없다. 계약서를 읽을 수 있게 만드는 데는 성공했지만, 아직 다 읽지는 못한 상태라는 것이 논문의 자기 평가다.
이번 편의 요점
- 1-레이어 모델은 직접 경로(자동완성)와 헤드 경로(스킵 트라이그램)의 합으로 완전히 전개된다.
- 스킵 트라이그램은 QK 표의 칸 하나와 OV 표의 칸 하나가 이어진 규칙이고, 대표 격은 복사다. 복사는 OV 표의 양의 고유값으로 요약된다(강한 증거이되 결정적 증명은 아님).
- 규칙이 두 표의 곱이라는 구조 자체가, 교차 조합이 새는 버그를 예측하게 해준다.
- 알고리즘적 미스터리는 사라졌지만, 25억 항목을 요약 없이 머리에 담을 수는 없다. 남은 것은 요약의 문제다.
다음 편이 이 논문의 절정이다. 층이 둘이 되면 헤드가 헤드와 합성되고, 인덕션 헤드가 등장한다.
원문: A Mathematical Framework for Transformer Circuits의 One-Layer Attention-Only Transformers 섹션. 본문의 도식은 모두 원문에서 가져왔다.