이 논문의 절정은 간단한 패턴 하나다.

[A] [B] ... [A]  →  [B]

앞에서 “Harry Potter”를 보았고 지금 “Harry”가 다시 나오면 “Potter”를 예측한다. 외운 문구가 아니라 현재 문맥에서 발견한 패턴을 이어 쓰는 능력이다. 이를 구현하는 부품이 인덕션 헤드(induction head)다.

이 능력은 헤드 두 개가 협업할 때 생긴다. 앞층 헤드가 단서를 써두고 뒷층 헤드가 그 단서를 읽는다. 논문은 이렇게 앞층의 출력이 뒷층의 계산에 들어가는 일을 합성(composition)이라고 부른다. 먼저 합성의 종류를 정리한 뒤 실제 작동 장면으로 돌아오자.

합성의 세 종류

2편의 그림을 떠올리자. 헤드는 잔차 스트림(residual stream)이라는 공용 보드에서 읽고, 보드에 쓴다. 층이 둘이면 새로운 일이 가능해진다. 앞층 헤드가 보드에 써둔 결과를, 뒷층 헤드가 읽어 갈 수 있는 것이다. 그런데 2편에서 봤듯 뒷층 헤드가 보드에서 읽는 자리는 세 군데다. 검색어(쿼리)를 만들 때, 색인(키)을 만들 때, 내용물(밸류)을 만들 때. 앞층의 결과가 어느 자리로 흘러드느냐에 따라, 합성도 세 종류가 된다.

  • Q-합성(Q-Composition): 앞층 출력이 뒷층의 쿼리에 영향을 준다. 어디를 볼지가 앞층의 계산에 의존하게 된다.
  • K-합성(K-Composition): 앞층 출력이 뒷층의 키에 영향을 준다. 각 위치가 얼마나 주목받을지가 앞층의 계산에 의존하게 된다.
  • V-합성(V-Composition): 앞층 출력이 뒷층의 밸류에 영향을 준다. 옮겨지는 내용물 자체가 앞층에서 가공된 것이 된다.

경로 전개(path expansion)를 2-레이어에 적용하면 길이 하나 더 생긴다. 직접 가는 길, 헤드 하나를 지나는 길에 더해, 앞층 헤드와 뒷층 헤드를 차례로 지나는 길이 생긴다.

V-합성은 두 번의 정보 이동을 한 번으로 묶을 수 있다. 택배를 중간 창고에서 다른 기사에게 넘겨도 출발지에서 목적지로 물건을 옮긴 일이라는 점은 같다. 수식에서는 $(A^{h_2} A^{h_1}) \otimes (W_{OV}^{h_2} W_{OV}^{h_1})$로 정리되고, 논문은 이 묶음을 가상 어텐션 헤드(virtual attention head)라고 부른다. 반면 Q-합성과 K-합성은 어디를 볼지 정하는 검색 규칙 자체를 바꾸므로 하나의 평범한 헤드로 줄일 수 없다.

2-레이어 모델의 경로 전개: 직접 경로, 개별 헤드 항, 가상 헤드 항

그림 1. 2-레이어 로짓의 경로 전개. 직접 경로 항은 바이그램을 담당하고, 개별 헤드 항은 1-레이어 때와 같으며, 마지막 가상 어텐션 헤드 항이 V-합성에 해당한다.

인덕션 헤드

학습된 2-레이어 모델을 분석하면, 합성을 실제로 사용하는 뚜렷한 패턴 하나가 반복해서 발견된다. 인덕션 헤드(induction head)다.

하는 일은 규칙 하나로 요약된다. 문맥에서 [A][B]를 본 적이 있고 지금 토큰이 다시 [A]라면, 다음 토큰으로 [B]를 예측하라. 앞에서 나온 패턴을 그대로 이어 쓰는 것이다. 찾는 토큰이 문맥에 없으면 시작 토큰에 주목한 채 아무 일도 하지 않는다. 일종의 대기 위치(resting position)다.

해리 포터 첫 문단에서 작동하는 인덕션 헤드: 지난번 등장 바로 다음 토큰에 주목해 그것을 예측한다

그림 2. 해리 포터 첫 문단에서 작동하는 인덕션 헤드 1:8의 어텐션 패턴과 로짓 효과. 현재 토큰의 쿼리가 지난번 등장 바로 다음 토큰의 키에 주목하고, 그 토큰의 로짓을 밀어 올린다.

작동 원리는 두 헤드의 협업이고, 핵심은 K-합성이다. 구체 예로 따라가 보자. 문맥에 “Harry Potter"가 나왔고, 지금 토큰이 다시 “Harry"다. 다음은 “Potter"여야 한다.

  1. 앞층의 **이전 토큰 헤드(previous token head)**가 준비 작업을 한다. 각 자리에서 바로 앞 토큰을 읽어 잔차 스트림에 써두는 것이다. 이제 “Potter” 자리에는 “내 앞은 Harry였다"는 쪽지가 붙어 있다.
  2. 뒷층의 인덕션 헤드는 색인(키)을 그 쪽지에서 만든다. 이것이 K-합성이다. 그래서 “Potter” 자리의 색인은 “앞이 Harry였던 자리"가 된다. 지금 토큰 “Harry"의 검색어(쿼리)는 정확히 그런 자리를 찾고, 어텐션은 “Potter"에 꽂힌다.
  3. OV 회로(output-value circuit)는 그냥 복사다. 주목한 “Potter"를 다음 토큰 예측으로 밀어 올린다.

일반화하면 [A][B] … [A] → [B]다. 2편의 언어로 요약하면, QK가 앞층의 쪽지 덕분에 훨씬 똑똑한 검색(“지난번의 나 바로 다음 자리”)을 하게 된 것이고, OV는 4편의 복사 그대로다. 새 부품이 생긴 게 아니라, 새 조합이 생긴 것이다.

