A Mathematical Framework for Transformer Circuits 이해하기

PreviousIntroduction to Mechanistic Interpretability (기계적 해석의 전반적인 설명)NextIn context learning and Induction heads 이해하기

Last updated 7 months ago

📝 Original authors: Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, 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 제외), (잘 알려진) transformer 구조를 재해석하여 mechanistic interpretability에서 바라보는 transformer를 정리한다. 내용이 워낙 방대하여 번역을 하는 동시에 핵심적인 부분 위주로 추려내려고 노력하였다. 뒷 부분의 Related works나 appendix등은 원문을 확인하기 바란다.

원 글: A Mathematical Framework for Transformer Circuits

원 번역 글 Notion link: A Mathematical Framework for Transformer Circuits (요약 번역) - coming soon

Introduction

Transformer기반 언어 모델들이 강력해짐에 따라 세계적으로 그 쓰임새 또한 증가하고 있다. 하지만 워낙 모델의 크기가 크고 복잡해지다보니, 모델을 만든 사람들과 사용하는 사람들 모두 출시가 된 지 몇 년이 지난 후에도 전에 발견하지 못했던 새로운 현상들을 발견하고는 한다. Mechanistic interpretatilbity는 이러한 예측 불가능성을 탐구하는 접근법 중 하나로, transformer 모델 안에서 일어나는 복잡한 계산들을 reverse engineer (역설계/리버싱)하고자 한다. 비유하자면, 인간이 읽기 힘든 이진 코드를 읽기 쉬운 소스 코드로 변환하는 과정이라고 생각해볼 수 있다. 이것이 가능하다면, 이러한 transformer 모델들을 배포하는 과정에서 발생할 수 있는 다양한 안전성 문제를 포착/설명/해석하거나 심지어 사전에 예상할 수 있을 것으로 생각한다. 이와 비슷한 시도로는 이전에 vision 모델을 기준으로 작성된 Distill circuits thread가 있으나, transformer 및 언어모델에 대해서 이런 접근이 이루어진 적은 없었다.

이 글에서는 transformer 모델에 대한 mechanistic interpretability의 초석을 놓는 것을 목표로 잡았다. 구체적으로는, 작은 모델에서 알고리즘적인 패턴이나 framework 등을 발견하는 것이 목표이며, 나중에는 발견한 내용들을 점점 더 복잡한 모델들에 적용할 수 있으리라 기대한다. 이런 의미에서, 이 글에서는 2 layer 이하의 (MLP 블록이 없는) attention-only transformer 모델에 대해서만 이야기 할 것이다.

우리는 transformer 모델을 수학적/기능적으로는 동일하지만 새로운 방식으로 표현하게 됨으로서 이런 작은 모델들에 대한 이해도를 매우 높일 수 있다는 것을 발견하였다. 특히, 우리는 attention head들이 소위 “induction head”이라는 개념을 통해 in-context learning을 발현한다는 것을 발견하였고, 이는 적어도 attention layer를 두 개를 쌓아야 발현된다는 것 또한 알아내었다.

결론 요약

Reverse engineering으로부터 얻은 것들

우리는 앞서 언급했던 작은 transformer 모델들을 뜯어보았고, 다음과 같은 결론들을 얻을 수 있었다.

0-layer transformer 모델들은 bigram 통계를 모델링한다. 그리고 이 bigram table은 모델 weight으로부터 직접적으로 추출이 가능하다.
1-layer attention-only transformer은 bigram + “skip-trigram (A…BC의 패턴)”의 앙상블이다. 그리고 이 bigram 및 skip-trigram table들 또한 모델 weight으로부터 직접적으로 추출이 가능하다. 이 skip-trigram들은 의외로 expressive하며, 아주 간단한 형태의 in-context learning을 구현하기도 한다.
2-layer attention-only transformer 모델은 attention head들을 조합 (composition)하게 됨으로써 더욱 복잡한 알고리즘들을 구현할 수 있게 된다. 그리고 이런 알고리즘 또한 모델 weight들을 살펴봄으로서 관찰이 가능하다. 특히 in-context learning 알고리즘을 구현하는 “induction head”를 구현한다.
1-layer 및 2-layer attention-only transformer들은 in-context learning을 구현하기 위해 서로 다른 알고리즘을 사용한다. 2-layer의 경우 inference시 상대적으로 더욱 정교한 알고리즘을 사용하며 —induction head라는 특별한 attention head를 이용하게 된다—, 더 큰 모델들로 확장할 수 있는 연결고리를 제공한다.

모델 재해석을 통해 얻은 것들

Transformer 모델이 이전 Distill circuits thread에서 다룬 InceptionV1과는 다른 방식으로 짜여진 만큼, 수학적으로 재해석을 하면서 새로운 것들을 알 수 있는 기회가 되었다. 추후에 자세히 다룰 예정이지만, 일단 여기에 그 내용을 압축하여 정리한다. 각 섹션에 가면 새롭게 정의한 용어들에 대한 구체적인 설명이 있을 것이며, 우리가 이 곳에 제시한 개념들이 완전히 novel하다는 주장을 하는 것이 아니다 (다른 논문들이 implicit하거나 explicit하게 많이 다루었다).

Attention head는 모델의 독립적인 기능으로 볼 수 있으며, 각각 계산하는 결과물은 residual stream에 더해진다. Attention head은 “concatenate 이후 곱한다”라는 맥락으로 설명이 되기는 하나, 수학적으로는 이렇게 해석해도 동일하다.
Attention-only 모델들은 독립적인 end-to-end (token으로부터 logit까지 연결해주는) 함수들의 합으로 나타낼 수 있다. 이 함수들은 모델 안에서 일종의 “경로 (path)”으로 생각할 수 있으며, attention 부분을 고정시키면 linear 함수가 된다.
Transformer들은 (생각보다) 엄청난 양의 linear한 구조를 가지고 있다. 이는 모델 안의 행렬 곱들과 그들의 합을 뜯어보면 알게 된다.
Attention head들은 크게 두 가지의 독립적인 계산으로 이루어져있다고 생각할 수 있다: Attention pattern을 담당하는 QK (”query-key”) circuit과 각 token이 output에 얼만큼 영향을 미치는지 계산하는 OV (”output-value”) circuit 이 두가지이다.
Key, query, 그리고 value 벡터들은 low-rank인 $W_Q^\top W_K$ 와 $W_O W_V$ 를 계산을 하는 도중 발생하는 중간 결과물이라고 생각할 수 있다.
Attention head들을 조합하면 transformer의 expressivity를 크게 향상시킬 수 있다. 크게 key, query, value에 해당하는 세 가지 방식으로 조합을 할 수 있으며, key & query 조합은 value 조합과 매우 다르다.
Token embedding, attention head, MLP layers, 그리고 unembedding을 포함한 transformer의 모든 부분은 residual stream의 다른 부분에 쓰기/읽기를 하는 방식으로 소통하게 된다. Residual stream 벡터를 직접 분석하기보다, 이런 “소통 창구”에 해당하는 모델 안의 다양한 경로를 파악하는 것이 더욱 효과적이다.

Transformer 모델 재해석

시작하기 전에…

Transformer 모델에 대한 설명을 하기 이전에, 우리가 이 모델에 대해 어떻게 생각하는지 잠깐 정리해보는 시간을 가져보는게 좋다고 생각한다.

많은 경우에, 우리는 transformer 모델을 수학적으로는 동일하지만 흔하지 않은 방식으로 표현하는 것이 해석적인 측면에서 도움이 된다는 것을 발견했다. Mechanistic interpretability에서는 모델을 인간이 이해할 수 있는 정도의 작은 단위로 쪼개는 것이 중요하고, 그러다보면 (모델을 효율적으로 구현하는데 초점이 맞춰져있는) 대중적인 모델 description과 달라질 확률이 높다.

또한 재해석을 하면서 많은 notation들과 개념들을 소개하게 될텐데, 구체적인 정의는 notation appendix을 참고하면 되겠다.

단순화 작업

이 글에서 다루게 될 개념들을 투명하게 전달하기 위해, 우리는 단순화 작업을 거친 “toy transformer” 모델에 집중한다.

이 글의 대부분의 경우, 우리는 attention-only transformer에 집중하도록 하겠다. 이 가정은 꽤나 극단적인 단순화이며, 이런 가정을 깔고 가는 이유에는 attention mechanism에 집중할 수 있다는 점도 있지만 더 중요하게는 우리가 아직 MLP layer들에 대한 이해가 부족하다는 것도 있다. 이는 이 글의 큰 한계점 중 하나이며 추후에 다룰 수 있기를 희망한다. 그럼에도 불구하고 MLP layer가 있는 모델에 대한 논의 또한 후반부에 넣었다.

명료함과 단순함을 위해 몇 가지 추가적인 수정 또한 가한다. 우리가 다룰 모델에서는 (마치 linear model에서 bias를 input 벡터에 차원을 하나 추가해주는 방식으로 무시할 수 있는 것 처럼) bias를 무시한다. 또한 scaling에만 영향을 미치는 layer norm도 복잡도만 증가시킨다고 판단하여 무시한다.

High-level architecture

Transformer 모델도 다양한 variant가 있지만, 우리는 GPT-3와 같이 autoregressive하고 decoder만 있는 transformer (언어) 모델에 집중한다 (원 transformer 논문 [1]은 언어 번역에 염두를 둔 encoder-decoder 형태를 가지고 있으나, 많은 현대 언어 모델을은 그렇지 않다).

