In context learning and Induction heads 이해하기
📝 Original authors: Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, Chris Olah
📝 Translated by: Yong-Min Shin
📝 역자 노트: 이 글은 Transformer Circuits Thread의 두 번째 글이며 (original distill circuits thread 제외), 첫 글에서 나온 induction head에 대해 더 깊게 파헤친 글이다. 상세한 내용과 일러스트가 엄청남에도 아직 글에서 모호한 개념과 설명, 그리고 가끔은 불필요하거나 반복적인 설명이 글의 난이도를 올리고 지나치게 길게 만드는 느낌이 있어 단순 요약이 아닌 번역 밑 요약을 하는 방식으로 진행하였다 (물론 이는 저자들이 확실한 증거 없이 성급한 결론을 지으려는 것을 굉장히 조심스럽게 여기며 이를 막기 위한 노력의 일부이기도 하다). 또한 역자도 아직 mechanistic interpretability를 공부하고 있으며 번역 또한 그 일환으로서 진행하는 것이기에 내용적 오류가 있을 수 있다. 용어들은 억지로 한글로 변환하면 더 어려워지는 경향이 있어서 (예: co-perturbation → 동시-섭동) 웬만하면 그대로 두었고, 또한 때에 따라서는 주관적 판단 하에 글을 재편집하거나 추가 설명 등을 붙이는 등 단순 번역을 넘어서는 작업이 되도록 노력하였다.
원문 링크: transformer-circuits.pub
원 번역 글 Notion link: In context learning and Induction heads (요약 번역) - coming soon
1. Induction head 한 줄 요약
Transformer 안에서 “복사 + 붙여넣기” 역할을 수행하는 circuit을 말한다.
2. Induction head 세 줄 요약
Transformer의 attention head들 중 “(이전 context로부터) 복사 + 붙여넣기” 역할을 수행하는 head로 이루어진 circuit을 말한다.
In-context learning을 가능하게 하는 주 원인일 가능성이 높으며, 작은 모델 (특히 attention-only)이면 더욱 그렇다.
학습 중 phase change라는 구간에서 갑작스럽게 형성이 된다.
3. 원문 내용
전반적인 내용 요약 (Introduction)
다른 글들과 마찬가지로, Anthropic에서 Mechanistic interpretability [1] 로 연구를 진행중에 작성된 article이다.
이러한 복사 + 붙여넣기는
[A][B] … [A] → [B]
로 표현되며, 이를 풀이하면…[A]
,[B]
는 token을 나타낸다.작동 서순을 반영하여 이를 읽으려면 다음과 같다.
중간에 나오는 두 번째
[A]
가 현재 보고 있는 token이다.Previous token head라는 또 다른 attention head가 이전 token들을 훑어보며 (현재 보고 있는)
[A]
랑 같은 (혹은 비슷한) token이 있는지 본다 (첫 번째[A]
).Induction head가 previous token head가 찾아낸 첫 번째
[A]
다음에 나오는 token (첫 번째[B]
)을 보고 다음에 올 token으로[B]
를 예측한다 (→ [B]
).
즉 정확히 말하자면 induction head와 previous token head의 합작으로 복사 + 붙여넣기를 진행하는 것이다.
해당 글에서 대부분의 실험은 2-layer (attention-only) transformer에서 실험을 진행했기 때문에, induction head가 더 현실적인 모델에도 존재하는지에 대한 추가적인 증거가 필요함을 역설한다. 하지만 레이어를 더 쌓고 MLP 모듈이 추가되면 해석의 난이도가 매우 상승하는 문제가 있다 (의외로 MLP가 mechanistic interpretability 관점에서 다루기 매우 까다롭다고 한다). 그럼에도 불구하고, 이 article에서는 induction head가 더 크고 현실적인 transformer에도 존재하며, 많은 in-context learning에 결정적인 역할을 한다는 가설에 대한 (간접적인) 증거 6개를 제시한다.
In-context Learning
현재 token이 주어졌을 때, 다음에 올 가능성이 높은 token을 예측하는 transformer의 입장에서, 지금까지 이미 주어진 token (context라고 불린다)이 많으면 많을수록 예측이 쉬워지는 특징이 있다. 이런 특성을 이용하여 모델의 능력치를 (re-training이나 tuning 없이) 더욱 끌어올릴 수 있으며, 그 정도가 유의미하기 때문에 in-context learning이라고 명명되었다. 이 현상은 GPT-2 [2]에서 처음 언급되었으며, GPT-3 [3]부터 본격적으로 주목을 받게 되었다.
이 글에서는 in-context learning의 성능을 측정하는 방법으로 ([3]에서 사용된 few-shot learning과 유사한 방법론보단) [4]에서 제시된 token index에 따른 loss값을 관찰하는 방식으로 in-context learning의 성능을 정의한다.
구체적으로, in-context learning을 측정하는 기준은 다음과 같이 세웠다:
🔍 In-context score = 500번째 token에서의 loss 값 - 50번째 token에서의 평균 loss 값 (여기서 데이터셋 평균을 의미)
물론 ‘500’과 ‘50’이라는 값을 채택한 데에는 큰 이유는 없지만, 이 숫자들을 바꾼다고 해서 전반적인 경향성이 달라지지는 않으며, 이 또한 추후에 보이게 된다.
Induction Heads
이전 글에서 “induction head”를 (2-layer transformer에서 발견되는) 소개했었다. Induction head는 각자 다른 layer에 속하는 두 개의 attention head가 서로 합동하여 나타나는, 복사 + 붙여넣기 혹은 내용 완성의 역할을 하는 circuit이다 [5, 주의1]. 첫 layer의 attention head는 직전 token에서 현재 token으로 내용을 복사하고, 두 번째 layer의 attention head는 이렇게 복사된 정보를 바탕으로 이전 token 정보들에 접근을 할 수 있다.
구체적으로, 현재 주어진 [A]
라는 token이 있다고 하자. 두 번째 layer에 속하는 어떤 attention head (이를 induction head라 한다)가 앞선 token sequence들을 훑어보며 [A]
가 등장한 적이 있는지 확인한다. 만약 그렇다면, 그 당시 다음에 어떤 token이 왔는지 확인하고 (이를 [B]
라고 하자), 이 [B]
token에 대한 logit 값 (혹은 그 확률값)을 높이는 역할을 수행하게 된다. 즉, …[A] [B] … [A]
라는 상황에서 두 attention head가 합작하여 다음에 [B]
가 오게 될 확률을 높이는 것이다.
주목할 만한 특징은, induction head가 N-gram [참조]의 통계적 모델이 아닌, (간단한) 알고리즘을 구현해낸다는 것이다. 즉, [A][B] … [A] → [B]
패턴은 [A]
와 [B]
자리에 어떠한 token이 와도 작동하는 ‘규칙’인 것이다. 이러한 특징으로 인해 induction head는 OOD (out-of-distribution) 상황에서도 어느 정도 작동할 것으로 예상할 수 있으며 (초반과 뒷 부분의 context가 너무 다르지 않다는 가정), 따라서 더 일반적인한 상황에서도 좋은 성능이 나올 수 있다는 것이다.
이번 글의 목표는 다음 가설에 대하여 조금 더 확장적인 증거들을 제시하는 것이다:
📝 단순 `[A][B] … [A] → [B]`이라는 패턴에서 그치치 않고, induction head가 더 큰 transformer 모델에서까지 in-context learning이 이루어지는 데에 주요한 역할을 담당하고 있다.
이를 명확하게 진행하기 위해서는 일단 induction head에 대한 더욱 구체적인 정의부터 필요하다. 이전 글에서는 induction head의 몇 가지 예시 및 분석에 집중하였다 (2-layer attention-only transformer 모델에서 관찰된 induction head들과 transformer를 재해석한 그림 위에서의 분석, 그리고 eigenvalue 기반 분석). 하지만, 아직 induction head의 정확한 정의를 내리지는 않았고, 관측된 내용에 기반하여 모델 내부적인 요소들이 합작하여 이루어낸 현상이라는 (다소 모호한) 설명만 남기는 정도에서 끝났다.
조금 더 구체적으로 정의를 내리기 위해서, 이 글에서는 induction head를 일단 눈으로 관찰할 수 있는 [A][B] … [A] → [B]
패턴에 기반한 정의를 먼저 내린다. 그 뒤에 (1) in-context learning과 연관이 있는 몇몇 기능으로 확장할 수 있으며 (2) 작은 모델에서 분석되었던 기존의 개념들과도 일치한다는 것을 보이도록 하겠다. (이는 Mechanistic (기계적인) 혹은 eigenvalue에 기반한 기존의 접근법으로는 모델 사이즈가 커지고 MLP가 포함되기 시작하면 훨씬 더 까다로워지기 시작한다는 한계가 있다는 이유도 포함되어 있다)
우리는 induction head를 임의로 반복되는 token sequence에 대해 다음 두 성질을 보이는 attention head으로 정의한다:
Prefix matching (접두사 매칭): 현재 보고 있는/최근에 등장한 token이 등장하는 (복수의) 이전 token들 다음에 따라붙는 token들에 attention 값들이 반응을 한다 (일종의 fuzzy matching이라고 볼 수 있음). 즉, induction head가 다음에 올 token으로 추천할만한 후보 token들을 과거 sequence에서 골라내는 것이다. (밑 그림에서 붉은색 부분에 해당한다.)
Copying (복사): Prefix matching이 추천한 후보 token에 대한 logit 값, 즉 다음에 올 token으로 추천될 확률을 높인다. (밑 그림에서 보라색 부분에 해당한다)
따라서, induction head는 이전 sequence들을 통째로 반복하는 데에 특화가 되어있다. 예를 들어, “The cat sat on the mat. The cat …”이라는 문장이 주어져있을 때, induction head는 이전 문장을 계속 반복하여 “…set on the mat”로 완성시키려고 한다는 것이다. 임의의 token sequence를 반복하는 능력을 학습하기 때문에, in-context learning 및 (일부 케이스에서는) few-shot learning에 induction head가 어떠한 역할을 할 지에 대해 가설을 세울 수 있다.
추후에는 2-layer 모델보다 더욱 큰 모델에서 induction head에 관여하는 attention head들이 단순 복사 + 붙여넣기를 넘어서서 analogical sequence copying (추상적인 의미를 가지고 하는 복사 + 붙여넣기) 및 in-context nearest neighbors와 같은 추가적인 역할에도 가담하는 증거들을 볼 것이다. 즉, [A][B] … [A] → [B]
형태의 sequence completion 뿐만 아니라, [A*][B*] … [A] → [B]
와 같은 형태에서도 비슷한 역할을 수행한다는 것이다 (예를 들어, [A*]
는 [A]
와 같은 token은 아니지만, 같은 단어이지만 언어가 다르다던지 등의 이유로 embedding space에서 비슷할 수 있다). 아직 mechanistic interpretability에서 이러한 induction head가 모든 transformer 모델에서 일반적으로 이러한 행동을 한다는 것을 증명하지는 못했지만, 증거 4에서 몇 가지 긍정적인 예시들을, 증거 5에서는 (작은 모델에서 발견한) 기존의 induction head들이 더 큰 모델들로 확장 가능하다는 것을 보여준다.
Per-Token Loss Analysis
모델 학습 중 일어나는 변화를 더욱 면밀히 관찰하기 위해, 우리는 “per-token loss vectors”라는 것을 관찰하기로 하였다. 핵심 아이디어는 [6]으로 거슬러 올라가며, 수학에서는 함수 공간이라는 개념과 관련이 있다 [부연 설명].
일단 세팅은 다음과 같다.
데이터셋에서 무작위로 10,000개의 token sequence들을 준비하고 (무작위로 선별된 10,000개 샘플은 비교 기준을 세우기 위해 다른 실험들에서도 동일하게 고정되었으며, 모든 token sequence 샘플들은 길이가 512이다), 이에 대한 token 별 log-lilkelihood를 기록한다. 이 단계 이후부터는 무엇을 관찰할 것인지에 따라 추후 분석이 다르다 (이 부분과 3번에 대한 구체적인 내용은 뒷 부분 실험에서 최대한 각주를 달아보도록 하겠다).
PCA를 활용하여 가장 변화 폭이 큰 방향을 찾아내고 이를 기준으로 분석을 진행한다 [부연 설명].
이를 기반으로, 메인 주장과 세부적인 증거 6개를 나열하면 다음과 같다.
주장하고자 하는 내용: Induction head가 대부분의 in-context learning에 있어서 주요한 역할을 한다.
Introduction에도 이야기가 되었듯이, 이 가설은 다소 포괄적인 성격이 있으며 제시하는 증거들은 대부분 간접적이다. 그럼에도 불구하고, 증거들을 종합하면 (확정적이지는 않더라고) 이 주장에 대한 꽤나 강력한 증거가 될 것이라 생각된다.
구제척인 내용으로 들어가기 전에, 증거들이 더 확정적일 때 vs. 덜 확정적일때를 구분하고자 한다 (밑의 table을 참조). 종합적으로 보면, 작은 attention-only 모델들의 경우 ablation 및 reverse engineering등을 통해 메인 주장에 대한 강한 증거들을 모으는데 성공했다고 생각한다. 또한, 모든 크기의 모델 (작은 attention only 모델, 작지만 MLP도 포함된 모델, 대형 모델)에서도 해당 주장이 어느 정도는 참이라는 것은 다양한 예시 및 여러 상관관계를 보여주는 방식으로 보여주었다고 생각한다. 다만, 모델이 커지게 될 수록 induction head가 in-context learning의 대부분의 비중을 차지한다는 주장에 대한 입증이 힘들어진다. (MLP까지 포함된) 대형 모델의 경우 대부분의 증거들은 (인과관계보다는) 상관관계를 보고 있으며, 교란(confounded)되었을 가능성이 있다. 이와 관련해 다른 가설들도 증거 1 및 증거 6에서도 다룬다.
우리가 제시할 6가지 증거에 대한 요약을 하자면 다음과 같다:
증거 1 (거시적인 동시-발생 관찰 / Macroscopic co-occurence): Transformer 모델들은 학습 초기 단계에 “위상 변화 (phase change)”를 겪게 되면서 induction head가 형성된다 (그리고 그에 따라 in-context learning 능력이 크게 향상된다).
증거 2 (거시적인 동시-섭동 관찰 / Macroscopic co-perturbation): Induction head 형성에 영향을 미치는 방향으로 transformer 아키텍쳐를 변형하면 in-context learning 성능에서도 이에 상응하는 변화가 일어난다.
증거 3 (Direct ablation): Test time때 induction head들을 직접적으로 처내면, in-context learning 성능이 급격히 저하된다.
증거 4 (일반적인 경우에서의 induction head 예시): Induction head들이 (비슷하지만) 더 복잡한 형식의 행동에도 관여한다는 것을 보인다.
증거 5 (Induction head의 mechanistic 타당성): 작은 모델에서 induction head가 mechanistic하게 설명이 됨을 보이고, 더 일반적인 모델에서도 가능케 하는 자연스러운 확장이 있음을 보인다.
증거 6 (작은 모델에서 큰 모델으로의 연속성): In-context learning 밑 induction head와 관련한 패턴들과 데이터들이 작은 모델에서부터 큰 모델까지 연속적인 변화를 보이며, 따라서 큰 모델에서도 사실상 같은 방식으로 작동하고 있음을 제시한다.
증거 1: Transformer 기반 언어 모델들은 학습 과정 중 phase change를 겪게 되며, 이때 induction head가 형성되며 in-context learning 성능이 향상하게 된다.
해당 증거는 모델 학습 도중 in-context learning과 induction head의 존재에 대한 상관관계를 분석하면서 나온 결과에 기반한다. 결론적으로, 이들 사이에 강한 상관관계가 (모델 크기와 데이터셋에 상관없이) 관측된다.
[실험 결과 1-1]
In-context learning 성능 (구체적인 정의는 앞선 내용 참조)을 transformer 모델의 학습 과정 동안 꾸준하게 관찰하면, 학습 과정 중 꽤나 좁은 구간 (~2.5에서 5B token) 사이에 급격하게 성능이 좋아지는 것을 보인다 (위 figure 참조: 2&3-layer transformer에서 노란 구간에서 in-context loss가 급격하게 떨어진 이후, 어느 정도 유지된다). 비록 위 예시에서는 3개의 모델만 보여줬지만, 이러한 경향성은 매우 일반적으로 나타난다 (참조). [구체적 데이터 처리 방법1]
[실험 결과 1-2]
학습이 진행될수록 각 token 위치 별로 loss를 찍어봐도 같은 결론에 도달할 수 있다 (윗 figure 참조). [구체적 데이터 처리 방법2]
[실험 결과 2]
갑작스럽게 in-context learning 성능이 좋아지는 구간에서 attention head를 살펴보자. Induction head를 정의하는데 쓰인 성질 중 하나인 prefix matching을 척도로 삼아 score로 환산해서 attention head를 분석할 수 있다. 역시나 해당 구간에서 induction head로 추정되는 몇몇 attention head들이 급격하게 생성되는 것을 확인할 수 있다 (밑 figure 참조). 물론 이 현상 또한 더 큰 모델에서도 일반적으로 나타난다 (참조).
[종합 1]
실험 결과 1-1, 1-2, 그리고 2를 종합해볼 때, induction head와 in-context learning 사이에 어느 정도 관계가 있음을 짐작할 수 있다. 이는 loss function에서도 (살짝 아래로 푹 꺼지는 형태로) 나타나기는 하나, training loss에는 전반적으로 기여하는 token의 수가 워낙 많기 때문에 in-context learning의 성능 향상에 의한 loss는 전체 training loss에 비해 미시적일 수 밖에 없다 (밑 figure 참조).
[실험 3]
모든 모델들에 대한 (model snapshot X 10,000 token sequence 샘플) 데이터를 모두 모아 하나로 합쳐보았다 (즉, 작은 모델 12개에 대해서는 snapshot이 각 200개였으므로 2400 X 10,000 짜리 loss matrix가 나온다). 이 matrix에 대해 PCA를 진행하여 첫 번째 및 두 번째 principal component를 구하고, 이를 기준 축으로 하여 다시 각 모델별로 plot한다. 왜 PCA를 하며 이렇게 하는 의미가 무엇인지는 이 [부연 설명]을 참조하면 된다. 한 마디로 학습 과정 중 함수가 어떻게 변하는지를 압축적으로 나타내준다고 이해하면 되며, 각 축이 특별하게 의미하는 것은 없다.
여기에서도 비슷한 결론이 나온다. 즉, in-context leanring 성능이 향상되는 시점에서 모델의 학습 궤적이 급격하게 꺾이는 모습이 보이며, 이러한 현상은 2/3-layer 모델들에게서만 관찰된다.
[종합 2]
결론적으로, 다음 4 가지 현상이 학습 도중 비슷한 시기에 겹친다는 것을 관찰하였다:
In-context learning 성능의 향상.
Induction head의 생성.
Loss에서 살짝 파이는 모양이 생김.
모델의 학습 궤적이 갑자기 틀어짐.
이러한 모델의 변화를 “phase change (위상 변화)”로 명명하도록 하겠다.
Phase change 더 깊게 파헤쳐보기
지금까지 본 결과들을 보아하니, induction head가 in-context learning 특성을 만들어내고 이에 따라 다른 변화들이 관찰되는 것으로 해석하는 것이 자연스러워보인다. 이 가설에 힘을 실어주기 위해 몇 가지 사항들을 추가적으로 확인하였다:
첫 번째로, phase change가 일어나는 시기는 learning rate, warmup, 혹은 weight decay의 영향을 크게 받지 않는 것으로 보인다.
두 번째로, 이러한 phase change가 다른 데이터셋으로 학습할 때에도 관측이 된다 [9].
세 번째로, 학습이 진행되면서 모델이 예측을 더 잘하거나 못하는 token들을 따로 분석하였다 (이 부분을 밑에서 자세히 소개한다).
마지막 사항의 분석을 위해, 일단 테스트용 텍스트를 하나 골라보자. 재미를 위해, 해리포터 1권의 첫 문단을 분석해보겠다.
[실험 4]
Phase change가 사작될 때 (밑 figure에서 A 지점)와 끝났을 때 (밑 figure에서 B 지점), 이 두 시점에서 테스트용 텍스트에 대한 모델의 log-likelihood 차이를 보겠다.
위 figure 왼쪽에는 기존 PCA 분석에서 phase change 전과 후 지점을 나타냈으며, 오른쪽 해리포터 1권 첫 문단 텍스트에서는 phase change 전후 log-likelihood의 차이를 보여준다. 붉은색이 진해질수록 phase change 이후 해당 token들을 더욱 쉽게 예측할 수 있었다는 것을 의미하고, 초록(?)색이 진해질수록 예측이 더 어려워졌음을 의미한다.
결론적으로, 비슷한 순서의 token들이 반복되는 경우 모델이 phase change 이후 더욱 쉽게 예측하게 되며, 반복되지 않는 token sequence에서는 반대의 결과가 나온다. 이에 대한 가장 대표적인 예시는 이름과 같은 고유명사들이다. Tokenizer들의 특성상 이름의 경우 조금 특이하게 찢어지는 경향이 있다. 예를 들어, ‘Dursley’ → ‘[D][urs][ley]
’, 혹은 ‘Potters’ → ‘[Pot][ters]
’가 되는 식이다. 이름은 바뀌지 않기 때문에, 해리포터 내에서 [D]
뒤에 [urs]
, 그리고 그 뒤에 [ley]
가 나오는 패턴은 정말 많을 것이다. 그러기에 phase change가 끝난 뒤에는, (아마도 induction head에 의해) 이러한 반복적인 패턴은 쉽게 예측이 되는 것이다.
[실험 5]
비슷한 방식으로, 모델이 학습하는 과정중 전체적인 training loss가 아닌 특정 token에 대한 loss에 대해서만 집중해서 분석을 진행하였다.
분석 대상이 되는 두 token sequence와 그 위치는 다음과 같다.
‘The Dursleys’ (token 단위로는
[The][D][urs][leys]
), 마지막으로 나타나는 위치 (’…The Dursleys shuddered to think what the neighbors would say if the Potters…’)에서 분석.‘ Mrs Potter’ (token 단위로는
[ Mrs][Pot][ters]
), 첫 번째로 나타나는 위치 (’…Mrs Potter was Mrs Dursley’s sister, but they hadn’t met for several years;…’)에서 분석.
해당 예시에서도 비슷한 결론에 도달한다: 이미 앞서 몇 번 반복된 적이 있는 ‘The Dursleys’는 phase change 이후 loss가 급격하게 낮아지는 반면, 처음 ‘ Mrs Potter’는 이전 context에서 살짝 다른 ‘ Mrs Dursley’를 본 적이 있기 때문에 오히려 다른 패턴이 나타나자 loss가 올라간다.
지금까지의 증거들에 대한 평가
앞서 증거들을 보아 induction head와 in-context learning이 학습 과정 중 아주 비슷한 시기에 함께 발현된다는 경황이 뚜렷하지만, 그럼에도 불구하고 아직 induction head가 in-context learning을 일으키는 주요한 작동 원인이 된다는 것을 직접적으로 보인 적은 없다. 오히려, induction head가 in-context learning의 원인 (인과적인 관계)이라는 결론을 내리기 조심스러운 몇 가지 counfounding의 가능성들을 생각해볼 수 있다.
먼저, ‘induction head가 in-context learning의 원인이 된다’라는 가설에 힘을 실어주는 주장 몇 가지를 살펴보면:
Induction head가 in-context learning 성능이 크게 향상되는 시기에 형성된다는 것 (모델 크기와 상관없이 관측 가능).
induction head와 in-context learning이 아무런 인과관계 없이 순수 행운으로 겹치기엔 너무 다양한 모델에서 일관적으로 관찰이 됨.
적어도 어느 정도의 연결고리가 있다는 것은 높은 확률로 참이며, 아마 그 연결고리는 ‘induction head가 직접적으로 in-context learning을 일으킨다’일 것이다 (하지만 이 부분은 밑에서 반대되는 입장으로 다시 다룰 것).
Phase change 구간 동안 in-context learning의 75% 정도가 학습이 되기에, induction head가 이 정도의 역할을 수행할 것이라고 생각할 수 있다.
반대로 다음과 같은 이유들로 인과관계를 확정할 수 없다:
큰 모델들에 대해서는 학습 과정이 상대적으로 낮은 해상도에서 분석되었다 (snapshot이 15개 밖에 안된다는 이야기).
Phase transition때 아직 파악하지 못한 다른 기능이 형성이 되었을수도 있으며, 그 기능이 induction head와 in-context learning 모두에게 영향을 미치는 것일 가능성을 배제할 수 없다 (confounding의 전형적인 예시).
In-context learning을 측정할 때 상대적인 loss 값을 측정했기에 (50th vs. 500th. token), 비록 이 값이 phase change 이후에 크게 변하지 않더라도 모델 안에서 in-context learning을 담당하는 부분 또한 변하지 않는다는 것은 아니다. 학습이 진행될수록 50th token의 (절대적인) loss가 계속 좋아진다는 것과, 500th token loss를 줄이는 것이 학습 초기와 비교했을 때 더 어려울 것이다. 그러기에 학습 뒷 부분에서는 in-context learning 외적으로 loss를 줄이는 변수가 있을 수 있다.
한 가지 주목해볼만한 것은 ‘induction head가 phase change 시점에서 in-context learning의 대부분에 기여한다’가 ‘학습이 다 끝났을 때 induction head가 in-context learning의 대부분에 관여한다’보다 가능성이 높다는 것이다. 학습이 진행될수록 다른 요소들이 loss를 개선하는데 기여할 가능성이 높기 때문이다.
증거 2: Induction head 형성에 영향을 미치는 방향으로 transformer 아키텍쳐를 변형하면 in-context learning 성능에서도 이에 상응하는 변화가 일어난다.
증거 1의 가장 큰 단점을 꼽자면, 결국 관찰에 기반한 정황증거라는 것이다. 이러한 관찰보다 더 확실한 것은 우리가 직접 모델 아키텍쳐에 어떠한 수정을 가한 후에 무슨 일이 일어나는 지를 보는 것이다 (*개입적). 다시 말해, 이번에는 induction head가 형성되기 더욱 용이한 방향으로 모델에 수정을 가한 후 in-context learning 성능이 어떻게 변화하는가를 관찰할 것이다. 이러한 실험은 작은 모델에서 더욱 긍정적인 결과로 나타나며, 큰 모델에서도 어느 정도 같은 경향성을 보인다.
실험 설계를 하기 전에, induction head가 생성되기 위해서 적어도 layer를 2개 이상 쌓아야 한다는 점부터 시작하자. 즉, attention layer간의 조합이 요구 된다는 것이다 [Composition]. (1-layer 모델을 포함하여) 임의의 layer를 가지는 transformer 모델이 induction head를 가진 것 처럼 행동할 수 있게, 우리는 “smeared key [10]”라는 것을 이용해 모델을 변형시켰다. 일단, 각 attention head 마다 학습 가능한 parameter인 가 있다고 가정하고, 실제로는 로 0 과 1 사이로 변환한뒤 사용한다 [11]. 이 값들은 현재 (-th position) 및 이전 token (-th positon)의 key vector 사이를 interpolate 하는데 사용된다: [12]
즉 2-layer 모델 (혹은 그 이상)에서만 존재했던 previous attention head의 역할을 (라는 변수를 통해) 직접 해주는 것이다. 기존의 가설 하에서는 이러한 변형을 거쳤을 시 1-layer 모델에서도 phase change가 일어날 것이라고 예측을 할 수 있다. 나아가, 기존의 2-layer 모델 (혹은 그 이상)에서도 phase change가 더 앞당겨져서 일어날 것으로 예상이 된다. 만약 induction head가 in-context learning에 영향을 미치는 주된 요인이라기보다는 그저 다양한 요인 중 하나에 불과한 경우라면, 약간의 in-context learning 성능 향상이 초기 단계에 일어나고, 나머지 성능 향상은 원래 phase change가 일어나던 시점에 일어날 것이다.
실험 결과는 우리가 예상했던 결과대로 나왔다. 즉, smeared key를 사용하게 되면 1-layer 모델에서도 in-context learning이 가능해지고, 2-layer 모델에서는 그 시점이 앞당겨진다.
하지만 더 큰 모델들에 대해서도 같은 현상이 일어날 것이라는 결론을 내리는 데에는 조심스러워야 할 것이다. 이 실험으로 결론을 내릴 수 있는 것은 induction head가 in-context learning을 발현시킬 수 있는 가장 작은 단위의 메커니즘이라는 것이며, 큰 모델에 대해서는 in-context learning에 개입하는, 아직 논의되지 못한 더 복잡한 메커니즘이 충분히 있을 수 있다. 또한 학습 과정 중에 in-context learning의 메커니즘이 고정되어있지 않고 계속 변하고 있는가에 대한 단서를 제공하는 것도 아니다.
증거 3 (Direct ablation): Test time때 induction head들을 직접적으로 처내면, in-context learning 성능이 급격히 저하된다.
증거들을 종합적으로 모았을 때, 이번에 다룰 ablation 기반 증거들이 제일 강력하다고 할 수 있다 [Ablation]. 기본적으로, induction head를 (먼저 어떤 head가 induction head인지 파악한 후) 어떻게든 제거하게 되면 in-context learning의 성능이 크게 저하된다는 추론에 기반하였다. “제거”한다는 것은, test time때 해당 attention head를 모델에서 말 그대로 삭제한 뒤 평상시처럼 input을 통과시켜 그 결과값을 본다는 것이다.
위에 정리된 실험은 attention head들을 하나씩 제거해나가면서 in-context learning에 어떠한 영향을 주는지 관찰한 결과를 보여준다. 적어도, 작은 모델들 안에서는 거의 모든 in-context learning이 induction head로부터 나온다는 것을 볼 수 있다! 이러한 경향은 phase change가 시작될 즈음에 같이 보이고, 학습이 끝날 때 까지 지속된다.
안타깝게도, 우리는 더 큰 full-scale 모델들에 대해서는 ablation 결과가 없다 [Attention head ablation 방법론] [13]. 그래도 실험 결과들이 있는 작은 모델들에 대해서는 induction head와 in-context learning의 관계를 상당 부분 정리해주는 것으로 보인다. 하지만 “induction head가 in-context learning의 주된 요인인가?”에 답하기 위해서는 몇 가지 사항을 고려해봐야한다:
Attention-only 모델에서는 거의 필연적으로 in-context learning 현상이 attention head들의 복합적인 현상이지만, MLP까지 포함되어있는 모델에서는 MLP-attention 사이의 상호작용이라는 변수를 고려해야하며, 이 상호작용을 분석하는 일은 확연히 더욱 복잡하다. 이러한 상황에서 attention head ablation를 하는 것은 어느 정도는 유효할 것이지만 완전한 분석을 할 수 있을 것이라는 보장은 할 수 없다 (MLP가 포함된 모델에 대한 실험들은 본문에서 모델 분석 테이블을 참조해야한다).
우리가 진행한 실험은 각 attention head가 모델에 주는 영향을 하나씩 분석한 것으로, 두 개 (혹은 그 이상)의 attention head가 비슷한 역할을 한다거나 layer norm이 중요도를 왜곡하는 가능성에 대해서는 여전히 분석의 취약점이 있다.
종합적으로, 우리는 induction head가 작은 attention-only 모델에 한해서는 in-context learning의 주된 원이다라는 결론을 내릴 수 있다.
증거 4 (일반적인 경우에서의 induction head 예시): Induction head들이 (비슷하지만) 더 복잡한 형식의 행동에도 관여한다는 것을 보인다.
지금까지 보아온 증거 1~3은 induction head의 형성 과정에서의 관측 혹은 변형을 통해 전반적인 in-context learning을 보았다. 이와 다른 접근법이긴 하지만, induction head가 상당히 어려워보이는 in-context learning 현상들을 일으키는 예시들을 찾아내볼 수 있다. 이러한 예시들은 결국 induction head들이 대부분의 in-context learning의 원인이 될 수 있다는 가능성을 시사할 수 있으며, 이러한 예시들은 매우 큰 모델 (여기에서는 최대 12B 까지의 모델에 대하여 실험한다)까지 적용된다. 하지면 결국 일부 task에 대해서만 결과를 보여주는 것이기에 이에 기반하여 더욱 일반적인 주장을 펼치는 것은 무리다.
앞서 우리가 induction head를 임의의 token sequence들을 ‘prefix matching’ 패턴에 따라서 복사하는 현상으로 정의했음을 떠올려보자. 이번 증거에서 우리의 목표는 정의에 부합하는 모습 뿐만 아니라 다른 흥미로운 행동을 하는 head를 찾아내는 것이고, 결국 이러한 induction head들이 ‘일반화’ 될 수 있는 가능성을 제시하는 것이다.
이번 증거에서는 더 큰 transformer (40-layer 13B 모델)에서 앞서 말한 유용한 예시들을 찾아나설 것이다 (예를 들어, 복사 붙여넣기, 번역, 그리고 일종의 패턴 매칭 현상 등). 이러한 예시들은 모두 [A*][B*]…[A][B]
(즉 ‘fuzzy nearest neighbor match’, 혹은 ‘앞선 token에서 비슷한 것을 찾아내 현재 sequence를 비슷한 의미로 완성하는 형태’)의 패턴을 가지고 있다는 공통점이 있다. 이러한 head들은 우리가 고안한 “복사하기” 및 “prefix matching” 실험에서 좋은 성능을 낸다는 것을 보인다 (즉, 무작위 텍스트에서도 [A][B] … [A] → [B]
의 현상이 보인다). 이는 곧 우리의 (empirical한) 정의에 의하면 induction head라는 것이며, 이들은 다른 신기한 현상들에도 관여한다.
다음 표는 몇 가지 예시가 되는 head들을 종합하여 보여주며, 이에 대한 자세한 설명을 진행하도록 하겠다.
[Literal sequence copying head]
가장 간단한 케이스로, 반복되는 텍스트를 말 그대로 복사해주는 head를 예시로 보여준다. 이는 visualization 인터페이스 및 기본적인 메커니즘에 익숙해질 수 있다는 점도 존재한다 (실제 원문에서 제공하는 interactive visualization을 반드시 확인해보도록 하자). 우리는 아주 기초적인 복사를 담당하는 것으로 보이는 induction head를 골라내어 해리포터의 첫 문단을 어떻게 받아들이는지 보기로 하였다. 첫 문단 이후에는 더 긴 문장에서 어떤 식으로 작동하는 지 보기 위하여, 똑같은 텍스트를 반복하였다.
이번 visualization에는 다음과 같은 것들이 포함되어있다:
붉은 색으로 표기된 attention 값: 다음에 올 token을 예측하기 위해 뒤따라오는 token을 보고 있다.
청색으로 표기된 logit 값: 현재 token 예측에 기여를 한 이전에 등장했던 같은 token을 보여주고 있다.
스크린샷으로 대체한 몇 가지 예시는 다음과 같다:
[Translation head]
이러한 모델들이 번역을 할 수 있다는 사실 또한 잘 알려져있다. 흥미롭게도, 우리는 번역을 할 수 있는 많은 induction head들을 찾을 수 있었다. 여기에서 우리는 한 예시로, 7번째 레이어에서 찾은 head를 가지고 와 영어, 프랑스어, 그리고 독일어로 번역하는 것을 관찰한다 (추가적으로, 이 head 또한 우리가 내린 induction head 정의에 부합한다).
📝 Input token sequence: EN: This is the largest temple that I've ever seen. FR: C'est le plus grand temple que j'ai jamais vu. DE: Das ist der größte Tempel, den ich je gesehen habe.
스크린샷으로 본 attention 패턴을 보면, 대부분 대각부분 (diagonal)에서 조금 떨어진 위치에서 값을 가지는 경우가 많은 것을 볼 수 있다. 또한 대각선과 나름 평행한 부분들도 완벽한 대각이 아닌 약간씩 뒤틀린 모습을 보이고 있다. 이는 언어끼리 대체적으로 어순은 비슷하지만, 문법적인 이유로 완벽하게 단어 순서가 맞는 일은 드물기 때문에 일어나는 일이다. 이 head는 과거 (다른 언어로 쓰여진) token들을 보면서 의미를 고려했을 때 다음에 와야할 token들을 찾으면서 조금씩 왔다갔다 거린다.
Logit 패턴들 또한 완벽하게 도드라지지는 않는다. 즉, attention이 다른 언어로 쓰여진 정확하게 매칭이 되는 token을 보고 있어도, 직후에 오는 token의 logit 값들을 항상 올려주지는 않는다. 우리는 head의 output이 뒤따라오는 layer들에 의해 추가적으로 처리가 되어야하기 때문에 이러한 일이 일어나고 있다고 추측하고 있다. 하지만, 전반적으로는 logit 패턴이 정확한 번역에 어느 정도 기여를 하고 있다는 것을 명확하게 보여준다.
[Pattern matching head]
마지막으로, 26번째 레이어에서 찾은 head를 보여주도록 하겠다. 이 head는 조금 더 복잡한 형태의 패턴 매칭을 하며, 간단한 함수를 구현하고 있다고도 생각할 수 있다 (그리고 이 head 또한 induction head 정의에 부합한다).
이를 탐구하기 위해, 단순한 규칙성을 지니는 텍스트를 만들어내었다. 각 줄은 네 가지 중 하나의 템플릿을 따르며, 어느 템플릿에서 왔는가를 나타내주는 label이 뒤따라오도록 만들어졌다. 템플릿 또한 무작위로 선택되며, 템플릿에 채워지는 단어들 또한 무작위이다.
[월] [동물]: 0
[월] [과일]: 1
[색] [동물]: 2
[색] [과일]: 3
다음은 이러한 텍스트에서의 attention 및 logit 값들을 보여준다. 조금 더 visualization을 단순화 하기 위해, 현재 보고 있는 token이 [:]
인 경우에 대해서만 보여주도록 처리하였다. 예시 스크린샷은 다음과 같다 (실제 원문에서 제공하는 interactive visualization을 반드시 확인해보도록 하자):
마지막 “March apple: 1”의 [:]
token에 대해서 살펴보자. Attention을 보면 1번 템플릿을 사용했던 이전 예시들에 해당하는 [:]
token 다음 값 (즉, “1”)을 보고 있고, logit값은 해당 값 (”1”)을 올려준다. 물론 이 head는 완벽한 성능을 보여주는 것은 아니지만, 실험적으로 약 65%의 경우에 정답에 해당하는 예시들을 보고 있다는 것을 알아내었다.
[추상적인 패턴을 학습한 induction head에는 무슨 일이 일어나고 있는 걸까?]
앞서 논의된 attention head들은 우리가 앞서 정의가 된 induction head의 정의에 부합함과 동시에 추가적으로 예시에 나오는 추상적인 패턴 또한 보이고 있다는 것을 한 번 더 강조한다. 이는 비유적인 표현이거나 정의를 흐리는 것이 아니라 말 그대로 두 가지의 행동 패턴을 모두 가지고 있다는 것인다.
이는 궁금증 하나를 자아낸다: 무작위의 텍스트를 던져줘도 반복되는 부분들을 찾아내 복사하는 head들은 왜 이러한 추가적인 행동을 보이는가? 한 가지 힌트가 될 수 있는 부분은 이러한 추가적인 행동이 사실 ‘어떻게 보면 복사와 비슷하다’고 볼 수 있다는 관찰에서 온다. 원래의 정의에서는 [A][B] … [A] → [B]
과 같은 형태를 기준삼았지만, 이번 증거 4에서 보여주는 행동들 또한 [A*][B*]…[A][B]
이라는, 다소 비슷한 패턴에 기반하고 있다. 특히, [A][B] … [A] → [B]
는 [A*][B*]…[A][B]
의 특별한 케이스라고 해석하면, induction head는 사실 더욱 일반적인 알고리즘을 구현한 것이라고 해석할 수 있다. 또 다른 가능성으로, induction head만 따로 residual stream에 미치는 영향만 따로 떼어놓고 보면 기존의 복사 + 붙여넣기를 행하지만, 앞선 layer들 중 어떤 추상적인 representation이 계산된다고 가정했을 때 이를 induction head가 처리하는 과정에서 더욱 복잡한 패턴들이 나타나는 것일 수 있다.
증거 5에서는 induction head가 작동하는 방식이 더 일반적인 패턴 매칭으로도 확장될 수 있다는 것을 보여줌으로서 해당 추측에 힘을 실어줄 것이다. 이번 증거 4에서 이야기 하고 싶은 것은 단순히 말해 이러한 induction head들이 추상적인 패턴에도 관여 하는 것이 가능하다는 것을 보여주는 것이다.
증거 5 (Induction head의 mechanistic 타당성): 작은 모델에서 induction head가 mechanistic하게 설명이 됨을 보이고, 더 일반적인 모델에서도 가능케 하는 자연스러운 확장이 있음을 보인다.
우리가 induction head에 대한 탐구를 진행하기 위한 원래의 목적에는, induction head의 이해가 결국 in-context learning의 이해로 이어질 것이라는 생각이 포함되어있다. 하지만 그 반대의 방향도 생각해볼 수 있다: induction head에 대해 우리가 얻은 지식으로 in-context learning에 induction head가 기여를 하고 있다는 순수하게 논리적인 의견을 제시하는 것이다.
일단 semi-empirical한 주장부터 시작하도록 하자. Induction head가 지금까지 우리가 보아왔던 대로 행동한다는 것을 전제로 하면, 이는 모델이 다음에 올 token들을 예측하는데 도움을 줄 것이라고 예상할 수 있다. 이것은 결국 이전 context (이전 token)들을 일종의 data point라고 생각했을 때 nearest neighbor 알고리즘들 돌리는 것이라고 생각할 수 있고, data point가 추가가 될수록 nearest neighbor의 성능도 향상될 것이다. 따라서, induction head가 존재한다면, in-context learning에도 기여를 할 것이다.
어떤 의미에서는 이런 주장은 induction head가 in-context learning에 기여하는 ‘예시들이 존재한다’라는 범위 안에서는 꽤나 강력한 주장이 된다 (이전 예시에서도 봤듯이 그러한 경우들은 실제로 존재한다). 또한, 우리가 행동적인 측면에서 정의를 내렸기에 이러한 경향성이 일반적일 가능성이 높다. 이런 주장은 in-context learning에 induction head가 얼만큼 기여하는지에 대해서는 말할 수 있는 바가 없지만, (모델 크기와 상관없이) 최소한 어느 정도는 기여하고 있다는 것은 거의 확실하게 말할 수 있다.
이러한 논리적 접근을 했을 때 가장 만족스러운 부분-즉, 우리의 induction head에 대한 이해를 바탕으로 in-context learning에 미치는 영향을 생각하는 것-은 우리가 실험적인 증거에 의존하지 않고도 같은 결론에 도달한다는 것이다 (물론 더욱 복잡한 논리를 펼쳐야 한다는 값을 지불해야 한다). 이전 글에서 우리는 induction head를 reverse engineer 할 수 있었다 (즉, 파라미터 수준에서 분석하였다). 이 분석을 그대로 가지고 온다면, 실제로 induction head를 돌리지 않고도 그들이 어떻게 행동하는지 예측을 할 수 있을 것이며, 이전 문단에서 세운 주장들을 펼칠 수 있다. 물론, 몇 가지 한계점들이 있다. 이전 글에서는, 작은 attention-only 모델에서 induction head 하나에 대해서만 reverse engineer를 진행하였다 (물론 더 큰 모델에서도 진행하였다). 더 큰 문제점은 MLP가 포함되어있는 모델에서는 reverse engineer가 불가능하다는 것이다. 하지만 적어도 우리가 살펴본 몇 가지 경우에 대해서는, 프로그래머가 소스 코드를 보고 알고리즘을 파악하는 것과 비슷하게 transformer의 파라미터들을 보고 induction head를 알아차릴 수 있다.
사실 이전 글에서 reverse engineer했던 2-layer attention only 모델에 대해서는 이 주장을 조금 더 밀고 나갈 수 있다. Induction head가 어떻게 in-context learning에 기여하는지 알고 있을 뿐만 아니라, 사실 in-context learning을 일으킬만한 다른 메커니즘이 없어보이기 때문이다. 이는 induction head가 in-context learning의 주된 원동력이 된다는 주장에 더 힘을 실어주게 된다.
이제 우리는 induction head가 어떻게 더욱 추상적인 행동을 일으킬 수 있는지에 대해 주장을 펼쳐볼 것이다. 그 전에, 이전 글에서 다뤄졌던 induction head를 reverse engineering하면서 얻은 결과들을 간략하게 정리를 할 필요가 있어보인다. 만약 밑의 글에서 다루는 내용들이 처음이라면, 중간중간에 걸린 하이퍼링크를 읽지 않고 이해를 하는 것은 불가능할 것이라고 생각한다.
Induction head의 reverse engineering: 요약본
이전 글의 내용을 아주 밀도있게 요약한 것으로, 자세한 사항은 원 글을 읽는 것을 추천한다.
우리가 induction head를 복사하기와 prefix matching이라는 행동을 보이는 것으로 정의했다는 것을 기억하자.
복사하기는 OV (”Output-Value”) circuit에 의해 행해진다: Induction head의 주요 특징 중 하나이긴 하지만, 다른 head들도 복사를 할 수 있다! 구체적으로 말하자면, transformer 모델들은 복사를 할 수 있는 head들이 꽤나 있는 것으로 보이고, induction head는 그 중 일부이다. 이는 “복사 행렬”로 이루어진 OV circuit에 의해 만들어지고, 양의 eigenvalue들을 관찰하여 발견할 수 있다.
Prefix matching은 QK circuit의 K-composition (그리고 약간의 Q-composition)에 의해 행해진다: Prefix matching을 다시 생각해보면, attention이 앞선 context에서 보고 있는 token의 key 벡터가 그 token보다 앞서 나온 token들의 정보를 가지고 있는 구조여야 한다 [14]. 사실, attention이 앞선 context에서 보고 있는 그 token 자체는 induction head에서 필요한 attention 패턴을 계산할때는 개입을 하지 않는다 [15]. 우리가 다루고 있는 모델들에 대해서는 이런 “key shifting”이 K-composition이라는 원리에 의해 작동된다. 이는 induction head의 파라미터 중 가 이전 layer에서 attention head에 의해 쓰여진 subspace에서 정보를 읽는 것을 말한다. 가장 단순한 형태의 induction head는 순수 K-composition을 “previous token head”와 함께 QK circuit을 형성하고 있으며, 양의 eigenvalue를 가지는 를 가지고 형태가 된다. 이렇게 되면 induction head가 현재 보고 있는 token과 앞선 context에 나온 모든 token를 다 비교하면서 비슷한 부분들을 볼 수 있게 된다. 더 복잡한 형태의 QK circuit은 단순이 직전 token만 보는 것이 아닌 더 많은 token을 볼 수 있게 한다.
종합하자면, induction head는 이러한 메커니즘을 통해 직접 관찰할 수 있다. 이전 글에서 탐구에 사용된 작은 모델들에 대해서, 모든 induction head들은 앞서 설명된 QK circuit을 가지고 있고 양의 eigenvalue 특성 또한 지니고 있다.
어떤 모델들은 induction head를 다른 메커니즘을 통해 구현하기도 한다. GPT-2 [2]에서, 우리는 “pointer-arithmetic” 이라는 또 다른 메커니즘을 통해 induction head가 형성되는 증거들을 관찰할 수 있었다. 이 메커니즘은 positional embedding 및 Q-composition을 주로 사용한다는 특성이 있다. GPT-2에서, 앞선 layer의 attention head는 현재 보고 있는 token과 같은 token들을 찾아내고, 그 head의 OV circuit이 찾아낸 token의 positional embedding을 현재 보고 있는 token의 subspace로 복사한다. Induction head는 Q-composition을 이용해 positional embedding을 한 token 뒤로 돌리게 되고, 뒤따르는 token이 attend를 하게 된다. 이 메커니즘은 지금 보고 있는 모델들에서는 residual stream에 positional 정보를 더하지는 않기에 보이지 않는다.
더 복잡한 induction head에 대해서는 어떤가?
이전 증거 4에서 보았던 더 복잡한 induction head에 대해서는 어떤 논리를 전개할 수 있는가[16]? 여전히 같은 메커니즘을 사용하는가? 현재로서는, (그런 큰 모델에서) reverse engineering을 완벽하게 진행하는 것은 불가능하다. 아직 mechanistic interpretability 하에서 해석이 잘 되지 않는 MLP가 존재하기 때문이다. 하지만, 두 가지 정도의 가설은 세워 볼 수 있을 것이다:
(1) 단순히 직전 token을 보는 것 보다 복잡한 형태의 QK circuit을 사용할 수 있다.
(2) 단순히 token을 직접 복사하고 붙이기보다, (언어적으로) 더 추상적이고 복잡한 의미를 고려한 복사를 진행할 수 있다.
증거 6 (작은 모델에서 큰 모델으로의 연속성): In-context learning 밑 induction head와 관련한 패턴들과 데이터들이 작은 모델에서부터 큰 모델까지 연속적인 변화를 보이며, 따라서 큰 모델에서도 사실상 같은 방식으로 작동하고 있음을 제시한다.
이 증거는 이전까지 나온 증거들의 종합 및 그의 연장선이라고 볼 수 있다. 증거 1~5는 작은 transformer (특히 attention-only의 경우)에 대해서는 induction head가 in-context learning의 주 원인이 된다는 꽤나 강한 증거들을 제시할 수 있었다. 반면, 더 큰 모델들에 대해서는 증거들이 강하지는 않았다. 작은 모델들에서 일어나고 있는 일들이 큰 모델에서도 일어날 수 있다는 연결고리를 얼만큼 추론해 낼 수 있을까? 당연하게도 어느 정도 주관적인 기준을 적용할 수 밖에 없을 것이다.
모델 분석 테이블 [17]에서는 작은 attention-only 모델, MLP가 있는 모델, 그리고 풀스케일 모델 사이를 계속 비교한다. 레이어 갯수가 1개보다 많기만 하면, 모두 phase change를 겪게 된다. 모두 in-context learning 성능이 특정 구간에서 빠르게 올라가는 모습이 보이고, 성능 상승 폭 또한 비슷하다. 모두 PCA를 적용했을 때 모델이 겪는 경로가 비슷하다. 모두 induction head를 형성한다.
(작은 모델과 큰 모델에서 induction head와 관련하여 딱히 유의미한 차이점이 관찰이 되지 않기 때문에) 작은 모델에서 큰 모델로 분석 대상을 바꿀 때 (근본적으로) 바뀌는 것이 있다면, 무엇이 바뀌는가? 또한 왜 우리의 관측에서는 잘 나타나지 않는 것인가?
조금 다른 면들을 보자면, 큰 모델이 작은 모델과 다르게 행동하는 다른 예시들은 많다 (자세한 사항은 related work 참조). 그렇기에 작은 모델에서 얻은 결론을 큰 모델로 확장하는 것은 상당한 주의를 요한다.
가장 설득력이 있는 또 다른 가능성을 생각하자면, 또 다른 composition 메커니즘 또한 phase change 때 형성되는 것이다. 큰 모델은 head 갯수도 더 많고, 따라서 작은 모델에서는 불가능한 Q 혹은 K-composition 메커니즘이 생성이 될 가능성도 충분하다. 만약 모든 “composition에 의해 생성된 head”가 phase change때 남김없이 형성 된다고 하면, 어느 정도의 모델 크기 이상부터는 induction head에 관여하지 않는 head들이 오히려 phase change에 발생하는 모델의 변화에 더 큰 기여를 할 수 있다.
아직 규명되지 않은 부분들에 대하여
모든 과학적 탐구가 그렇듯, 이번 연구에서 규명하지 못한 부분들이 존재했고 이에 대해 다루도록 한다. 또한, 해결하지 못했지만 초기적인 실험 결과 중 흥미로운 부분이 있어 이 또한 소개한다.
In-context learning score가 유지되는 현상
우리가 설명할 수 없었던 첫 번째 현상으로는 모든 모델들이 in-context learning score가 어느 정도 다 비슷하다는 것이 있다. 즉, 모델이 고작 2 layer밖에 없던, 13B짜리 모델이던, 심지어 phase change가 끝난 직후이거나 그 이후 training이 꽤 오래 진행되었던 다 비슷한 정도의 score를 보인다는 것이다. 다른 말로 표현하자면, phase change를 겪었는가 아직 겪기 전인가 만이 in-context learning score에 영향을 미친다는 것이다.
자연스럽게 이 현상이 우리가 in-context learning score를 정의한 방식 (즉, 50번째와 500번째 token을 비교)에 문제가 있는가에 의문을 제기할 수 있다. 하지만 추가적으로 실험을 진행한 결과 이에 큰 문제가 없는 것으로 보인다.
위에 보여지는 plot에는 in-context learning score의 정의를 각 token과 마지막 token (8192-th position)과의 loss 차이로 변경한 이후 다시 측정한 결과가 정리되어있다. 모델 사이즈 별로 조금씩 다른 차이가 존재하기는 하지만 — 어쩔 때는 작은 모델이 in-context learning을 더 잘 할 때도 있다! — 전반적으로는 그 차이가 미미한 수준인 것으로 보인다.
어떻게 이러한 현상이 발생하는 것일까? 일단, 큰 모델이 여전히 작은 모델에 비해 다음에 올 token을 상대적으로 잘 예측한다는 것은 변함이 없다 (마지막 token으로 갈수록 더 잘한다). 자세히 보면 작은 모델에 비해 score를 context index 앞 부분에서 이미 앞서나가고 있는 것이 큰 작용을 하고 있다. 사실 대부분의 score 차이는 이 앞 부분에서 (첫 10개의 token) 발생한다.
더 큰 모델들은 주어진 문장 초반에 더 많은 정보를 파악할 수 있는 것으로 보인다 (이는 더 큰 모델들이 world knowledge가 더 많은 덕에 context의 도움 없이도 내용을 파악하는 것일수도 있다). 그 뒤로는 얼추 비슷한 loss drop을 context 끝까지 유지하게 된다. Loss가 떨어지는 양이 비슷한 현상으로는 어쩌면 큰 모델의 입장에서 더 어렵기 때문에 일어나는 현상일수도 있다 (이미 더 낮은 loss 값으로부터 시작했으므로). 왜 모델들이 비슷한 in-context learning score을 가져야 하는지는 명확하지 않은 부분이 있다.
Loss derivative에 phase change가 미치는 영향
또 다른 재미있는 현상은 학습 도중 loss가 떨어지는 정도를 모델 사이즈 별로 비교하면 phase change가 일어나는 시점에서 모델 별 loss가 감소하는 속도가 역전된다는 것이다. 이는 를 관찰하면 손쉽게 관찰할 수 있다. 핵심은 phase change에서만 큰 모델이 loss가 더 빠르게 감소하고, 나머지 부분에서는 작은 모델이 더욱 빠르게 감소한다.
학습 초기 단계에서는 작은 모델들이 더 빠르게 training loss를 떨어뜨린다는 것이 딱히 놀랍지는 않으나, 역전이 일어나는 부분과 phase change가 일어나는 부분이 일치한다는 점은 놀랍다고 할 수 있다.
남아있는 궁금증
모델 분석 table에서:
6-layer attention-only 모델에서 학습이 절반 이상 경과하였을 때 특이한 head가 학습된다. 이 head는 induction head는 아니지만, 이 head에 대한 ablation을 진행하면 phase change의 효과를 무마하는 듯한 현상이 나타난다. 이 head는 도대체 무엇인가?
4-layer MLP 모델에 대해서만 다른 모델의 실험결과와 비교했을 때 경향이 크게 두드러지지 않는다. 이 모델만 다른 무언가가 있는가?
6-layer MLP 모델은 ‘loss spike’가 일어나지만, 그 원인이 무엇인지는 알 수 없다.
6-layer MLP 모델은 induction head가 하나만 존재하며 이 head를 ablation하면 in-context learning score에서 반대의 현상이 일어난다. 이 head는 도대체 무엇인가?
Appendix에서:
Full-scale 모델의 16 layer 이상부터 몇몇 head는 “prefix search”에는 좋은 성능을 보이지만, 복사에는 negative score를 보인다. Induction head도 아닌 이 “안티-복사 prefix-search” head으로부터 어떤 정보를 얻어낼 수 있을까?
결론
안전성 (safety)에 대하여
이 연구의 가장 최종적인 목표는 neural network를 reverse engineering 한다면 이러한 모델을 사용하는 데에서 안전성에 대한 확신을 가질 수 있다는 것이다. 이번 연구는 그 방향성으로 가는 초기 단계이지만, 몇몇 안정성과 관련된 이슈들에 대해 다음과 같은 논의를 할 수 있다:
Phase change: 모델이 스케일이 변함에 따라 급작스럽게 변화하는 양상이 존재한다면, 미래에 문제점들이 발생한다면 연구자들 및 사회 전반적으로 이에 대해 대비하는 것이 까다로워질 것이다.
In-context learning: In-context learning은 안전성 커뮤니티 안에서 주의깊게 바라보고 있는 주제이다. 조금 작은 규모의 모델들에 대해서는 학습을 통해 어느 정도 통제가 가능할 것으로 생각할 수도 있다 (하지만 adversarial reprogramming은 이 가능성에 대해서도 의문을 표한다). In-context learning은 모델의 행동 양식이 (추가적인 학습 없이도) inference 단계에서 어느 정도 “변화”할 수 있음을 보여주는 현상이다. In-context learning이 이미 학습된 행동 양식을 “확인”하는 행위라고 해석을 한다 하더라도, 이러한 행동 양식은 의도하지 않은 off-distribution 상황에서의 일반화라고 볼 수 있다.
Mesa-optimization [18]: In-context learning이 어쩌면 mesa-optimization과도 연관이 있을 수 있다는 우려가 존재한다. Mesa-optimization이란, 학습된 모델이 내부적으로 (우리가 의도하지 않은) 어떤 optimization을 스스로 돌리고 있는 현상이다. 이 연구는 적어도 작은 모델에 한해서는 in-context learning의 원인이 induction head이며, mesa-optimization과는 거리가 멀다는 것을 보여준다.
Learning dynamics, scaling law, 그리고 mechanistic interpretability 사이의 관계에 대하여
In-context learning에서 보이는 phase change가 mechanistic interpretability와 learning dynamics, 그리고 neural network의 통계물리적인 특성 (예를 들어 scaling law 혹은 phase change)간의 관계를 규명할 초석이 될 수 있다. 만약 이러한 현상들의 관계를 연구하고 싶다면, phase change에 대한 규명을 하는 것이 좋은 출발점이 될 수 있다.
각주
[1] Mechanistic interpretability: Explainable AI (XAI)의 sub-area 중 하나로 비교적 최근에 떠오른 분야이다. (모델 작동의) 상당히 면밀한 관찰을 바탕으로 딥러닝 모델들을 해석하고자 하며, 모델을 ‘reverse-engineering’하고자 하는 움직임으로도 많이 불린다. Case-by-case의 성격이 강하며 주로 transformer 모델들이 해석 대상이 된다.
[2] Radford et al., “Language Models are Unsupervised Multitask Learners”, 2019
[3] Brown et al., “Language models are few-shot learners”, 2020
[4] Kaplan et al., “Scaling laws for neural language models”, 2020
[5] 매우 포괄적인 용어로, 모델 안에 있는 (임의의) 두 뉴런이 상호작용을 하면 이를 circuit이라고 부른다. 모델 안에 circuit은 무수히 많으나, 그 중 유의미한 것들은 따로 이름이 붙는다 (정의 참조: Distill article). 단순하게 모델 안에서 일어나는 유의미한 계산이라고 보아도 문제 없다.
[6] Erhan et al., “Why does unsupervised pre-training help deep learning?”, 2010
[7] 1~6-layer transformer를 나타낸다. MLP가 포함 되어있을수도 있고 포함 안될수도 있으며, 따라서 모두 합치면 12개의 모델이 나온다 (출처: 원문 Model analysis table).
[8] https://transformer-circuits.pub/2022/in-context-learning-and-induction-heads/index.html#models
[9] 글 뒷쪽 SMALL ATTENTION-ONLY TRANSFORMERS (DIFFERENT DATASET) 참조
[10] 얼추 “얼룩진/오염된 key”라고 번역된다.
[11] 여기에서 사용된 이 정확하게 어떤 함수인지는 나오지 않는다. 대략적으로 sigmoid 함수 같긴 한데, 그렇다고 하기에는 범위가 이 아닌 이다. 일단 적당히 differentiable하고 인 함수 정도라고 생각해도 나머지 내용을 이해하는데 큰 어려움을 주지 않는다.
[12] 원문에서 수식은 등호를 사용한 이지만 기존 key vector인 를 바꾼다는 의미에서 로 바꾸었다.
[13] 모든 attention head에 대하여 ablation을 수행하는 것은 superlinear complexity인 의 cost가 든다 (은 모델 파라미터 갯수). Attention head가 만큼 있고 모델을 돌리려면 이 걸리기 때문이다. 추가적으로 10,000개의 샘플에 대해서 각 모델 체크포인트마다 ablation이 진행되고 있으니 cost가 비싼 편에 속한다.
[14] 앞서 논의되었듯이, [A][B] … [A] → [B]
에서 일단 현재 보고 있는 token (두 번째 [A]
)이 query로 들어오고, 원하는 위치의 token (첫 번째 [B]
)에서 key 값이 반응을 하기 위해서는 하나의 attention으로 부족하다. 왜냐하면 prefix matching의 정의 상 첫 번째 [B]
가 “적절한 token인지” 판단하기 위해서는 바로 앞에 (첫 번째)[A]
가 있는지에 대한 정보가 있어야하고, 이를 담당하는 또 다른 attention head가 필요하기 때문이다.
[15] 구체적으로 말하자면 QK circuit을 통해 attention 패턴을 계산할때만 바라보는 token이 무시된다. 이전 글에서 밝혀냈듯이, attention이 바라보는 곳을 계산하는 부분인 QK circuit과, 실제 output을 결정하는 OV circuit은 독립적으로 생각할 수 있고, 이런 접근 방식이 유용할 때가 많다.
[16] 원문에서는 “증거 2”라고 되어있으나 하이퍼링크는 “증거 4”로 향하고, 실제로 문맥상에서도 증거 4가 맞다.
[17] 모델 분석 테이블의 일부 내용은 번역되어 이 글에도 실려있으나, 전문은 번역하지 않고 원문 링크를 걸어둔다.
[18] Hubinger et al., Risks from Learned Optimization in Advanced Machine Learning Systems, arXiv 2019
[Ablation] Ablation study는 어떤 아키텍쳐를 이루고 있는 여러 부품 (혹은 모듈들)을 하나씩 (혹은 여러 개를 동시에) 제거하거나 무효화하면서 그 반응을 살펴보는 실험을 의미한다. ML 논문에서도 분석에 자주 쓰이는 툴이며, 직역하면 제거나 절개에 가깝다. 저자들이 ablation을 사용하는 증거 3이 제일 강력하다고 하는 이유가 causal inference에서 intervention (실험자가 직접 모델에 개입하는 방식) 기반의 실험이 주가 되기 때문이다 (따라서 증거 3 평가 table에서도 다 causal 증거라고 표기되어있다).
[Attention head ablation 방법론] 각 attention head의 역할을 알아낼 수 있는 방법 중 하나로는, 일단 온전한 모델의 output을 관찰한 뒤 해당 attention head만 빠져있는 모델의 output 차이를 보는 것이다. Attention head를 이렇게 ‘삭제’하기 위해 해당 attention head가 residual stream (참조)에 기여하는 정보를 0으로 대체한다. 이런 방법론은 생물학에서 특정 유전자를 제거하면서 그 역할을 알아내는 “knockout” 방식과 굉장히 유사하다. Ablation은 크게 “full” 과 “pattern 보존” 이렇게 두 가지 방식으로 진행할 수 있으며, 대부분의 ablation은 후자를 선택하였다.
Full ablation: 해당 attention head의 최종 결과물 (representation vector)를 0으로 대체한다. 이에 따라 그 뒤에 나오는 모든 계산에 영향을 미치게 된다.
Pattern 보존형 ablation: 두 단계로 나누어진다.
일단 아무런 모델 수정 없이 output을 계산한다. 그와 동시에 모델의 attention weight 값 (attention 패턴)을 저장한다. Attention head의 최종 output 값들은 저장을 하지 않는다.
모델을 한 번 더 돌린다. 이번에는 해당 attention head의 최종 결과물을 0으로 대체하면서 진행하되, attention 패턴들은 단계 1에서 저장했던 값들로 강제로 고정시킨다. 그 결과, 해당 attention head의 V에 대해서만 수정이 가해지고 모델에 존재하는 모든 Q, K 값들은 변화하는 것이 없다. 즉, 굉장히 local하고 타겟팅이 되는 ablation이 된다.
[Composition] 왜 적어도 두 개의 layer가 필요한 것일까? Induction head의 패턴인 [A][B] … [A] → [B]
을 살펴보자 (중복되는 token 사이의 구별을 쉽게 하기 위해 [A1][B1] … [A2] → [B2]
이라고 해보자). 현재 바라보고 있는 [A2]
token을 기준으로, induction head는 [B1]
를 찾아야 한다. 문제는 이 [B1]
token 위치는 [A1]
에 의해 결정된다는 것이다. 각 attention head는 두 개의 token pair 사이의 관계만 볼 수 있기 때문에, induction head가 [A2]
을 보고 [B1]
을 찾기 위해서는 또 다른 attention head가 [A1]
과 [B1]
사이의 관계를 알려줘야 한다.
[Confounding variable] 쉽게 말해, 인과적인 관계 (원인과 결과)가 없고 단순히 상관관계만 존재하는 두 사건들 사이에 존재하는 제 3의 사건이다. Brady Neal의 causal inference course에 나오는 예시를 보자. 어느 사람이 1) 아침에 일어났을 때 숙취를 경험하는 것과 2) 전날 신발도 벗지 않고 잠을 자는 행위 사이에 큰 상관관계가 있다는 것을 관찰했다고 하자 (즉, 전날에 신발을 벗고 자면 높은 확률로 다음날 숙취가 있다). 이 것을 곧대로 인과관계로 해석하면 다음과 같은 이상한 결론이 나온다: “신발을 신고 잠을 자는 행위는 다음날 숙취의 원인이 된다”. 실제로는 3) 전날 술을 진탕 마시는 행위라는 변량을 고려해야하며, 실제 인과 관계는 3 → 1 (술을 마셨기 때문에 다음날 숙취가 있음)과 3 → 2 (술을 마셨기 때문에 신발도 벗을 겨를이 없음)이며, 이러한 공통적인 변수 때문에 1과 2 사이에 큰 상관관계가 있는 것 처럼 보이는 것이다. 여기에서 3)이 confounding variable이 된다.
[“복사하기” 및 “prefix matching” 실험] 해당 실험은 특정 head가 다음과 같은 성질을 얼만큼 보이는지를 activation을 통해서 확인하기 위해 고안이 되었다. 관심 대상이 되는 현상들은 다음과 같다:
복사하기: 현재 head가 residual stram에 미치는 영향이 현재 보고 있는 token의 logit 값에 직접적인 영향을 미치는가? 관측 방법) 길이가 25인 token sequence를 무작위로 만들어낸다 (가장 흔하거나 희귀한 token 제외). 해당 head가 residual stream에 기여하는 representation을 구하고, 그대로 unembedding matrix를 통과시켜 logit에 미치는 영향을 구한다 (즉, 나머지 attention layer들을 무시하고 residual stream을 쭉 따라서 output까지 통과시킨다). 이 값과 logit의 평균값과의 차이를 계산하고 ReLU를 통과시켜 logit값이 증가한 부분에만 분석을 집중시킨다. (역자: 원문으로는 정확히 어떻게 분석이 이루어지는지 알기 어려웠다)
Prefix matching: 반복되는 랜덤 token sequence가 주어졌을 때, 현재 보고 있는 token과 같은 token이 직전에 등장하는 이전 token에 attention이 가해지고 있는가? 관측 방법) 길이가 25인 token sequence를 무작위로 만들어낸다 (가장 흔하거나 희귀한 token 제외). 이 sequence를 4번 반복하고 맨 앞에
[start of sequence (SOS)]
token을 붙인다. 그리고 attention 패턴을 구한다. Prefix matching score는 현재 보고 있는 token과 같은 token이 뒤따라서 오는 모든 token들에 해당하는 attention 값들의 평균으로 정의한다.Previous token attention: 현재 보고 있는 token 직전의 token에 attention이 가해지고 있는가? 관측 방법) Training 데이터셋에서 샘플을 하나 고른다. 그리고 attention 패턴을 구한다. Previous token score는 각 position 마다 번째 token에 대한 attention 값들의 평균으로 정의한다.
[부연 설명] (원문에도 각주가 있으나 이를 기반으로 부연 설명을 더한다) 수학적으로, 연속 함수를 무한 차원을 가지는 벡터로 표현할 수 있다. 예를 들어보자. 이라는 함수를 벡터로 어떻게 나타낼 수 있을까? 어찌 되었건, 기능적으로 보면 함수는 어떤 input (여기서는 )이 들어왔을 때 어떻게 반응하는지 (즉, )를 특정지으면 된다고 생각할 수 있다. 이를 정말 단순하게 이용해서 를 표현하는 다음과 같은 벡터를 생각해 볼 수 있다: . 즉, 실수 (real number)에서 몇 가지 input들을 선택한 뒤 (여기서는 0,1,2,3) 이에 대한 함숫값을 모아서 벡터로 만드는 것이다. 물론 이건 아주 조악한 표현 방식이다. 실제 정의역인 무한히 많은 실수들이 아닌, 그 중에서 고작 4개만 골라서 표현한다는 분명한 한계가 있다. 즉, 더욱 정확하게 표현하기 위해서는 무한히 많은 실수들을 다 끌고 와서 함수의 반응을 다 봐야하고, 결론적으로 벡터의 크기가 무한히 커지게 된다. 그러한 의미에서 무한 차원의 벡터로 함수를 표현할 수 있다는 것이다. 이러한 아이디어는, 벡터끼리의 내적을 어떻게 하는지 생각해본 다음 함수끼리의 내적에서 왜 적분이 등장하는지 생각해보면 또 납득이 갈 것이다.
다시 원 글로 돌아와서, 저자들의 의도는 training 중 여러 개의 transformer 모델들을 비교하고 싶은것이다. 구체적으로는 각 (transformer 모델 + loss function)끼리의 비교인 것이다. Loss function이 윗 예시의 와 크게 다르지 않다는 것을 이해한다면 이미 거의 다 이해한 것이나 마찬가지이다. 당연히, 실수 전체를 가져와서 진짜로 무한 차원 벡터를 만들 수는 없기에 윗 예시에서 4개들의 값 (0,1,2,3)들에 대해서 표현했듯, 저자들도 같은 맥락으로 몇 개의 (무작위로 골라진 10,000개의) input 샘플들을 기준으로 비교하고 싶은 것이다 (그림 참조). 여기서 말하는 ‘무작위로 골라진’은, 10,000개의 길이 512짜리 샘플 각각에서 하나의 token만을 선택하는 식으로 구성하는 것을 말한다 (즉, 각 샘플 당 나머지 511개는 버린다). 이렇게 모아 만든 10,000 차원짜리 벡터가 ( 와 같이) transformer 함수 그 자체를 표현한다고 생각한다는 것이다.
마지막으로, 이렇게 되면 왜 굳이 PCA를 하고 싶은지도 이해가 된다. 모델 하나를 표현하는데 10,000 차원 짜리 벡터가 사용되었으니, 이를 (visualize 등을 포함하여) 손쉽게 분석하려면 10,000차원 → 2차원으로 줄이는 편이 자연스럽다. 그 중에서 가장 대표적인 방법이 PCA이기에 이러한 맥락으로 사용되는 것이다.
[주의1] Induction head가 하나의 attention head를 지칭하는지 (2-layer transformer의 2번째 layer에 속하는 attention head), 아니면 이를 포함하는 circuit 전체를 지칭하는지에 대해 글 전반적으로 하나로 통일되지 않는 모습을 가끔 보인다. 저자들도 이를 인지하고 있으며 뒷 부분에서 조금 더 구체적인 정의가 나오기 때문에 번역에서도 이를 굳이 구분하지 않고 원문을 따라가는 것을 지침으로 한다 (물론 더 구제척인 정의에서도 이 모호성은 완벽히 해결되지는 않는다). 특별한 언급이 없는 한 “induction head”는 1-layer에 있는 previous token head와 2-layer에 있는 head가 같이 상호작용하면서 만들어지는 circuit을 지칭한다.
[구체적 데이터 처리 방법1] 해당 내용은 글 아랫 부분의 Analyses - Per-token analysis - Loss curve에 해당한다. (N개 모델 snapshot X 10,000개의 token sequence별 loss) matrix를 만들고, token sequence를 기준으로 평균을 내어 (N X 1) 짜리 데이터를 만든다. 작은 모델은 N=200, 큰 모델은 N=15이다. 즉 학습이 진행되면서 평균적으로 모델들이 보이는 loss를 산출하는 셈이다.
[구체적 데이터 처리 방법2] 해당 내용은 글 아랫 부분의 Analyses - Per-token analysis - 2D in-context learning plot에 해당한다. (N개 모델 snapshot X 512 token index 위치) matrix를 만들고 (참고: 각 token sequence 샘플의 길이가 512이다) 이를 그대로 2D figure로 그린다.
Last updated