인덕션 헤드의 작동 원리: 키가 한 토큰 뒤로 밀려 계산된다

그림 3. QK 회로를 헤드 대신 토큰 단위로 전개한 그림. 키와 쿼리의 세기는 각 토큰이 어텐션 점수를 얼마나 올리는지를 나타낸다. 키가 한 토큰 뒤로 밀려 계산되기 때문에, 같은 토큰을 찾는 쿼리는 그 다음 토큰 자리에 꽂힌다.

1-레이어의 복사 헤드와 비교하면 질적 차이가 보인다. 1-레이어 복사는 “나왔던 토큰이 또 나온다"는 수준이다. 인덕션 헤드는 “나왔던 패턴이 이어진다"를 구현하고, 그 패턴이 무엇이든 상관없다. 완전히 무작위한 토큰 수열을 반복시켜도, 즉 학습 분포에서 한참 벗어난 입력에서도 작동한다. 통계를 외운 것이 아니라 문맥에서 규칙을 읽는 알고리즘을 배운 것이다.

논문은 이 이론을 두 방향으로 검증한다. 하나는 가중치 검증이다. 이 알고리즘이 맞다면 인덕션 헤드의 OV 회로는 복사 행렬이어야 하고, K-합성이 만드는 QK 항은 같은 토큰을 찾는 일치 행렬이어야 한다. 지난 편의 고유값 잣대로 두 축을 재보면, 인덕션 헤드들만 양수-양수의 극단 구석에 모인다. 다른 하나는 절제 실험이다. 전개 항을 차수별로 지워보는 항 중요도 분석(term importance analysis)은, 성능의 큰 몫이 2층 헤드 항에 있고 가상 헤드 항의 몫은 작다는 것을 보여준다.

OV/QK 고유값 양수성 평면: 인덕션 헤드들만 극단의 구석에 모인다

그림 4. 2층 헤드들을 OV 회로의 고유값 양수성(가로)과 K-합성 QK 항의 고유값 양수성(세로)으로 배치한 평면. 인덕션 헤드들만 오른쪽 위, 복사와 일치가 모두 강한 극단의 구석에 모인다.

합성은 생각보다 희소하다

주의할 경험적 발견 하나. 이들의 2-레이어 모델에서 합성은 어디에나 있는 것이 아니라 희소했다. 대부분의 헤드는 합성을 거의 쓰지 않고 1-레이어처럼 스킵 트라이그램(skip-trigram)을 계산한다. 모델 전체는 대체로 큰 1-레이어 모델 위에 인덕션 헤드를 얹은 것처럼 행동한다. 새 능력은 층을 채운 물이 아니라, 몇 개의 가는 물줄기로 들어온 셈이다.

정오표 하나. 논문 공개 후, 이 합성 측정에 쓰인 저자들의 라이브러리에서 버그가 발견되어 수정 도표가 문서에 추가되었다. 실제 모델에는 처음 보였던 것보다 합성이 더 많았고, 일부 인덕션 헤드는 이전 토큰 헤드만이 아니라 최근 몇 토큰을 보는 헤드와도 합성하고 있었다. 다만 핵심 결론, 즉 이전 토큰 헤드와의 K-합성이 인덕션 헤드의 뼈대라는 사실은 그대로다. 주장과 정정을 같은 문서에 붙여두는 이 스레드의 문화는 마지막 편에서 다시 만난다.

정오표의 원본(위)과 수정본(아래) 합성 도표

그림 5. 정오표의 합성 도표. 위가 버그가 있던 원본, 아래가 수정본이다(청록: 인덕션 헤드, 빨강: 이전 토큰 헤드). 수정본에는 합성이 더 많이 보이지만 K-합성이라는 뼈대는 그대로다. 두 도표의 선 굵기는 직접 비교할 수 없다.

가상 어텐션 헤드도 마찬가지다. 이론상 가상 헤드는 헤드 수의 곱만큼 존재할 수 있어서, 큰 모델에서는 엄청난 표현력의 원천일 수 있다. 논문이 드는 정황 하나. 이전 토큰을 보는 헤드는 흔한데, 두 칸 앞을 보는 헤드는 잘 발견되지 않는다. 두 칸 앞의 예측력은 이전 토큰 헤드 둘의 가상 합성으로 얻으면 되기 때문일 수 있다. 인칭이나 시제처럼 옮길 정보가 아주 적은 일에는, 통째로 자리를 차지하는 정규 헤드보다 가상 헤드 쪽이 경제적이기도 하다. 다만 이 논문이 분석한 작은 모델에서 뚜렷하게 큰 역할을 하는 것은 인덕션 헤드 쪽이고, 가상 헤드의 몫은 작았다. 가능성의 표시와 확인된 사실을 구분해 두는 논문의 태도는 배울 만하다.

이번 편의 요점

  • 2-레이어부터 Q-, K-, V-합성이 가능해지고, V-합성 경로는 가상 헤드라는 새 항이 된다.
  • 인덕션 헤드는 이전 토큰 헤드와의 K-합성으로 [A][B] … [A] → [B]를 구현한다.
  • 이것은 임의 패턴에서 작동하는 인컨텍스트 러닝(in-context learning) 알고리즘이며, 이후 대형 모델 연구의 중심 소재가 된다.
  • 다만 관찰된 합성은 희소하다. 2-레이어 모델은 대체로 큰 1-레이어 모델 + 인덕션 헤드처럼 행동한다.

다음 편은 논문의 결산이다. 이 프레임워크가 어디까지 왔고, 무엇이 남았는가.

원문: A Mathematical Framework for Transformer Circuits의 Two-Layer Attention-Only Transformers 섹션. 본문의 도식은 모두 원문에서 가져왔다.