Transformer에서의 계산은 token embedding에서 시작하고, 몇 개의 연속적인 “residual blocks”들로 이어지게 되며, 마지막으로는 token unembedding으로 끝나게 된다. 각 residual block은 attention layer 하나에 MLP 레이어 하나가 뒤따르는 형태로 구성된다. Attention과 MLP 레이어 모두 residual stream으로부터 (정보를) “읽고” (이는 linear projection을 통해 이루어진다), 정보를 처리한 결과를 (다시 linear projection을 써서) residual stream에 “쓴다”. 각 attention layer는 동시에 돌릴 수 있는 여러 개의 attention head이 존재한다.

커뮤니케이션 채널으로써의 virtual weight과 residual stream

Transformer 구조의 가장 큰 특징 중 하나는 각 레이어들의 구성품들이 내뱉는 결과물들이 “residual stream”에 더해진다는 것이다. Residual stream은 단순히 생각하면 이전 레이어들의 모든 결과물, 그리고 원래의 embedding 결과들을 모두 합한 결과물이라고 보면 된다. 보통 residual stream은 그 자체로 어떤 계산을 진행하지 않기 때문에, 이를 레이어들끼리 상호작용하는 일종의 ‘커뮤니케이션 채널’로 생각한다.

Residual stream은 굉장히 선형적인 구조를 지니고 있다 [2]. 모든 레이어는 embedding 직후 (시작 부근)부터 residual stream에서 “정보를 읽어내기” 위해 어떤 linear transformation (선형변환)을 적용하며, 그 결과 새롭게 생성한 정보를 “반영하기(쓰기) 위해” residual stream에 더하기 전에 다시 linear transformation을 적용한다. 이 linear transformation은 여러 중요한 역할들을 수행한다. 대표적으로, residual stream은 “privileged basis”이 없기 때문에, 이와 상호작용을 하는 모든 행렬들에 대해 rotation을 먹여도 모델은 실질적으로 바뀌지 않는다 [추가 설명1].

Virtual weights

Residual stream이 선형적이라는 사실이 특별이 유용할 때에는 우리가 임의의 두 레이어 사이에 대한 연결 관계를 “virtual(가상) weight”을 통해 표현할 때이다. 이 virtual weight들은 한 쪽 레이어의 output weight과 나머지 한 쪽 레이어의 input weight을 곱한 형식으로 생각할 수 있다 (예를 들어, $W^{2}_{input} W^{1}_{output}$ ). 이를 통해 뒷단에 나오는 레이어가 앞쪽 레이어에서 생성된 정보를 어떻게 읽는지를 파악할 수 있다.

Subspace와 residual stream의 bandwidth

Residual stream은 고차원의 vector space이다 [3]. 작은 모델에서는 수백 차원일수 있고, 큰 모델에서는 몇 천, 혹은 그 이상까지 올라갈수도 있다. 이렇게 되면 레이어들은 subspace를 통해 서로 다른 정보들을 각자 서로 다른 레이어에 보낼 수 있다는 것을 의미하게 된다 [추가 설명2]. 이 성질은 attention head에게 있어서 상당히 중요한데, 이는 각 attention head가 상대적으로 작은 subspace (보통 64에서 128 차원 정도)에서만 작용하고 residual stream의 일부 차원에 독립적으로 정보를 읽고 쓰는 것이 쉽기 때문이다.

한 번 정보가 residual stream의 어떤 subspace에 더해지면 다른 레이어가 이를 의도적으로 삭제하지 않는 한 그대로 보존이 된다. 이렇게 보면 슬슬 residual stream의 차원이 일종의 “메모리” 혹은 “bandwidth (대역폭)” 처럼 여겨진다. Token embedding 및 unembedding 행렬은 이 residual stream의 차원 중 상대적으로 일부만 사용하고 있는 것으로 확인되었다. 이렇게 되면 나머지 차원들은 다른 레이어들이 정보들을 가공하면서 이를 저장할 수 있는 “여유 공간” 처럼 되는 것이다.

이러한 residual stream의 넓은 bandwidth는 모델 안에서 수요가 높아 보인다. 특히, attention head나 MLP 레이어 등에서 결과물로 내뱉는 벡터들의 차원들의 합이 residual stream의 차원보다 훨씬 더 많은 것을 생각하면 더욱 그렇다. 즉, 25~50 레이어짜리 transformer 안의 특정 위치에서 residual stream을 보면 residual stream의 차원의 약 100배 정도 되는 수의 뉴런들이 그 전에 존재하고, 그 뒤로도 100배 정도 되는 수의 뉴런들이 존재하며, 이들이 서로 superposition [4]을 하면서 소통하고 있는 것이다! 이러한 종류의 텐서를 “bottleneck (병목) activations”라고 하며 해석이 유난히 더 까다로울 것으로 생각된다. (이것이 residual stream을 직접적으로 해석하기 보다는 다른 레이어들이 이를 통해서 어떻게 소통하는지에 집중하기로 한 큰 이유 중 하나이다.)

이러한 이유 때문인지 우리는 몇몇 MLP 뉴런들과 attention head들이 일종의 “메모리 관리” 역할을 하는 듯한 모습을 관측하기도 하였다. 즉 다른 레이어에 의해 사용된 residual stream의 차원들을 반대되는 부호의 정보를 다시 쓰는 것을 통해 지우는 것이다 [5].

Attention Head는 독립적으로 더해진다.

위에서 본 대로, 우리는 transformer의 attention layer를 완전히 독립적으로 작동하는 attention head $h \in H$ 들이 각자의 결과물을 residual stream에 더하는 (쓰는) 것으로 이해하고 있다. 하지만 보통 transformer를 사람들이 이야기 할 때는 이런 식의 설명이 흔하지 않고, 사실 (수학적으로) 같다는 것 또한 짚고 넘어가볼 만 하다.

원래의 Vaswani 논문 [1]에서는, attention layer의 output은 결과로 내뱉는 벡터들 $r^{h_1}, r^{h_2}, \cdots$ 들을 쌓고 output matrix $W^H_O$ 를 곱하는 형식으로 제시하였다. 여기서 행렬 $W^H_O$ 을 (attention head의 차원에 맞춰) 균등하게 잘라서 $[W^{h_1}_O, W^{h_2}_O, \cdots]$ 으로 생각해보자. 그러면 같은 계산을 다음과 같이 다시 쓸 수 있다:

W^H_O \begin{bmatrix} r^{h_1}\\ r^{h_2}\\ \cdots \end{bmatrix} = \begin{bmatrix} W^{h_1}_O,W^{h_2}_O,\cdots \end{bmatrix} \cdot \begin{bmatrix} r^{h_1}\\ r^{h_2}\\ \cdots \end{bmatrix} = \sum_i W^{h_i}_O r^{h_i}.

마지막 결과를 보면, 결국 원래 논문에서 제시한 방식과, 우리의 해석(즉, attention head가 독립적으로 작동하고 마지막에 다 한꺼번에 합치는 해석)이 동일하다는 것을 알 수 있다. 원래 논문과 같이 concatenation을 하는 버전으로 아키텍쳐를 설명하는 것이 유용한 것은 더 효율적인 구현을 하기 편하기 때문인 것도 있다. 하지만 이론적으로 transformer를 이해하는데는 독립적으로 더해지는 우리의 해석이 더 편리하다.

정보 교환의 수단으로써 attention head를 이해하기

하지만 attention head들이 독립적으로 작동한다면, 과연 어떤 역할을 수행하고 있는 것일까? Attention head의 가장 기본적인 역할은 정보를 움직이는 것이다. 한 token으로부터 시작한 residual stream으로부터 정보를 읽고 다른 token의 residual stream으로 정보를 쓰는 것이다. 여기에서 알아야 할 가장 중요한 것은 어느 token으로부터 정보를 가져올지 선택되는 과정과 어떤 정보가 실제로 읽히고 쓰여지는지가 완전히 독립적이라는 것이다.

이걸 편하게 이해하기 위해 attention을 조금 다른 방식으로 재해석해보자. 어떤 attention pattern이 주어졌을 때, attention head의 결과값을 계산하기 위해선 보통 세 가지 단계를 통해 표현되곤 한다:

Value vector를 각 token의 residual stream으로부터 계산한다 ( $v_i = W_V x_i$ ).
그 다음, 여러 value vector를 주어진 attention pattern에 따라 선형 조합 (linear combination)을 통해 계산한다 ( $r_i = \sum_j A_{i,j} v_j$ ).
최종적으로, 각 token마다 결과 vector를 각 head마다 계산할 수 있다 ( $h(x)_i = W_O r_i$ ).

이제 이 과정을 다시 한 번 보자. 모든 단계는 결국 행렬곱인 만큼, 이걸 하나의 계산으로 한꺼번에 표현할 수 있지 않을까? 만약 $x$ 을 2차원 행렬 (각 token 별로 벡터 하나씩 모은 셈)이라고 확장해서 생각해보자. $W_V$ 와 $W_O$ 는 “token 별 해당하는 벡터에 대하여 (즉 context position 마다)” 곱해지는 것을 원하고, 반면 $A$ 는 “context position에 걸쳐” 곱해지는 것을 알 수 있다. Tensor를 이용하면 이런 식으로 계산되는 여러 행렬 곱들을 한 번에 표현할 수 있는 자연스러운 방법을 제시한다 (만약 tensor product가 익숙하지 않다면, 관련하여 [짧은 introduction]을 읽어보아라). 이렇게 하는 또 다른 동기는 다음과 같다. 우리는 실질적으로 행렬에서 행렬로 보내는 선형 변환 (linear mapping)을 원하고 있다는 걸 상기하자 (예를 들어, $[n_{\text{context}}, d_{\text{model}}] \rightarrow [n_{\text{context}}, d_{\text{model}}]$ ). 수학자들은 이러한 linear mapping을 “(2,2)-tensor”라고 부르며, 따라서 이러한 상황을 다루기에 적절한 도구이다.

