이번 편의 핵심은 트랜스포머를 “공용 작업 공간과 여러 작업자”로 다시 보는 것이다. 잔차 스트림은 모든 층이 함께 쓰는 작업 공간이고, 어텐션 헤드는 그 공간에서 정보를 찾아 다른 위치로 옮기는 작업자다.
이 그림을 잡으면 뒤의 수식은 두 질문으로 정리된다. 헤드는 어디를 볼까. 그리고 거기서 무엇을 가져올까. 논문의 Transformer Overview 섹션은 구현용 구조도를 이 두 질문에 답하기 좋은 형태로 다시 그린다.
단순화
논문은 분석 대상을 어텐션 전용 트랜스포머로 좁힌다. MLP 층은 떼어내고, 레이어 정규화와 바이어스는 인접 가중치에 접어 넣을 수 있다고 보고 생략한다. 남는 것은 토큰 임베딩 $W_E$, 어텐션 층들, 언임베딩 $W_U$뿐이다. MLP를 빼는 것은 물론 큰 단순화다. 하지만 어텐션만이라도 완전히 이해하는 것이 목표라면 좋은 출발점이다. MLP의 어려움은 6편에서 다시 만난다.
그림 1. 토큰 임베딩에서 시작해 잔차 블록들을 지나 언임베딩으로 끝나는 구조. 어텐션 헤드와 MLP는 각자 계산 결과를 잔차 스트림에 더한다.
잔차 스트림: 아무것도 하지 않는 통로
다시 그린 그림의 중심에는 잔차 스트림(residual stream)이 있다. 각 층은 잔차 스트림에서 자기 입력을 읽고, 계산 결과를 스트림에 더한다. 스트림 자체는 비선형성이 없다. 읽기와 더하기만 있는 공유 메모리, 또는 모든 층이 함께 쓰는 통신 채널이다.
이 관점의 힘은 선형성에서 나온다. 어떤 층의 출력은 이후 모든 층의 입력에 그대로 더해져 있다. 그래서 멀리 떨어진 두 구성요소 사이의 상호작용을, 사이에 있는 층들을 건너뛰고 가중치 행렬의 곱 하나로 요약할 수 있다. 논문은 이것을 가상 가중치(virtual weights)라고 부른다. 예를 들어 임베딩이 곧바로 언임베딩으로 이어지는 직접 경로(direct path)는 $W_U W_E$라는 하나의 행렬이다.
그림 2. 왼쪽: 각 층은 선형 사상으로 잔차 스트림에서 읽고, 결과를 더해 쓴다. 오른쪽: 모든 연산이 선형이므로 가중치를 곱해 나가면 층 쌍을 암묵적으로 잇는 가상 가중치가 드러난다. 층들은 서로 다른 부분공간을 써서 특정 층에만 정보를 보낼 수도 있다.
선형성에는 해석할 때 주의할 점도 따른다. 잔차 스트림의 “17번째 숫자”처럼 좌표 하나에 고정된 의미를 붙이면 안 된다. 스트림과 연결된 행렬들을 함께 회전시켜도 모델 행동은 같기 때문이다. 논문은 이를 특권적 기저(privileged basis)가 없다고 표현한다.
노트: 특권적 기저란? 공간의 좌표축이 특별한 의미를 갖게 되는 것은, 활성 함수처럼 좌표축 단위로 적용되는 연산이 있을 때다. 예를 들어 MLP의 활성값은 뉴런마다 비선형이 걸리므로, “3번 뉴런"이라는 축이 실제 계산의 단위가 된다. 이런 공간에는 특권적 기저가 있다. 반면 잔차 스트림은 읽기와 쓰기가 전부 선형이라, 공간 전체를 회전시켜도 계산이 달라지지 않는다. 어느 축도 특별하지 않으니 “몇 번째 차원이 무엇을 뜻하나"라는 질문 자체가 성립하지 않는다. 의미를 찾으려면 축이 아니라 방향, 즉 여러 차원의 선형 결합을 봐야 한다.
비좁은 채널: 병목, 그리고 중첩의 예고
트랜스포머에는 모든 층이 함께 쓰는, 크기가 정해진 메모리(잔차 스트림)가 하나뿐이다. 그런데 거기 담아야 할 계산 결과는 그보다 훨씬 많다. 그래서 여러 정보가 같은 공간에 겹쳐 저장될 수밖에 없고, 이것이 트랜스포머 해석을 어렵게 만드는 가장 큰 이유 중 하나가 된다. 이 한 문단이 이 섹션의 전부다. 이제 하나씩 풀어 보자.
잔차 스트림의 차원은 유한하다. 히든 차원이 4096인 모델이라면 잔차 스트림은 길이 4096짜리 벡터 하나다. 숫자 4096개를 적을 수 있는 화이트보드 한 장이 정보가 오가는 유일한 통로인 셈이다. 층들은 이 보드를 통째로 쓰는 대신, 서로 다른 부분공간(칸 묶음)을 나눠 쓰며 통신한다.
그런데 적겠다는 쪽의 수를 세어 보면 상황이 심상치 않다. 어텐션 헤드가 수천 개, MLP 뉴런은 수십만 개다. 한 층의 MLP 뉴런만 해도 보통 스트림 차원의 4배다. 저마다 “내 결과도 적어야 한다"고 줄을 서는데, 보드 칸은 4096개뿐이다.
그래서 병목이 생긴다. 50층 모델의 한가운데, 25층쯤의 잔차 스트림을 보자. 앞쪽 25개 층이 만든 계산 결과가 전부 이 벡터 하나를 거쳐 간다. 앞쪽 뉴런 수를 합치면 40만 개쯤, 스트림 차원의 100배다. 도로로 치면 100차선이 4차선으로 합쳐졌다가 다시 100차선으로 퍼지는 지점이다. 논문은 이런 텐서를 병목 활성(bottleneck activations)이라 부르고, 해석이 유독 어려울 것이라 내다본다.
40만 개의 결과가 어떻게 4096개의 숫자에 들어가는가. 여기서 예고되는 개념이 중첩(superposition)이다. 스트림의 한 차원이 “강아지” 정보만 담는 것이 아니라, 강아지에 문법에 감정에 문맥까지 여러 의미를 동시에 담을 수 있다는 것. 하나의 음성 신호에 여러 주파수가 섞여 있는 것과 같다. 벡터 속 2.31이라는 숫자 하나에도 여러 의미가 겹쳐 있을 수 있다. 이 주제는 후속 논문 Toy Models of Superposition이 정면으로 다룬다.
공간이 부족하면 청소도 필요하다. 논문은 어떤 헤드가 “이 정보는 이제 필요 없다"고 판단해 잔차 스트림에서 지우는, 메모리 관리(memory management) 역할을 할 가능성도 언급한다. 쓰는 헤드와 읽는 헤드만이 아니라, 지우는 헤드도 있을 수 있다는 것이다.
그림 3. 고차원 잔차 스트림은 부분공간들로 나뉜다. 층들은 겹치는 부분공간으로 상호작용하고, 서로소인 부분공간에 쓰면 상호작용하지 않는다. 오른쪽처럼 부분공간을 읽어 음수로 되써서 정보를 지우는 층도 있다.
헤드는 독립적이고 가산적이다
한 어텐션 층에는 헤드가 여러 개 있다. 구현에서는 헤드 출력들을 이어 붙여 큰 행렬 하나를 곱하지만, 수학적으로는 각 헤드가 독립적으로 계산해 잔차 스트림에 더하는 것과 동등하다.
$$\text{AttnLayer}(x) = x + \sum_{h} h(x)$$층이라는 단위는 구현의 편의일 뿐이고, 분석의 자연스러운 단위는 헤드다.
QK 회로와 OV 회로
어텐션 헤드 하나는 사실 두 개의 독립적인 질문에 답하는 장치다. 어디를 볼 것인가. 거기서 무엇을 가져올 것인가. 보통 어텐션을 쿼리, 키, 밸류라는 계산 절차로 배우기 때문에 이 둘이 한 덩어리로 보이지만, 논문의 재해석은 서로 독립적인 두 기능이 묶여 있을 뿐이라는 것이다. 이 분해가 이후 모든 분석의 지렛대가 된다.
구체적인 예로 보자. “Tom likes apples. He …“까지 읽고 다음 토큰을 예측하는 참이다. 어텐션 헤드는 두 단계로 움직인다.
1단계, 어디를 볼까. 현재 토큰 “He"가 검색어(쿼리)를 내고, 앞선 토큰들은 저마다 색인(키)을 갖고 있다. 쿼리를 색인들과 대조해 점수를 매기고, 소프트맥스로 배분한 것이 어텐션 패턴이다.
Tom 0.93
likes 0.02
apples 0.05
이 단계가 정하는 것은 “Tom을 봐라"뿐이다. Tom이 사람인지, 남성인지, 주어인지는 여기 없다. 오직 어디를 볼지만 정한다.
2단계, 무엇을 가져올까. Tom을 보기로 했다. 그런데 Tom 자리의 잔차 스트림에는 정보가 여러 겹 쌓여 있다. 사람, 남성, 주어, 이름, 단수. 헤드는 이것을 전부 옮기지 않고 자기 몫만 골라낸다. 예컨대 “3인칭 단수"만 뽑아 변환해서 “He” 자리의 잔차 스트림에 써 넣는 식이다. 다음 예측이 eat이 아니라 eats가 되는 데 필요한, 바로 그 정보다.
비유하자면 1단계는 내비게이션이고 2단계는 택배 기사다. 어느 집으로 갈지 정하는 일과, 그 집에서 어떤 상자를 싣고 올지는 서로 다른 일이다.
이제 이 두 단계를 행렬로 보자. 계산에는 네 행렬 $W_Q, W_K, W_V, W_O$가 쓰이는데, 각자 따로는 의미를 갖지 않는다. $W_Q$와 $W_K$는 1단계에서 항상 붙어 다니고, $W_V$와 $W_O$는 2단계에서 항상 붙어 다닌다. 그래서 실체는 곱으로 묶인 두 행렬이다.
QK 회로(query-key circuit) $W_{QK} = W_Q^\top W_K$. 토큰 쌍마다 “이 쿼리가 이 키에 얼마나 끌리는가"라는 점수 하나를 내는 표다. 위의 1단계 전체가 이 표 하나로 요약된다.
OV 회로(output-value circuit) $W_{OV} = W_O W_V$. “이 토큰이 주목받으면 출력에 어떤 영향을 주는가"를 담은 표다. 위의 2단계 전체가 이 표 하나로 요약된다.
크기 감각도 챙겨 두자. 스트림이 4096차원이고 헤드가 128차원이라면, $W_{QK}$는 4096 x 4096짜리 표처럼 굴지만 실제 자유도는 랭크 128뿐이다. 거대해 보여도 실체는 작은 저랭크 행렬 둘이라는 것. 4편에서 이 표들을 실제로 읽을 때, 이 사실이 계산을 가능하게 해준다.
이 관점에서 보면 키, 쿼리, 밸류 벡터는 부산물이다. 두 행렬을 다르게 인수분해해서 전혀 다른 중간 벡터를 만들어도 모델은 동일하게 작동한다. 논문이 이 벡터들을 아예 언급하지 않고 트랜스포머를 서술하는 쪽이 유용할 때가 많다고 말하는 이유다.
분해의 쓸모는 조합에서 드러난다. “항상 바로 앞 토큰을 본다"는 같은 QK 회로를 가진 두 헤드라도, 하나의 OV가 품사 정보를 옮기고 다른 하나의 OV가 감정 정보를 옮긴다면 둘은 전혀 다른 기능의 헤드다. 거꾸로 OV가 같아도 QK가 이전 토큰, 주어, 같은 단어, 문장 시작으로 바뀌면 역시 완전히 다른 헤드가 된다. 헤드의 기능은 두 회로의 조합이고, 그래서 헤드를 설명하는 데는 두 문장이 필요하다. “이 헤드는 이전 토큰을 본다"는 QK에 대한 설명이고, “이 헤드는 문법 정보를 전달한다"는 OV에 대한 설명이다.
분리는 분석의 지렛대이기도 하다. 어디를 볼지 정해지고 나면 옮겨지는 효과는 소스 토큰만의 함수라서, 두 표를 각각 따로 읽을 수 있다. 4편에서 스킵 트라이그램을 표로 읽어내고 그 버그까지 예측하는 것이 전부 이 분리 덕분이다. 그리고 이 분해에 익숙해지면 5편의 인덕션 헤드도 쉬워진다. QK가 “지난번의 내가 있던 자리"를 찾아가고, OV가 그 옆 토큰을 복사하는 조합으로 읽히기 때문이다.
그래서 논문의 슬로건이 나온다. 어텐션은 정보의 이동이다. 헤드는 한 위치의 잔차 스트림에서 정보를 읽어 다른 위치의 잔차 스트림으로 옮기는 장치다.
그림 4. 어텐션 헤드는 한 토큰의 잔차 스트림에서 다른 토큰의 잔차 스트림으로 정보를 복사한다. 보통 읽은 곳과는 다른 부분공간에 쓴다.
덧붙일 중요한 사실 하나. 어텐션 전용 모델에서 비선형성은 어텐션 패턴을 만드는 소프트맥스뿐이다. 그래서 패턴 $A$를 고정된 것으로 취급하는 순간, 모델의 나머지 전부는 선형이 된다. 다음 편부터 쓸 경로 전개(path expansion)는 정확히 이 사실 위에 서 있다.
이번 편의 요점
- 잔차 스트림은 계산하지 않는 선형 통신 채널이고, 그 선형성이 가상 가중치 분석을 가능하게 한다.
- 분석 단위는 층이 아니라 헤드다. 헤드들은 독립적이고 가산적이다.
- 헤드는 어디를 볼지(QK)와 무엇을 옮길지(OV)로 분해된다.
- 비선형성은 소프트맥스뿐이다. 패턴을 고정하면 모델은 선형이고, 이것이 경로 전개의 근거다.
- 스트림은 모두가 나눠 쓰는 비좁은 메모리다. 겹쳐 쓰기, 즉 중첩이 여기서 예고된다.
다음 편은 가장 작은 모델, 0-레이어 트랜스포머다.
원문: A Mathematical Framework for Transformer Circuits의 Transformer Overview 섹션. 본문의 도식은 모두 원문에서 가져왔다.