Tensor product를 이용하면, attention을 다음과 같이 다시 쓸 수 있다:

그리고 결과를 자세히 보면 다음과 같다는 것 또한 알 수 있다:

그런데 attention pattern에 대해서는 어떤 이야기를 할 수 있는 건가? $A$ 는 key vector $k_i = W_K x_i$ 와 query vector $q_i = W_Q x_i$ 를 계산하고 이 둘을 내적하는 형태로 계산된다 ( $A = \text{softmax}(q^T k)$ ). 이것도 한 방에 행렬곱으로 설명이 가능하다: $A = \text{softmax}(x^T W^T_Q W_K x)$ . 참고로, 이런 표현들은 수학적으로는 문제가 없지만 실제 구현단계는 비효율적이다.

Attention head에 대한 관찰

이렇게 attention 구조를 재해석하면 기존에는 잘 보이지 않았던 몇 가지 구조적 특징들이 두드러진다는 장점이 있다:

Attention head는 한 token으로부터 다른 token으로 정보를 운반한다.
- 이에 따라 residual stream의 vector space — 이는 “contextual word embedding”이라고도 불린다 — 안에는 현재 token의 정보가 아닌, 다른 token으로부터 옮겨져 온 정보가 포함되어있는 subspace들이 있다는 것을 예상해볼 수 있다.
Attention head는 (독립적으로, 다른 subspace들에 대하여) 사실상 두 가지 linear operation (즉, $A$ 와 $W_O W_I$ )을 먹인다.
- $A$ 는 token의 정보가 어디로 옮겨질지를 결정한다.
- $W_OW_V$ 는 source token(의 residual stream)으로부터 어떤 정보를 읽고 destination token(의 residual stream)에 어떤 정보를 쓸 지를 결정한다.
유일하게 non-linear한 부분은 $A$ 의 계산이다 (softmax가 있기 때문). 이를 다시 생각하면 attention pattern (즉 $A$ )를 고정시켜놓으면, attention head는 그냥 통째로 선형적(linear)인 계산이 된다. 또 다른 말로는, $A$ 를 고정 시키지 않으면, attention head는 “어느 정도만 선형적”인 셈이다 ( $A$ 말고 나머지 부분은 그래도 선형적)
$W_Q$ 와 $W_K$ , 그리고 $W_O$ 와 $W_V$ 는 서로 항상 붙어다닌다.
- 실제로는 두 개의 행렬의 곱으로 표현은 되지만, 그냥 $W_O W_V$ 와 $W^T_Q W_K$ 는 각각 하나의 low-rank 행렬이라고 생각해도 된다.
- 이렇게 생각하게 되면 “query vector”와 “key vector”는 그렇게 큰 의미가 있는 개념은 아니고, 단지 $W^T_Q W_K$ 행렬을 계산하는 도중에 파생되어 나오는 결과물 정도라고 봐도 무방하다.
- 이제부터 편의상 $W_OW_V = W_{OV}$ , $W^T_Q W_K = W_{QK}$ 이라고 지칭한다.
Attention head 두 개를 곱한 결과는 그냥 하나의 새로운 attention head가 된다 [여기 참조]: $(A^{h_2} \otimes W^{h_2}_{OV}) \cdot (A^{h_1} \otimes W^{h_1}_{OV}) =(A^{h_2} A^{h_1} \otimes W^{h_2}_{OV} W^{h_1}_{OV})$ . 즉 attention pattern이 $A^{h_2} A^{h_1}$ 이고 output-value matrix가 $W^{h_2}_{OV} W^{h_1}_{OV}$ 인 셈. 우리는 이를 “virtual (가상의) attention head”라고 부르며, 뒤에서 따로 다루도록 하겠다.

Zero-layer Transformers

더 복잡한 모델에 대한 이야기를 하기 전에, 간단히 “0-layer” transformer에 대하여 몇 가지 짚고 넘어갈 점이 있다. 0-layer transformer이라고 하면, token을 embedding하고, 바로 unembedding 통해 다음 token을 예측하는 모델을 말한다.

이 모델은 다른 token으로부터 어떠한 정보도 직접적으로 얻을 수 없기 때문에, 그냥 지금 token만 보고 바로 다음에 오는 token을 예측하는 모델이 된다. 즉 $W_U W_E$ 은 bigram (연속적으로 등장하는 두 단어 사이의 관계)를 학습하게 된다.

더 복잡한 모델으로 넘어가서도 $W_U W_E$ 는 계속 등장은 할 것이다. 여기에서 본 예시와 같이, 복잡한 transformer 모델에서도 token embedding에서부터 unembedding까지 residual stream을 타고 direct하게 연결되는 경로를 이야기 하면 $W_U W_E$ 이 나오게 될 수 밖에 없다. 물론 이런 모델에서는 다른 레이어들이 residual stream에 개입하면서 bigram 예측에도 영향을 미치겠지만, 여전히 “residual”을 개념적으로 표현하는 점은 유효하다. 특히, $W_U W_E$ 는 문법 구조상 표현될 수 없는 bigram 관계를 표현하는 것으로 보인다 (예를 들어 “Barack” 뒤에 “Obama”가 오는 것).

One-Layer Attention-Only Transformers

우리는 one-layer attention-only transformer 모델은 bigram 및 “skip-trigram” 모델의 앙상블 (“A…BC” 형태의 패턴)이라고 주장하려고 한다. 직관적으로, (bigram은 이미 설명이 되었고) attention layer가 하나 있으면 현재 token (”B”)로부터 이전 token (”A”)으로 attention을 통해 보게 되고 이에 따라 다음 token (”C”)를 예측할 수 있기 때문이다.

여기서의 목적은 이런 패턴이 가능하다는 것을 직접적으로 보이고, 어떻게 raw weight으로부터 skip-trigram 확률분포에 영향을 미치는지 해석하는 것이다.

The Path Expansion Trick

One-layer attention-only transformer 모델 아키텍쳐를 그림을 통해 구체적으로 확인해 보자:

위에서 사용한 tensor notation 및 tensor product를 이용해 표현한 attention head를 사용하면, transformer를 다음 세 항들의 합으로 나타낼 수 있다:

여기서 product를 확장하는 식으로 조금 더 변형할 수 있다. 중간의 덧셈이 있는 괄호를 풀어서 각 항이 end-to-end 계산을 나타내게 되면 다음과 같이 다시 쓸 수 있다.

이렇게 다시 쓰고 난 수식에 대해 우리는 각 항이 end-to-end 계산 경로를 나타내며, 독립적으로 이해 및 탐구할 수 있으며 각각 합쳐져서 최종 모델의 결과값에 기여한다고 주장한다.

Direct path 항 (첫 번째 항)은 zero-layer transformer에서도 본 모양과 같다. Token끼리 정보교환이 없기 때문에 (즉 Kronecker product의 첫 항이 $\text{Id}$ ), 기여를 할 수 있는 것은 bigram statistic이다.

더 흥미로운 패턴들은 attention head에서 나온다.

Attention head 항들을 Query-Key 및 Output-Value Circuit으로 나누기

각 attention head $h$ 를 나타내는 $A^h \otimes (W_U W^h_{OV}W_E)$ 가 있다는 것을 다시 상기해보자 (그리고 $A^h$ 는 $\text{softmax}(t \cdot W^T_{E} W^h_{QK} W_E \cdot t)$ ). 이 안에 있는 항들을 어떻게 모델의 행동과 연관지을 수 있을까? 거기에다가, 왜 이런 행렬곱들이 나타나게 된 것일까?

이를 이해하기 위해 가장 중요한 것 중 하나는 이 계산은 두 가지 주요한 계산으로 구성되었다는 것을 알아차리는 것이다. 각 (차원이 $[n_{\text{vocab}}, n_{\text{vocab}}]$ 인) 행렬들을 보면:

$W^T_{E} W^h_{QK} W_E$ — 우리는 이 행렬을 “Query-Key (QK) circuit”이라고 부른다. 모든 query 및 key token에 대한 attention 값을 계산해주며, 다른 말로 한 query token이 주어졌을 때 다른 key token을 얼마나 “보고 있는지”를 의미한다.
$W_U W^h_{OV}W_E$ — 우리는 이 행렬을 “Output-Value (OV) circuit”이라고 부른다. 이는 어떤 token을 attention이 “보고” 있을 때 output logit값에 어떠한 영향을 미칠지를 결정하게 된다.

이런 행렬곱들을 조금 더 직관적으로 이해하기 좋게, 실제 모델을 가져다가 놓고 어떤 경로를 표현하고 있는지를 그려보도록 하자. QK circuit은 query 및 key vector가 attention head를 통과하면서 내적곱의 형태로 형성된다. OV circuit은 value vector를 따라 logit까지 경로를 따라가면서 생기게 된다.

Attention 패턴은 source 및 destination token에 대한 함수이지만 [6], destination token이 어떤 source token을 얼마나 ‘바라볼지’ (attention score를 말한다) 결정한 이후에는 그 source token이 output에 어떤 영향을 미칠지는 오로지 그 source token에 의해서만 결정이 된다. 즉, 여러 개의 destination token이 같은 source token을 보고 있다면, 각자의 destination token에 그 source token이 logit에 미치는 영향은 모두 동일하다.

OV 및 QK의 독립성 (Attention pattern 얼리기)

어차피 OV 및 QK circuit은 따로 떼어서 이해할 수 있는 함수이기 때문에, 그 둘을 독립적으로 생각하는 것은 유용할 때가 많다. 하지만 이렇게 생각하는게 정말로 괜찮은 것일까? 이걸 상상해보자: 머릿속으로 같은 input에 대하여 모델을 두 번 돌리는 것이다. 첫 번째 돌릴 때 head별로 attention 패턴들을 따로 저장해놓는다. 이 attention pattern들은 QK circuit에만 영향을 받는다 [7]. 그리고 두 번째 돌릴 때는, 아까 따로 저장해놓은 attention pattern 값들을 불러와 일종의 “얼어있는” attention pattern으로 대체하고 돌려도 아무 문제가 없을 것이다 (즉 임시적이지만 일종의 상수 취급을 하는 것이다). 이렇게 되면 모델이 logit이 token에 대한 선형함수가 된다고 생각할 수 있다! 이러한 사고 방식은 우리가 transformer를 해석할 때 유용할때가 많을 것이다.

Skip-trigram으로서 해석하기

Mechanistic interpretability의 주요 도전 과제 중 하나는 neural network의 parameter를 적절한 시나리오 안에서 해석을 해내는 것이다 (자세한 논의는 이 논문을 참고). 모델 안에서 OV 및 QK circuit을 발견해냄으로써 우리는 token에 대한 linear 혹은 bilinear 함수로 바라보게 되었고 이러한 해석에 성공했다고 볼 수 있다. 위에서 언급한 것 처럼, QK 및 OV circuit이 destination, source, 그리고 output token 사이의 관계를 정립하게 되었고, [source] … [destination] [out] 형태의 skip-trigram 패턴에 대해 본격적으로 이야기 할 수 있게 되었다.

한 가지 짚고 넘어갈 것은 이러한 해석이 trivial한 것이 아니라는 것이다. 일단, 우리가 다루게 될 matrix들은 그 크기가 엄청날 것이고 (vocabulary가 약 ~50,000 token이 될 것이니까 OV matrix 하나에는 최대 2,500억개의 숫자들이 포함되어있을 것이다); 우리는 1-layer attention-only 모델이 일종의 압축된 중국어 방이라는 것을 알게 되었다. 또한 일반적인 선형모델들을 상관관계가 있는 두 변수에 관련하여 해석을 하려고 할 때 발생하는 문제점들을 그대로 안고 오게 된다. 예를 들어, attention head의 weight가 0인걸 발견했다고 하자. 그 이유가 사실 그 head가 중요하지 않은게 아니라, 사실은 같은 역할을 하는 다른 attention head가 존재했기 때문일수도 있다. 마지막으로, 서로 다른 query에 대해서 QK의 weight들끼리 비교하는 방법은 명확하지 않다.

이럼에도 불구하고, 우리는 그래도 transformer의 parameter들이 해석 가능하도록 정리하는데 성공하였다. 몇몇 사소한 문제점들이 있을 수 있어도, 우리는 단순하게 OV 및 QK matrix를 보고 skip-trigram 패턴을 곧장 파악할수 있다. 특히, 이러한 행렬들 안에서 유난히 큰 값을 가지는 곳들을 찾아내면 흥미로운 성질들을 발견할 때가 많다. 우리가 발견한 몇 가지 흥미로운 모델의 행동들을 정리해보도록 하겠다.

복사하기 및 원시적인 형태의 in-context learning

이러한 1-layer 모델들을 뜯어보고 가장 놀라웠던 것은 상당한 attention head들을 정보 복사를 하는데에 쓰고 있다는 점이었다. OV circuit은 일단 head가 attention을 통해서 어떤 token을 보게 되면, 해당 token (혹은 비슷한 token)의 logit 값을 높이게 설정이 되어있다. QK circuit은 다음에 올 법한 token들을 attention을 통해 바라본다. 즉, bigram스러운 패턴으로 다음에 올 법한 token들이 복사되는 것이다.

대부분의 예시들은 자연스럽게 이해가 된다. 두 가지 예시 정도를 추가적으로 다뤄보면: 네 번째 행 같은 경우는 LaTeX를 학습하게 된 결과인 것으로 보이고, 다섯 번째 행은 HTML 코드를 학습하게 된 결과인 것으로 보인다.

이 모든 것들은 복사에 해당하며, 꽤나 자주 보이는 행동 패턴이다.

우리는 또한 더 미묘한 형태의 복사도 볼 수 있었다. 특히 흥미로웠던 예시는 tokenizer가 transformer 안에서 어떻게 행동하는지와 관련이 있다. Tokenizer들은 보통 띄어쓰기를 단어와 합쳐서 처리하는 경향이 있다. 하지만 가끔씩 문단의 앞부분이나 따옴표 직후에 나오는 단어들의 경우 띄어쓰기가 붙어있지 않는 경우들도 존재한다. 이런 경우들은 흔하지 않고, 이에 따라 tokenizer들도 이에 최적화가 잘 되어있지는 않다. 그래서 흔하지 않은 단어들 (예를 들어 고유명사) 같은 경우, 띄어쓰기를 단어의 일부로 통째로 처리하는 경우들이 있다 (” Ralph” → [” Ralph”]). 반면 앞에 띄어쓰기가 없는 경우에는 결국 단어를 분리시킨다 (”Ralph” → [”R”, “alph”]).

Skip tri-gram 패턴들이 이러한 상황을 처리해야 하는 경우들이 자주 생기게 마련이다. 게다가, 우리는 이따금씩 이렇게 앞에 띄어쓰기가 없는 단어를 전담하는 attention head들이 생기는 경우도 보게 되었다. 이 attention head들은 위에서 설명한 경우에서 발생한 단어 파편을 맞이했을 때 (예를 들어, 아까 예시에서 “R”), 띄어쓰기가 포함된 완전한 단어를 찾아내고 (” Ralph”) 이를 통해 나머지를 유추해내는 (”alph”) 역할을 수행해낸다. (이러한 행동은 우리가 뒤에서 2-layer 모델에서 보게 될 induction head랑 비슷한 역할을 한다고 생각할 수 있다).

이러한 행동 패턴들을 함축적으로 다음과 같이 정리할 수 있다:

이러한 패턴들 모두 일종의 원시적인 in-context learning과 같다고 볼 수 있다. Transformer가 주어진 context에 따라 유연하게 적응할 수 있다는 것은 아주 흥미로운 성질이고, 이러한 행동 패턴들이 그러한 성질의 기초적인 버전이라고 볼 수 있을 것이다. 하지만, 뒤에서 2-layer 모델의 경우 in-context learning에 기여하는 훨씬 더 강력한 버전의 알고리즘이 존재한다는 것을 보게 될 것이다.

(다른 종류의 skip-trigram들은 원문을 확인)

Primarily positional attention heads

Attention head들이 token position을 어떻게 다루는지는 아직 다루지 않았는데, 그 이유는 이제 positional encoding과 관련된 여러 방법론들이 제시되었고 ([1], [8], [9]) 우리가 지금까지 발견해낸 식들이 복잡해질 우려가 있기 때문이다.

실제로는, 1-layer 모델 같은 경우 position 정보를 담당하는 attention head들이 적게나마 존재하는 것으로 보이고, 특정 상대적인 position들을 선호하는 것으로 보인다. 밑에서 우리는 현재 혹은 이전 token에 주로 관심을 가지는 attention head를 예시를 통해 보인다.

Skip-trigram “버그”들

QK 및 OV 행렬들을 관찰하는 일이 가장 흥미로웠던 일 중 하나는 겉으로 보기에 이해하기 힘든 현상들에 대한 해석을 가능하게 해줬을 때이다.

우리가 1-layer 모델을 해석한 방법은 skip-trigram들을 $f(a,b,c)=f_1(a,b)f_2(a,c)$ 처럼 일종의 쪼갠 형식으로 바라보는 것이기도 하다. 이렇게 쪼개는 방식은 skip-trigram의 유연성을 오히려 어느 정도 제한하기도 한다. 예를 들어, 만약 어떤 head가 keep… in mind와 keep… at bay의 확률을 올리면, 필연적으로 keep… in bay와 keep… at mind의 확률도 높을 수 밖에 없다. 이런 행동은 어떻게 보면 버그라고 봐도 무방하다. 우리는 이러한 행동을 attention head를 관찰하면서 지속적으로 볼 수 있었다.

OV/QK 행렬 요약

지금까지 우리는 1-layer transformer를 이해하는 문제를 그들의 OV/QK 행렬을 이해하는 문제로 바꿔서 진행하였다. 하지만 이전에도 언급했듯이, OV 및 QK 행렬은 그 크기가 어마어마하다. 가장 큰 값들을 찾아다니는 건 재미있어 보이지만, 이 방법 말고는 이해할 수 있는 방법이 없을까? 적어도 다음 사항들은 대안으로 고려해볼만하다:

OV 및 QK 행렬은 매우 low-rank이다. 크기는 50,000 X 50,000일지라도 결국 rank는 (최대) $d_{\text{head}}$ 일 것이다 (보통 64나 128).
값을 하나하나씩 뜯어보면 다른 패턴들도 보일 때가 있다. 예를 들어, 어떤 head는 사람 이름에 대해서 query가 “ by”인 경우에 강하게 반응하며 (예를 들어, “Anne… by → Anne”) 장소 이름에 대해서는 query가 “ from”인 경우에 대해 강하게 반응한다 (예를 들어, “Canada… from → Canada”). 이런 것을 보아 어떤 cluster 구조가 행렬 안에 존재할 수 있다고 유추해볼만하다.
복사하는 패턴은 특히 OV 행렬에 전방위적으로 나타나며 이를 조금 더 수학적으로 다룰법도 하다.

우리는 아직 정확히 어떤 방법이 옳다는 것을 발견하지 못했지만, 적절한 형태의 matrix decomposition이나 차원 축소 방법론을 적용하는 방향성에 대해서 긍정적으로 생각한다.

복사 패턴 찾아내기

이러한 복사 패턴을 자동적으로 찾아낼 수 있는 건 흥미로운 일일 것이다. 복사는 개념적으로 어떤 벡터를 그 자신에게 보내는 것이기에 (예를 들어, 어떤 token에 대해 그 자신의 확률값을 높이는 것), 적절한 요약통계량 (summary statistics)를 잡으면 이러한 패턴을 확인할 수 있을 거란 아이디어는 설득력이 있어보인다.

하지만, 복사하는 패턴에 대해 정확하게 정의를 내리기 시작하면 조금 골치가 아파진다; 이는 “복사 행렬”이 어떻게 생겨야 하는지에 대해 다양한 (그리고 그럴싸한) 정의들을 생각할 수 있는데, 가장 유용한 정의가 무엇인지 아직 결론을 내리지 못했기 때문이다. 예를 들어, 이 글에서 다루지 않는 조금 더 큰 모델들에 대해서 어떤 attention head들은 gender와 plurality (복수(複數)성), 그리고 tense (시제(時制))의 조합을 복사하는 것으로 보이며, 이는 모델이 정확한 대명사 및 verb conjugation을 사용하는데 일조를 하는 것으로 보인다. 이런 head들의 행렬들 살펴보면 token을 각각 그대로 복사한다기보다는, 조금 더 의미론적인 의미에서 복사한다라고 보는 것이 더 적절하다. 즉 복사라는 것도 우리가 생각한 것보다 더 복잡하다는 것이다.

이 문제를 해결하는데 조금 더 자연스러운 접근은 eigenvector 및 eigenvalue들을 이용하는 것이다. 잠깐 정리하자면, 어떤 벡터 $v_i$ 는 행렬 $M$ 과 적절한 값 $\lambda_i$ 에 대해 $M v_i = \lambda_i v_i$ 를 만족할 때 eigenvector라고 불린다 (그리고 $\lambda_i$ 는 eigenvalue라고 한다). 이제 OV circuit을 분석하기 위해 $M = W_U W^h_{OV} W_E$ 라고 두고, 만약 $\lambda_i$ 가 양의 실수라면 $M v_i = \lambda_i v_i$ 가 어떤 의미를 가지는지를 보자. 저 식에서 $v_i$ 는 token(들의 어떤 적절한 linear combination)일 것이고, $M$ 을 곱한 결과 원래 token 그대로 증폭된다는 것을 의미한다 [10] [추가 설명3]. 그리고 그에 따라 그 token들에 대항하는 logit(들의 어떤 적절한 linear combination) 또한 증폭이 될 것이다. 즉, 직관적으로 생각하면 서로서로 확률값을 높이는 token 그룹이 있다고 해석할 수 있을 것이다. 물론 일반적으로는 음의 eigenvalue들도 존재할 것이기에, 이에 따라 반대되는 해석 (확률값을 줄이는)이 적용되는 token 그룹들도 존재할 것이다. 하지만 복사 패턴을 찾아내는 것에 집중한다면 그러한 행렬들은 양의 eigenvalue들을 지니고 있을 것이다.

여기에서 조금 더 나아가, 총 eigenvalue 중 양의 값을 가지는 비중에 따른 histogram을 그려볼 수도 있겠다.

이를 보아하니 12개 중 10개의 head가 복사를 하고 있는 것으로 보인다! (이는 직접 관찰을 통한 정성 분석 결과와도 일치하는 부분이다)

하지만 복사 행렬이 양의 eigenvalue들을 가져야 하지만, 역으로 모든 양의 eigenvalue를 가지는 행렬들이 우리가 예상하는 복사 행동 패턴을 가지리란 보장은 없어보인다. 행렬들의 eigenvector들은 orthogonal(수직)이라는 보장은 없고 (사실 그럴 가능성은 낮다), 그렇게 되면 언뜻 이해가 안되는 현상들이 발생하곤 한다. 예를 들어, 어떤 token의 확률값을 높이면 특정 token의 확률값이 오히려 감소하는 일이 발생할 수 있다. 물론 그럼에도 양의 eigenvalue를 가지면 “평균적으로는” 복사를 한다고 말할 수 있지만, 보았듯이 양의 eigenvalue를 가진다고 해서 반드시 복사한다고 단정짓는 것은 불가능하다.

(Eigendecomposition이 아닌 다른 방식에 대한 논의는 본문을 참고)

완벽하지는 않지만 그래도 이 글에서는 eigenvalue를 통한 분석을 이용할 것이다 (그래도 강력한 상관관계를 보이므로).

(우리는 1-layer 모델을 “완벽하게” 이해하고 있는가?는 본문을 참고)

Two-Layer Attention-Only Transformers

딥러닝에서 layer를 깊게 쌓는 것은 많은 모델들이 뛰어난 성능을 낼 수 있는 중요한 이유 중 하나라고 바라볼 정도로 중요하다. 특히 depth가 증가할수록 조합 (composition)의 가능성이 생기면서, 모델의 표현력 (expressivity) 또한 올라가게 된다.

Attention head의 조합은 1-layer 및 2-layer attention-only transformer 사이의 큰 차이를 불러일으킨다. 조합이 없으면 사실 attention layer를 두 개 쌓더라도 그냥 skip-trigram의 성능에만 기여하게 될 테고 단순한 lookup table에 불과할 것이다. 하지만 실제로는 이러한 조합을 통해 모델이 마치 어떤 알고리즘을 돌리고 있는 듯한 새로운 행동 패턴이 나타나게 되고, 이는 in-context learning에 중요한 기여를 하게 된다.

세 가지 종류의 조합

우리가 residual stream을 일종의 커뮤니케이션 채널처럼 생각한다는 것을 다시 상기해보자. 모든 attention head들은 residual stream의 subspace로부터 정보를 ( $W_Q$ , $W_K$ , 그리고 $W_V$ 를 통해) 읽어들이고, 처리된 정보를 또 다른 subspace에 ( $W_O$ 를 통해) 쓴다. Attention head 벡터들은 보통 그 차원이 residual stream의 1/10에서 1/100 정도밖에 안되기 때문에, 작은 subspace에만 작용하게 되고 서로 간섭하지 않는 게 생각보다 쉬워진다.

Attention head들끼리 조합을 할 때, 세 가지 종류가 있다:

Q-composition: $W_Q$ 가 이전 head에 의해 영향을 받은 subspace로부터 정보를 읽어들인다.
K-composition: $W_K$ 가 이전 head에 의해 영향을 받은 subspace로부터 정보를 읽어들인다.
V-composition: $W_V$ 가 이전 head에 의해 영향을 받은 subspace로부터 정보를 읽어들인다.

Q- 및 K-composition은 V-composition과는 사뭇 다르다. Q- 및 K-composition은 attention pattern에 영향을 주고, 이에 따라 attention pattern이 더 복잡한 형태를 띄게 하는데 도움을 준다. 반면 V-composition은 attention head가 바라보는 position으로부터 정보를 어떻게 움직일지를 결정하게 되고, 그 결과 독립적인 유닛처럼 행동하게 되고 추가적인 “가상의 (virtual) attention head”가 있는 것 처럼 보이게 된다. 정보를 운반하는 뒤에 정보를 또 운반하는 식으로 조합을 하는 것은 attention head가 attention pattern에 영향을 미치는 것으로는 따라할 수 없는 또다른 패턴을 만들어내게 된다.

이러한 조합을 진정으로 이해하기 위해서, OV 및 QK circuit을 다시 살펴볼 필요가 있다.

Logit의 path expansion

“Logit은 어떻게 계산되는가?”에 대한 질문에 대해 우리가 1-layer 모델에서 했던 것처럼 end-to-end 경로들의 합으로 분석할 수 있다.

위 그림에서 direct path 항과 individual head 항은 1-layer 모델에서 본 것과 매우 비슷하다. 마지막 “virtual attention head” 항은 V-composition에 의해 발생했으며, 이에 대해서는 추후 따로 다루겠다. 참고로, 실제로는 작은 2-layer 모델에서는 이 마지막 항이 큰 역할을 하진 않는 것으로 보인다.

Attention QK circuit의 Path expansion

Logit 쪽에서의 path expansion만 보고서는 2-layer 모델에서 나타나는 고유의 특징을 놓치기 쉽다. 이 특징은 Q- 및 K-composition에 의해서 나타나는 복잡한 attention pattern으로, 이에 대해 더 자세하게 다루도록 하겠다.

예전 1-layer 분석에서 $A^h = \text{softmax}(t \cdot W^T_{E} W^h_{QK} W_E \cdot t)$ 였던걸 기억하자 (이제 QK circuit에 해당하는 $W^T_E W^h_{QK} W_E$ 를 $C^h_{QK}$ 라고 하자). 다시 2-layer 모델로 돌아와서, 일단 첫 번째 attention head에 대해서는 크게 변화한게 없다: $C^{h \in H_1}_{QK} = W^T_E W^h_{QK} W_E$ .

하지만 두 번째 layer의 QK circuit을 보면, Q- 및 K-composition을 고려하기 시작해야한다. $W^h_{QK}$ 를 중심으로 보면, 이 행렬은 결국 residual stream에 작용하는 행렬이라고 볼 수 있다. 첫 번째 layer에서는 $C^{h \in H_1}_{QK} = W^T_E W^h_{QK} W_E = x^T_0 W^h_{QK} x_0$ 이었다 ( $x_i$ 에 대해서는 이 그림을 참조). 하지만 두 번재 layer에서는 $C^{h \in H_2}_{QK} = W^T_E W^h_{QK} W_E = x^T_1 W^h_{QK} x_1$ 이고, $x_1$ 에 작용한다는게 결정적으로 다르다.

이걸 풀어서 수식으로 쓰려고 하면 조금 복잡해지는 건 결국 6차원짜리 tensor를 다뤄야 한다는 점에 있다. 이는 우리가 표현하려고 하는 mapping이 $[n_{\text{context}}, d_{\text{model}}] \times [n_{\text{context}}, d_{\text{model}}] \rightarrow [n_{\text{context}}, n_{\text{context}}]$ 처럼 생겼기 때문이다 [11]. 1-layer 모델에서는 외적을 통해 이걸 직접 계산하는 것을 피할 수 있었지만, 이제는 이게 불가능하다. 이걸 자연스럽게 표현하기 위해서는 (4,2)-tensor를 사용해야 한다. $C^{h \in H_2}_{QK}$ 를 표현하기 위해서는 $A_q \otimes A_k \otimes W$ 꼴의 항들의 합을 많이 사용해야하는데, $x(A_q \otimes A_k \otimes W)y$ 를 하면 $A^T_q xWy A_k$ 가 된다. 즉 $A_q$ 와 $A_k$ 는 각각 query와 key 쪽에서 정보가 어떻게 이동하는지 결정하고, $W$ 는 그 둘이 서로 어떻게 곱해지는지를 결정한다.

2-layer 모델 분석하기

지금까지, 우리는 2-layer attention-only transformer 모델들을 분석할 수 있는 이론적인 기반을 마련하였다 (Logit 부분을 설명하는 OV circuit과 attention head의 attention pattern을 담당하는 QK circuit). 그러면 실전에서 어떻게 분석할 수 있을까? 여기에서 2-layer 모델을 역설계 해보도록 하겠다.

[Attention head composition 분석 부분과 관련된 메모]

일단, 각 Q-, K-, 그리고 V-composition에 의해 두 layer 사이에 attention head가 어떻게 상호작용하는지 분석하는 것이 목표이다. 즉, “두 번째 layer의 attention head가 첫 번째 layer로부터 정보를 얼마나 읽어들이는가?”를 알고 싶은 것이다. 이는 두 번째 layer의 head $h_2$ 및 첫 번째 layer의 head $h_1$ 사이에 다음 값들을 계산하는 방식으로 관찰할 수 있다 [12]:

Q-composition: $||W^{h_2 T}_{QK} W^{h_1}_{OV}||_F / (||W^{h_2 T}_{QK} ||F ||W^{h_1}_{OV}||_F)$
K-composition: $||W^{h_2}_{QK} W^{h_1}_{OV}||_F / (||W^{h_2}_{QK} ||_F ||W^{h_1}_{OV}||_F)$
V-composition: $||W^{h_2}_{OV} W^{h_1}_{OV}||_F / (||W^{h_2}_{OV} ||F ||W^{h_1}_{OV}||_F)$

또한, 이 값들은 행렬들이 random일 경우를 baseline으로 잡아서 비교 된다. 즉, 밑의 그림에서 보여주는 값들은 행렬들이 random일 경우 대비 얼마나 더 정보를 많이 주고받는지를 보여준다. 우리의 해석에 연관이 높은 head들은 색깔을 다르게 표시해놓았다. 첫 번째 layer (아래쪽 점들)에서 색칠된 head (”Previous token heads”)는 단순한 attention pattern을 가지고 있다: 주로 직전 token을 바라보는 경향이 있고, 더 적은 빈도로 현재 token 혹은 이전 이전 token을 바라보는 경우도 있다. 두 번째 layer (윗쪽 점들)에서 색칠된 head는 ”Induction head”라고 부른다.

일단 처음부터 보이는 것은 K-composition이 제일 유의미한 정보들을 많이 보내고, V-composition은 생각보다 영향이 없다는 것이다. K-composition을 보면, 두 induction head들이 previous token head의 영향을 많이 받고 있다는 것이다. 추가적으로 그 두 induction head들이 직전 token만 보는게 아닌 조금 더 뒤까지 볼 수 있는 head와도 연결이 되고 있다.

이런 메커니즘이 어떻게 돌아가는지에 대한 이론을 다루기 전에, interactive diagram을 통해 이런 attention head들을 관찰할 수 있는 예시를 준비하였다. Attention pattern들을 해리포터와 마법사의 돌 첫 문단 위에 색으로 표시한 예시이다.

(실제 예시는 원문 글을 통해 확인하도록 하자)

Induction heads

작은 2-layer attention-only transformer에선, 조합 (composition)이 우리가 induction head라고 부르는 것을 만들기 위해 주로 쓰이는 것으로 보인다. 우리는 1-layer 케이스에서도 모델이 attention head의 상당 부분을 복사하는데에 사용하는 것을 이미 보았다. Induction head는 이보다 훨신 더 강력한 알고리즘으로 in-context learning에 크게 기여한다.

Induction head의 기능

Induction head는 현재 보고 있는 token이 과거에 나타난 적이 있는지 이전 context를 쭉 훑어본다. 만약 적절한 token을 발견하지 못한다면, 그냥 첫 번째 token을 바라보고 아무것도 하지 않는다. 하지만 적절한 token을 찾았다면, 그 token의 다음 token을 바라보고 그걸 복사한다. 이를 통해 이전에 나왔던 token sequence들을 정확하게든 비슷하게든 반복할 수 있게 된다.

1-layer 모델에서 봤던 패턴과 비교를 해보자:

1-layer 모델의 복사 담당 head: [b] … [a] → [b]
- 그리고 가끔 tokenization이 이상해질때는: [ab] … [a] → [b]
2-layer 모델의 induction head: [a][b] … [a] → [b]

2-layer 모델의 알고리즘이 훨씬 더 강력하다. 그냥 현재 token만을 보고 복사를 할지 결정하는 1-layer의 경우보다, 이전에 특정 token이 어떻게 사용됐었는지 파악하고 비슷한 경우를 찾아나서며, 이를 통해 훨씬 더 근거를 가지고 예측을 한다는 것이 근본적으로 다르다. 또한 학습 당시 데이터에 존재하는 특정 단어들끼리의 패턴에 의존하는 것도 아니기에 distribution shift에도 강한 모습을 보인다 (추후에 이런 induction head들이 random token sequence에서도 잘 작동한다는 것을 보일 것이다).

밑의 예시는 induction head들이 해리포터의 첫 문단을 처리하는데 도움을 주는 케이스를 보여준다.

또한 이전에 언급했던 것 처럼, induction head에 대해서 세운 이론을 무작위 token sequence에 적용했을 때를 보려고 한다. Induction head는 이전 context sequence에서 적절한 token을 복사해 현재 위치로 끌고 와 붙여넣는 기능을 구현하는 것이기에, 이론상 무작위적인 token sequence에 적용했을때도 문제가 없어야 한다.

[실제 demo를 보면 반복되는 무작위 token sequence를 던져주면 이전에 나왔던 패턴을 induction head가 잘 찾아낸다. 링크: https://transformer-circuits.pub/2021/framework/index.html#induction-heads]

Induction head의 작동 원리

Induction head의 핵심 원리는 attention이 강하게 반응하는 key가 직전 token에서 발현이 될 수 있는 기능에서부터 온다. Query는 “유사한” key 벡터들을 찾아내기를 원하고, 그 key 벡터가 애초에 한 칸 밀려있기 때문에 자연스럽게 다음에 오는 token을 바라볼 수 있게 된다.

Induction head를 재현하기 위한 최소한의 조건은 previous token head와의 K-composition을 통해 key 벡터를 다음 칸으로 이동시키는 것이다. 이는 $\text{Id}\otimes A^{h_{-1}} \otimes W$ 꼴의 계산에 해당이 된다 ( $A^{h_1}$ 은 이전 token을 가리키는 attention pattern을 가리킨다). 만약 여기에서 $W$ 가 copying matrix처럼 행동한다면 이 항은 source token 이전의 token이 destination 토큰과 동일할 때 attention score를 증가시킬 것이다.

Mechanistic theory 확인해보기

우리의 mechanistic theory에 의하면, induction heads는 두 가지 기능을 할 수 있어야 한다.

“복사”가 가능한 OV circuit matrix
이전 token들을 보면서 (현재 token과) “대응시키는” QK circuit

완벽하진 않지만, 이전에 사용했던 eigenvalue 기반 분석 방식으로 삼아 진행하도록 하겠다. 해당 분석에서 그랬던 것 처럼 양의 eigenvalue 비중을 QK 및 OV circuit에 대해 계산한 값을 2차원 plot 위에 그리면, induction head들은 오른쪽 상단 구석에 다 몰려있는 것을 확인할 수 있다.

(이런 분석이 전에 했던 분석과 비교했을 때 순환논리에 빠지는가?에 대한 논의는 원문을 확인하도록 하자.)

(Term Importance Analysis에서는, 이전에 언급했던대로 과연 virtual attention head가 정말로 덜 중요한지 ablation을 통해 실험적으로 보인다. 자세한 내용은 원문을 확인하도록 하자. 또한 ablation을 어떻게 수행하는지는 추후 번역할 causal scrubbing 글에서 더 자세하게 다뤄질 예정이다.)

그래서 앞으로 나아갈 길은?

이 글을 통해 우리는 1-layer / 2-layer attention-only transformer를 이해하는데 큰 진전을 보였다. 하지만 아무래도 우리의 최종 목표는 transformer 모델을 특정 레이어 수 등의 제한사항 없이 일반적으로 이해하는 것이다. 이번 분석이 이 최종목표로 가는데에 도움이 되었는가? 이런 특수 케이스에 대한 해석을 했다고 해서 일반적인 transformer를 이해하는데 도움이 되는가? 이에 대한 자세한 해답은 추후 연구에서 다루겠지만, 일단 기본적으로 우리의 대답은 ‘그렇다’ 쪽으로 모아지고 있다. 적어도 (더 큰 언어모델을 포함한) 일반적인 transformer의 일부라도 이해하는데 일조를 했다고 할 수 있다.

그 이유 중 하나로 일반적인 transformer 모델들도 거의 순수 attention에만 의존하는 circuit들이 존재하는 것으로 보이기 때문이다. (해석이 까다로운) MLP 레이어들이 있을 때에도, attention head는 여전히 residual stream에 작용을 하고 embedding을 통해 서로 직접적으로 상호작용을 한다. 실제로도, attention head와 embedding만으로 이루어진 해석 가능한 circuit들을 발견할 수 있다. 모델 전체를 이해하는데 아직 부족할지라도, 이런 해석 가능한 부분에 대해서는 역설계를 하기에 충분한 위치에 있다.

게다가, 우리가 이번의 toy 모델들에서 분석한 attention head 및 circuit과 비슷한 구조들이 거대 언어모델에도 나타나는 것으로 보인다! 특히나, 우리는 이런 큰 모델들이 많은 수의 induction head들을 만든다는 것을 발견하였고, 우리가 여기에서 본 것과 같이 이런 head들은 이전 head와 K-composition을 함으로써 구현이 된다. 이는 모델 크기와 상관없이 in-context learning의 핵심 엔진이 되는 것 같다 — 이에 대해서는 추후 글에서 더 자세하게 다룰 것이다.

그렇다고 해도, 사실 큰 모델들에 대해서는 해석할 수 있는 정도가 매우 제한적일 것이다. 일단, MLP 레이어들은 transformer 전체 parameter의 2/3 정도 차지한다. 이런 parameter들과 씨름하지 않고서는 이 모델들의 많은 부분들이 여전히 미스테리로 남을 것이다 (심지어 그 MLP parameter들이 attention과도 상호작용을 하고 있다는 것을 생각하면 더 힘들긴 하다). Mechanistic한 레벨에서 보면, 이런 circuit들은 수학적으로 꽤나 괜찮은 구조를 가지고 있다 (링크 참조). 그럼에도 결국 각 뉴런들에 대한 해석을 가능하게 하는 것이 앞으로 나아갈 수 있는 방향이라는 것은 확실해보이며, 사실 아직 이에 대해서는 큰 성공을 거두지 못하였다.

결론적으로, 이 글에서의 목표는 미래 연구를 위한 발판을 마련하는 것이며, 아직 연구가 되어야 할 것이 많이 남았다.

각주

[1] Vaswani et al., Attention is all you need, NeurIPS 2017

[2] 실제로는 완벽에 가까울 정도로 선형적인 residual stream은 존재하는 경우가 드물다. 가장 흔하게 쓰이는 아키텍쳐인 ResNet 또한 residual stream 그 자체에, 아니면 residual stream에 접근할때마다 비선형 함수를 (non-linear activation function) 먹인다.

[3] (역자) 그냥 여기에 상주하는 벡터의 차원이 높다라는 말과 같다고 보면 된다.

[4] (역자) Superposition (중첩)은 보통 물리학에서 많이 쓰이는 용어이며, 가장 대표적인 예시는 파동에서 일어나는 현상을 설명할 때 나온다. 수면 위 여러 개의 파동이 동시에 (”중첩된 상태”로) 같은 지점에 존재할 수 있으며, 이에 따라 나타나는 최종적인 파동은 각각 파동을 더한 형태가 된다. 이에 따라 보강이나 간섭 등의 현상 또한 발생한다. 여기에서도 “여러 개의 다른 정보가 동시에 같은 벡터에 존재한다”라는 개념으로 쓰인 것으로 이해하면 된다.

[5] 몇몇 MLP 뉴런들은 input과 output weight을 비교했을 때 상당히 음의 cosine similarity를 지니는 것으로 보이며 (즉 서로 거의 180도 틀어진 벡터라는 뜻), 이는 residual stream의 정보를 지우는 역할을 하는 것을 시사한다. 유사하게, 몇몇 attention head들은 해당하는 $W_O W_V$ 행렬이 큰 음의 eigenvalue들이 있고 대부분 현재 token을 바라보고 있으며, 이 또한 정보를 지우고 있는 것을 시사한다. 물론 이러한 것들이 “메모리 관리”에 해당하는 여러 메커니즘의 일부일 가능성이 있으며, 이런 정보 지우기 또한 선택적으로 일어날 가능성도 높다.

[6] 조금 더 정확히는, softmax normalization 때문에 단순히 source token 및 destination token 두 개에 대한 함수가 아니라 destination token과 destination token까지의 모든 source token에 대한 함수라고 해야한다.

[7] 물론 1-layer보다 더 레이어를 많이 쌓게 되면 단순하게 $W^T_{E} W^h_{QK} W_E$ 로만 쓸 수는 없다.

[8] Su et al., Roformer: Enhanced transformer with rotary position embedding, Neurocomputing 568: 127063 (2024)

[9] Press et al., Shortformer: Better language modeling using shorter inputs, ACL/IJCNLP (1) 2021: 5493-5505

[10] Token들도 엄청난 고차원 one-hot vector이라고 생각하면 당연히 그들의 linear combination도 만들 수 있을 것이다.

[11] (역자) (Query 쪽으로 흘러들어가는 residual stream) $\times$ (Key 쪽으로 흘러들어가는 residual stream) $\rightarrow$ (Attention pattern matrix)

[12] (역자) 예를 들어 $||W^{h_2 T}_{QK} W^{h_1}_{OV}||_F / (||W^{h_2 T}_{QK} ||_F ||W^{h_1}_{OV}||_F)$ 를 보자. $W^{h_2 T}_{QK} W^{h_1}_{OV}$ 는 첫 번째 layer의 attention head ( $h_1$ )에서 OV matrix를 통해 residual stream에 정보가 쓰이고, 두 번째 layer의 attention head ( $h_2$ )으로 query로써 정보가 흘러들어가는 경로를 나타낸다. 결국 저 식은 $W^{h_2 T}_{QK} W^{h_1}_{OV}$ 행렬의 크기 (norm)이 각자의 행렬 대비 얼마나 비중이 있는지를 계산한다. 이 값은 $||AB|| \leq ||A|| \cdot ||B||$ 이기에 항상 1보다 작거나 같다. (참고)

[privileged basis] (역자) 이를 설명한 더 자세한 글은 여기를 보면 된다. 선형대수적으로 간단히 말해서, 어떤 벡터가 “privileged basis”가 있다는 말은 각 차원이 서로 다른 독립적인 개념과 대응이 된다는 것을 의미한다. 예를 들어, token vector는 개념적으로 vocabulary만큼의 차원을 가진 벡터 중 해당 token만 1인데, 생각해보면 이게 가능한 이유가 각 차원 하나하나가 다른 token을 의미하기 때문이다. 또한 attention pattern 행렬도 privileged basis를 가지고 있는데, 각 element 하나가 명확한 의미가 있기 때문이다 ( $\alpha_{i,j}$ 는 $i$ 번째 token과 $j$ 번째 token 사이의 attention 값을 의미). 하지만 residual stream 같은 벡터는 각 차원이 무엇을 의미하는지는 알 수 없고, 보통 여러 차원에 걸쳐서 표현되고 있다.

[추가 설명1] (역자) 선형대수에서, 벡터 $v$ 는 정확히 말하면 몇 가지 공리를 만족하는 집합의 원소에 불과하다. 이렇게 생각하면 $[0,1,2,-1]$ 과 같은 익숙한 그 것과 굉장히 다른 정의하고 볼 수 있다. 사실 우리가 벡터를 네모난 괄호 안에 숫자를 나열하는 방식으로 표현하기 위해서는, 일종의 ‘기준축’이 되는 basis가 존재해야한다. 가장 접하기 쉬운 basis는 standard basis로, 2차원 공간에서 x축과 y축에 해당한다. 이 basis 위에서 어떤 벡터가 $[2,1]$ 이라는 것은 사실 첫 번째 standard basis (즉 $[1,0]$ ) 방향으로 2만큼, 두 번째 standard basis (즉 $[0,1]$ ) 방향으로 1만큼 뻗어있는 말이 숨겨져있다. 하지만 꼭 standard basis로 해야만 한다는 법이 있는가? 그런건 없다. 임의의 basis 두 개를 잡아도 상관없다 (심지어 두 개가 수직일 필요도 없다. 평행만 아니면 된다). 예를 들어, basis 두 개를 $[2/\sqrt{5}, 1/\sqrt{5}]$ , $[-1/\sqrt{5}, 2/\sqrt{5}]$ 로 잡으면, 원래 standard basis로 $[2,1]$ 이었던 벡터는 $[\sqrt{5}, 0]$ 이 된다. 여기서 중요한 것은 원래 벡터는 변화하지 않고 그대로 있다는 것이다. 우리가 방금 한 것은 standard basis를 시계 방향으로 90도 회전한 것이라고 생각할 수 있고, 원래 백터는 (선형 변환 하에서) 실질적으로 변하지 않았으니, 언제든지 반대 방향 (여기서는 반시계 90도)으로 회전시켜주면 다시 돌아온다. 아마 원 글에서는 residual stream에 어떤 rotation (rotation을 하는 것도 행렬으로 표현할 수 있다)을 적용하면 다른 부분에서 반대 방향으로 문제없이 rotation할 수 있다는 것을 말하는 것으로 보인다.

[추가 설명2] (역자) 쉽게 말해서, residual stream에 있는 벡터의 모든 차원을 사용하는게 아니라 일부만 사용할 수 있다는 것을 의미한다. 예를 들어, residual stream에 786차원짜리 벡터를 사용하는 모델이 있다고 가정하자. Input 쪽에 가까운 한 레이어가 residual stream의 10~~20 차원에 정보 A를, 311~~359 차원에 정보 B를 더했다고 생각해볼 수 있다 (이는 앞서 말한 “residual stream에 쓰기 전 linear projection을 거친다”에서 이 projection matrix가 어떻게 생겼는가에 따라 매우 가능한 시나리오다). 정보 A는 뒷쪽의 레이어에서 해당 linear projection이 선택적으로 통과시킬수 있고, 정보 B 또한 독립적으로 또 다른 레이어에 의해 이용될 수 있다.

[추가 설명3] (역자) 물론 증폭하는게 말이 되기 위해서는 $\lambda_i > 1$ 이어야 될 것이다.

[Short introduction to tensor products] (원문 번역 + 역자 추가 설명) 우리는 머신러닝에서 matrix와 tensor들을 특정한 방식으로 곱해야 하는 경우가 자주 발생한다. 특히 transformer에서는, activation은 보통 2차원이며 (context 위치에 따른 여러 벡터들의 모임), “(context) position 별로” 곱하거나, “여러 position에 걸쳐서” 곱하거나, 이를 동시에 행해야 할 때가 많다. Tensor product (혹은 Kronecker product —(역자: 엄밀하게 정의하면 둘이 다르지만 여기서는 Kronecker를 기준으로 설명하겠다)—)은 이러한 계산을 깔끔하게 표현해준다. 이를 위해 $\otimes$ 기호를 사용한다. (추가 설명) 원문에서는 생략이 된 것 같은데, 이제부터 벡터 $x$ 는 vectorized/flatten된 batch token embedding matrix라고 생각하면 된다 (밑 그림 참조).

단도직입적으로, 어떤 (vectorized 된) 벡터 $x$ 앞에…

$\text{Id}\otimes W$ 와 같은 것이 붙으면 이는 각 position마다 행렬곱을 하는 것과 같다 (batch-wise multiplication이랑 비슷하다고 생각할 수 있다).
$A \otimes \text{Id}$ 와 같은 것이 붙으면 position에 걸쳐서 곱하는 것과 같다.
$A \otimes W$ 은 두 가지를 합친다: 각 context position마다 일단 $W$ 를 곱한 다음에 $A$ 에 의해 서로 다른 context position끼리 섞인다.
$(A \otimes B) \cdot (C \otimes D) = (AB) \otimes (CD)$ 도 가능하다.

[Attention head composition 분석 부분과 관련된 메모] (역자) 원 글에서 보면, 저자들이 글을 쓴 이후 코드 상 버그로 인해 기존에 썼던 내용에 오류가 있다는 것을 명시하고 있다. 이에 따라 부록에서 교정된 내용을 서술하고 있는데, 따라서 이 부분은 부록을 위주로 참조하여 번역하였다.

Key, query, 그리고 value 벡터들은 low-rank인 $W_Q^\top W_K$ 와 $W_O W_V$ 를 계산을 하는 도중 발생하는 중간 결과물이라고 생각할 수 있다.

W^H_O \begin{bmatrix} r^{h_1}\\ r^{h_2}\\ \cdots \end{bmatrix} = \begin{bmatrix} W^{h_1}_O,W^{h_2}_O,\cdots \end{bmatrix} \cdot \begin{bmatrix} r^{h_1}\\ r^{h_2}\\ \cdots \end{bmatrix} = \sum_i W^{h_i}_O r^{h_i}.

Value vector를 각 token의 residual stream으로부터 계산한다 ( $v_i = W_V x_i$ ).

그 다음, 여러 value vector를 주어진 attention pattern에 따라 선형 조합 (linear combination)을 통해 계산한다 ( $r_i = \sum_j A_{i,j} v_j$ ).

최종적으로, 각 token마다 결과 vector를 각 head마다 계산할 수 있다 ( $h(x)_i = W_O r_i$ ).

Attention head는 (독립적으로, 다른 subspace들에 대하여) 사실상 두 가지 linear operation (즉, $A$ 와 $W_O W_I$ )을 먹인다.

$A$ 는 token의 정보가 어디로 옮겨질지를 결정한다.

$W_OW_V$ 는 source token(의 residual stream)으로부터 어떤 정보를 읽고 destination token(의 residual stream)에 어떤 정보를 쓸 지를 결정한다.

유일하게 non-linear한 부분은 $A$ 의 계산이다 (softmax가 있기 때문). 이를 다시 생각하면 attention pattern (즉 $A$ )를 고정시켜놓으면, attention head는 그냥 통째로 선형적(linear)인 계산이 된다. 또 다른 말로는, $A$ 를 고정 시키지 않으면, attention head는 “어느 정도만 선형적”인 셈이다 ( $A$ 말고 나머지 부분은 그래도 선형적)

$W_Q$ 와 $W_K$ , 그리고 $W_O$ 와 $W_V$ 는 서로 항상 붙어다닌다.

실제로는 두 개의 행렬의 곱으로 표현은 되지만, 그냥 $W_O W_V$ 와 $W^T_Q W_K$ 는 각각 하나의 low-rank 행렬이라고 생각해도 된다.

이렇게 생각하게 되면 “query vector”와 “key vector”는 그렇게 큰 의미가 있는 개념은 아니고, 단지 $W^T_Q W_K$ 행렬을 계산하는 도중에 파생되어 나오는 결과물 정도라고 봐도 무방하다.

이제부터 편의상 $W_OW_V = W_{OV}$ , $W^T_Q W_K = W_{QK}$ 이라고 지칭한다.

Attention head 두 개를 곱한 결과는 그냥 하나의 새로운 attention head가 된다 [여기 참조]: $(A^{h_2} \otimes W^{h_2}_{OV}) \cdot (A^{h_1} \otimes W^{h_1}_{OV}) =(A^{h_2} A^{h_1} \otimes W^{h_2}_{OV} W^{h_1}_{OV})$ . 즉 attention pattern이 $A^{h_2} A^{h_1}$ 이고 output-value matrix가 $W^{h_2}_{OV} W^{h_1}_{OV}$ 인 셈. 우리는 이를 “virtual (가상의) attention head”라고 부르며, 뒤에서 따로 다루도록 하겠다.

$W^T_{E} W^h_{QK} W_E$ — 우리는 이 행렬을 “Query-Key (QK) circuit”이라고 부른다. 모든 query 및 key token에 대한 attention 값을 계산해주며, 다른 말로 한 query token이 주어졌을 때 다른 key token을 얼마나 “보고 있는지”를 의미한다.

$W_U W^h_{OV}W_E$ — 우리는 이 행렬을 “Output-Value (OV) circuit”이라고 부른다. 이는 어떤 token을 attention이 “보고” 있을 때 output logit값에 어떠한 영향을 미칠지를 결정하게 된다.

OV 및 QK 행렬은 매우 low-rank이다. 크기는 50,000 X 50,000일지라도 결국 rank는 (최대) $d_{\text{head}}$ 일 것이다 (보통 64나 128).

Q-composition: $W_Q$ 가 이전 head에 의해 영향을 받은 subspace로부터 정보를 읽어들인다.

K-composition: $W_K$ 가 이전 head에 의해 영향을 받은 subspace로부터 정보를 읽어들인다.

V-composition: $W_V$ 가 이전 head에 의해 영향을 받은 subspace로부터 정보를 읽어들인다.

Q-composition: $||W^{h_2 T}_{QK} W^{h_1}_{OV}||_F / (||W^{h_2 T}_{QK} ||F ||W^{h_1}_{OV}||_F)$

K-composition: $||W^{h_2}_{QK} W^{h_1}_{OV}||_F / (||W^{h_2}_{QK} ||_F ||W^{h_1}_{OV}||_F)$

V-composition: $||W^{h_2}_{OV} W^{h_1}_{OV}||_F / (||W^{h_2}_{OV} ||F ||W^{h_1}_{OV}||_F)$

[7] 물론 1-layer보다 더 레이어를 많이 쌓게 되면 단순하게 $W^T_{E} W^h_{QK} W_E$ 로만 쓸 수는 없다.

[11] (역자) (Query 쪽으로 흘러들어가는 residual stream) $\times$ (Key 쪽으로 흘러들어가는 residual stream) $\rightarrow$ (Attention pattern matrix)

[추가 설명3] (역자) 물론 증폭하는게 말이 되기 위해서는 $\lambda_i > 1$ 이어야 될 것이다.

단도직입적으로, 어떤 (vectorized 된) 벡터 $x$ 앞에…

$\text{Id}\otimes W$ 와 같은 것이 붙으면 이는 각 position마다 행렬곱을 하는 것과 같다 (batch-wise multiplication이랑 비슷하다고 생각할 수 있다).

$A \otimes \text{Id}$ 와 같은 것이 붙으면 position에 걸쳐서 곱하는 것과 같다.

$A \otimes W$ 은 두 가지를 합친다: 각 context position마다 일단 $W$ 를 곱한 다음에 $A$ 에 의해 서로 다른 context position끼리 섞인다.

$(A \otimes B) \cdot (C \otimes D) = (AB) \otimes (CD)$ 도 가능하다.