A Mathematical Framework for Transformer Circuits 리뷰 - (1)

https://transformer-circuits.pub/2021/framework/index.html

A Mathematical Framework for Transformer Circuits

이번 포스팅에서는 Transformer 모델에 대해 복습하는 차원에서 Anthropic 사의 "A Mathematical Framework for Transformer Circuits"이라는 Article에 대해 리뷰해 보겠습니다.

 

이 paper는 Transformers 모델을 리버스엔지니어링하기 위한 매우 초기 단계의 예비작업을 시도합니다.

그들은 가장 단순한 모델부터 차근차근 단계를 밟아나가는 것이 가장 유익하다는 것을 알게 되었고

더 크고 복잡한 모델에 적용할 수 있는 간단한 알고리즘 패턴, 모티브, 프레임워크들을 발견하는 것을 목표로 하고 있습니다.

그래서 이 paper에서는 오직 attention block만을 사용하고 2 layer 이하 만들 사용 하여 Transformers를 공부합니다.

(GPT-3는 96개의 layer를 가지고 있고 attention block과 MLP block을 번갈아 사용)

 

그들은 Transfomers 모델의 작동을 새롭지만 수학적으로 동등한 방식으로 개념화하여 작은 모델을 이해하고

내부적으로 어떻게 작동하는지에 대해 상당한 이해를 얻을 수 있었다고 합니다.

특히 "induction heads"라고 부르는 특정한 attention heads는 작은 모델에서 in-context learning을 설명할 수 있고

이 heads는 최소 2개의 attention layer에서만 발생한다는 사실을 발견했습니다.

뒤에서 이러한 heads가 특정 데이터에서 작동하는 몇 가지 예시도 살펴볼 것입니다.

 

Summary of Results

Reverse Engineering Results

Transformers를 리버스 엔지니어링하기 위해, 여러 toy, attention-only 모델들을 리버스 엔지니어링을 하였고

몇 가지 사실을 알게 되었습니다 :

• Zero layer transformers model은 bigram 통계를 모델링합니다. bigram table은 가중치에서 직접 액세스 할 수 있습니다.
• One layer attention-only transformers은  bigram과 skip-trigram모델들의 앙상블입니다. bigram과 skip-trigram table은 모델을 실행하지 않고 가중치에서 직접 액세스 할 수 있습니다. 이러한 skip-trigrams는 놀라울 정도로 표현력이 뛰어납니다.
  여기에는 일종의 매우 간단한 in-context learning을 구현하는 것도 포함됩니다.
• Two layer attention-only transformers는 attention heads 구성을 사용하여 훨씬 더 복잡한 알고리즘을 구현할 수 있습니다. 이러한 알고리즘은 가중치에서 직접 감지할 수도 있습니다. two layer 모델은 attention head 구성을 사용하여 매우 일반적인 in-context learning 알고리즘인 "induction heads"를 생성합니다.
• One layer와 Two layer attention-only transformers는 매우 다른 알고리즘을 사용하여 in-context learning을 수행합니다.
  two layer attention heads는 질적으로 더 정교한 inference-time 알고리즘(특히 induction head라고 부르는 특수한 유형의 attention head를 사용)을 in-context-learning을 수행하는 데 사용하고 더 큰 모델에 관련해 중요한 전환점을 형성합니다.

 

Conceptual Take-Aways

또한 그들은 transformer 구조에 여러 가지 미묘한 세부사항 때문에 이전에 수행한 inceptionV1 Curcuit(링크)과는 상당히 다른 방식으로
리버스 엔지니어링에 접근해야 한다는 사실을 알게 되었습니다. 밑에서 더 자세히 다루겠지만 간단히 요약하자면 아래와 같습니다 :

• Attention heads는 각각의 결과를 출력하여 residual stream에 추가하는 독립적인 연산으로 이해할 수 있습니다.
  attention heads는 계산 효율성을 위해 "연결하고 곱하기"라는 다른 공식으로 설명하는 경우가 많지만 수학적을 동일합니다.
• Attention-only 모델은 토큰을 logits의 변화에 매핑하는 해석 가능한 end-to-end 함수의 합으로 작성될 수 있습니다.
  이러한 함수는 모델을 통과하는 "경로"에 해당하며 attention 패턴을 고정하면 선형적입니다.
• Transformers는 엄청난 양의 선형 구조를 가지고 있습니다. 합계를 분해하고 행렬의 chain들을 곱하는 것만으로도
  많은 것을 배울 수 있습니다.
• Attention heads는 두 가지 독립적인 계산으로 이해할 수 있습니다.
  1. attention 패턴을 계산하는 QK(query-key) 회로
  2. 각 토큰이 출력에 미치는 영향을 계산하는 OV(output-value) 회로
• key, query, value 벡터는 low-rank 행렬 계산의 중간 결과로 생각할 수 있습니다.
  이것은 reference 없이 transformers를 설명하는데 유용합니다
• attention heads를 구성하하면 transformers의 표현력을 크게 향상합니다.
  attnetion heads는 key, query, value에 따라 세 가지 방식으로 구성될 수 있습니다.
  key와 query 구성은 value 구성과는 매우 다릅니다.
• transformer의 모든 구성요소 (token embedding, attention heads, MLP layers, unembedding)는 residual stream의
  서로 다른 하위공간을 읽고 쓰면서 서로 통신합니다.
  residual stream 벡터를 분석하는 대신 residual stream을 모델의 경로에 해당하는 모든 다른 통신 채널로 분해하는 것이 도움이 될 수 있습니다.

 

Model Simplifications

그들은 이 paper의 아이디어를 가장 깔끔한 형태로 보여주기 위해, 단순화한 "toy transformers"에 초점을 맞춥니다.

MLP layer가 없는 attention-only transformers에 집중할 것이며 매우 드라마틱하게 단순화한 구조입니다.

Distill circuit 작업에서 마주하지 못한 새로운 도전과제에 직면하여 동기부여가 되었고
이를 단독으로 고려해 해당문제에 대해 우아한 처리 방법을 제공할 수 있었습니다. 
하지만 MLP layer 부분은 연구하기가 훨씬 어려워 이해하는 데 성공하지는 못했습니다.

그들 연구의 주요 약점이지만 앞으로 집중적으로 해결해 나갈 계획이라고 합니다.

 

또한 명확성과 단순성을 위해 몇 가지 변경사항을 적용했습니다. 

• bias를 고려하지는 않지만 편향이 있는 모델은 가중치로 접어서 항상 1이 되는 차원을 만들어 편향 없이 시뮬레이션할 수 있습니다.
• attention-only transformers의 bias는 대부분 기능적으로 logits의 bias로 곱해집니다.
• layer normalization을 무시합니다. (명시적으로 고려해야 할 복잡성을 상당히 증가시키며, variable scaling, layer norm을 인접한 가중치로 병합할 수 있습니다.)
• 구현상 몇 가지 불편함을 감안해 layer norm이 batch normalization(인접 파라미터로 완전히 확장가능한)을 대체할 수 있을 것으로 기대합니다.

High-Level Architecture

본문에서는 GPT-3와 같은 autoregressive decoder-only transformer language model만을 다룹니다.

 

transformer는 token embedding을 시작으로 residual block으로 이어지고 마지막에 token unembedding을 진행합니다.

각각의 residual block은 attention layer를 거치고 MLP layer를 거치를 구조로 구성되어 있습니다.

attention과 MLP layer 둘 다 각각 linear projection을 수행하여 residual stream에서 input을 "읽은" 다음,
linear projection을 다시 추가하여 결과를 residual stream에 "기록" 합니다.

각 attention layer는 병렬로 작동하는 multiple heads로 구성되어 있습니다.

Virtual Weights and the Residual Stream as a Communication Channel

transformer의 상위 레벨 아키텍처의 주요 특징 중 하나는 각각의 layer가 그 결과를 "residual stream"이라 부르는 곳에 추가한다는 것입니다. residual stream은 단순히 이전 layer의 모든 출력과 원본 embedding의 합계를 말합니다.

일반적으로 residual stream을 통신 채널로 생각합니다. 왜냐하면 residual stream 자체는 어떠한 처리도 하지 않고 모든 layers가 그것을 통해 통신하기 때문입니다.

 

residual stream은 깊은 linear 구조를 가지고 있습니다.

(완전히 linear 한 residual stream은 neural network 아키텍처 중에서 매우 이례적이라는 점은 주목할 가치가 있습니다. ResNet에서 조차 residual stream에 activation function을 갖거나 residual stream에 액세스 할 때마다 적용됩니다.)

 

모든 layers는 시작 시 residual stream에서 정보를 "읽어 들이기"위해 임의의 linear transformation을 수행합니다.

(이렇게 하면 각 layer가 시작될 때 layer normalization이 무시되지만, constant scalar까지 layer normalization는 constant affine transformation이며 linear transformation으로 접을 수 있습니다. layer normalization appendix 참조)

그리고 출력을 다시 residual stream에 "쓰기"위해 추가하기 전 또 다른 임의의 linear transformation을 수행합니다.

residual stream의 이러한 linear, additive 구조는 중요한 함의를 담고 있습니다.

 

기본적인 결과 중 하나는 residual stream에 "privileged basis"가 없다는 것입니다.

모델의 작동방식을 변경하지 않고도 residual stream과 상호작용하는 모든 행렬을 회전시켜 residual stream을 회전시킬 수 있습니다.

 

Virtual Weights

residual stream이 선형적일 때 특히 유용한 결과는 residual stream을 통해 interactions를 곱하여 모든 layer 쌍 (다른 layer로 분리된 layer 포함해서)을 직접 연결하는 암시적인 "Virtual Weights"를 생각할 수 있다는 것입니다. virtual weights는 한 layer의 출력 가중치와 다른 layer의 입력 가중치의 곱으로, 이전 layers에서 기록한 정보를 이후의 layers에서 읽는 정도를 설명합니다. 

 

Subspaces and Residual Stream Bandwidth

residual stream은 고차원 벡터 공간입니다. 작은 모델에서는 수백 차원이 될 수 있고, 큰 모델에서는 수만 차원에 이를 수 있습니다.

즉, layers는 서로 다른 하위 공간에 정보를 저장하여 서로 다른 layers에 다른 정보를 보낼 수 있습니다. 

이것은 attention heads의 경우에 특히 중요합니다. 모든 개별 head는 비교적 작은 하위 공간에서 작동하고 (종종 64차원 or 128차원)

완전히 분리된 하위 공간에 쉽게 쓸 수 있으며 상호작용을 하지 않기 때문입니다.

 

한 번 추가된 정보는 다른 layers에서 적극적으로 지우지 않는 한 하위 공간에 유지됩니다. 이러한 관점에서 residual stream의 차원은
'memory' 또는 'bandwidth'와 같은 개념이 됩니다. 

기존의 token embedding과 unembedding은 대부분 상대적으로 차원의 작은 부분과 상호 작용합니다.

따라서 대부분의 차원은 다른 layers가 정보를 저장할 수 있는 "여유(free)" 공간으로 남습니다.

 

residual stream 대역폭에 대한 수요가 매우 높을 것을 예상할 수 있습니다.

일반적으로 residual stream에는 정보를 이동하는 차원보다 훨씬 더 많은 'computational 차원'(neurons와 attention head 결과 차원 같은)이 있습니다. 단일 MLP layer에는 일반적으로 residual stream의 차원보다 4배 더 많은 뉴런이 있습니다. 

예를 들어 50 layer transformer의 25 layer에서 residual stream은 그 앞의 차원보다 100배 많은 뉴런을 가지고 있으며,

그 뒤의 차원보다 100배 많은 뉴런과 통신을 시도하면서 어떻게든 superposition 한 상태로 통신을 시도합니다.

이러한 텐서를 "bottleneck activations"라고 부르며, 이를 해석하기 매우 어려울 것이라 예상합니다.

(이것이 바로 residual stream을 통해 발생하는 다양한 communication stream을 직접 연구하기보다는 virtual weights 측면에서 분리하려고 하는 주된 이유입니다.)

 

아마 residual stream 대역폭에 대한 높은 수요 때문에 일부 MLP neurons와 attention heads는 정보를 읽고 negative version을 쓰는 방식으로, 다른 layers에서 설정한 residual stream 차원을 지우는 일종의 "메모리 관리" 역할을 수행할 수 있다는 힌트를 발견하였습니다.

일부 MLP 뉴런은 입력과 출력 가중치 사이에 매우 negative 한 코사인 유사성을 가지고 이는 residual stream에서 정보를 삭제한다는 것을 나타낼 수 있습니다. 마찬가지로 일부 attention heads는 매우 큰 negative 고윳값이 큰 행렬에 추가되며 주로 현재 토큰에 관여하여 잠재적으로 정보를 삭제하는 메커니즘으로 작용할 수 있습니다.

이러한 메커니즘은 "메모리 관리"를 위한 일반적인 메커니즘일 수도 있지만, 일부 경우에만 작동하는 조건부로 정보를 삭제하는 메커니즘일 수도 있다는 점에 주목할 필요가 있습니다.

Attention Heads are Independent and Additive

위에서 살펴본 바와 같이 Transformer attention layers는 완전히 병렬적으로 작동하고 각각의 출력을 residual stream에
다시 추가하는 완전히 독집적인 여러개의 attention heads로 생각하면 됩니다.

그러나 이는 Transformer layers가 일반적으로 표현되는 방식이 아니며, 서로 동등하다는 것이 명확하지 않을 수 있습니다.

 

Transformer에 대한 원본 논문에서 attention layer의 출력은 결과 벡터를 쌓은 다음 출력 행렬을 곱하는 것으로 설명됩니다.

그 결과를 각 heads에 대해 동일한 크기의 블록으로 분할해 보겠습니다. 그런 다음 이를 관찰합니다:

각 heads를 독립적으로 실행하고 각각의 출력 행렬을 곱한 다음 residual stream에 더하는 것과 같다는 것을 알 수 있습니다.

연결 정의는 더 크고 계산 효율이 높은 행렬 곱셈을 생성하기 때문에 종종 선호됩니다.

하지만 transformer를 이론적으로 이해하기 위해서는 독립적으로 덧셈하는 것으로 생각하는 것이 좋습니다.

 

Attention Heads as Information Movement

하지만 attention heads가 독립적으로 행동한다면 어떤 역할을 할까요? attention heads의 기본적인 활동은 정보를 이동하는 것 입니다.

attention heads는 하나의 토큰의 residual stream에서 정보를 읽고, 이를 다른 토큰의 residual stream에 씁니다.

이 섹션에서 주목해야 할 점은 어떤 토큰에서 정보를 이동할지, 어떤 정보를 '읽기'하고 어떻게 목적지에 '쓰기' 하는지는

완전히 분리될 수 있다는 것입니다.

이를 확인하려면 attention을 non-standard 방식으로 작성하는 것이 도움이 됩니다. attention pattern이 주어졌을 때,
attention head를 계산하는 것은 일반적으로 3단계로 설명됩니다:

1. residual stream에서 각 토큰의 value vector를 계산합니다.
2. attention pattern에 따라 value vector를 선형적으로 결합하여 "result vector"를 계산합니다
3. 마지막으로, 각 토큰에 대한 head의 출력 벡터를 계산합니다.

이러한 각 단계를 행렬 곱셉으로 작성할 수 있는데, 이를 하나의 단계로 축소해 보겠습니다. X를 2차원 행렬(각 토큰에 대한 벡터로 구성)로
생각하면, 서로 다른 변에 곱하는 것입니다. "토큰당 벡터" 쪽을 곱하고 "위치" 쪽을 곱합니다.

Tensor는 이러한 종류의 행렬간의 구도를 설명하는데 훨씬 더 자연스러운 언어를 제공합니다

(텐서 곱 표기법이 익숙하지 않은 경우 표기법 부록에 간단한 소개가 포함되어 있습니다)

행렬에서 행렬로 선형 구도를 표현하고 싶다는 것이 도움이 될 수 있습니다. 수학자들은 이러한 선형 구조를 (2,2) tensor라고 부릅니다.

(두개의 입력 차원을 두개의 출력 차원에 매핑합니다) 따라서 Tensor는 이러한 변환을 표현하기 위한 자연어입니다.

 

Tensor 곱을 사용하면 attention을 적용하는 과정을 다음과 같이 설명할 수 있습니다:

 

mixed product 속성을 적용하고 identites를 축소하면 결과가 나옵니다:

attention pattern은 어떨까요? 일반적으로 key를 계산하여 query를 계산한 다음 각 key와 query 벡터의 내적으로부터
attention pattern을 계산합니다. 하지만 key와 query를 참조하지 않고도 이 모든 작업을 한 번에 수행할 수 있습니다.

$$ A = softmax(x^{T}W_{Q}^{T}W_{kx}) $$

 

이 공식은 수학적으로 동일하지만, 실제로 이런식으로 attention을 구현하는것은 (ie. multiplying by $ W_{O}W_{V} $ and $ W_{Q}^{T}W_{K} $) 끔찍하게 비효율적일 것입니다.

 

Observations about Attention Heads

attention heads를 이 형태로 다시 작설할 때의 가장 큰 장점은 이전에는 관찰하기 어려웠던 많은 구조가 드러난다는 점입니다:

• Attention heads는 한 토큰의 residual stream에서 다른 토큰으로 정보를 이동합니다
  • 이에 따른 결론은 residual stream 벡터 공간 (종종 contextual word embedding으로 해석되는)은 일반적으로 현재 토큰에 대한 직접적인 정보가 아니라 다른 토큰에서 복사된 정보에 해당하는 linear supspaces를 갖게 된다는 것입니다.
• Attention head는 실제로 서로 다른 차원에서 작동하고 독립적으로 작동하는  두개의 선형 연산, 즉 $ A $ 와 $ W_{O}W_{V} $를 적용하고 있습니다.
  • $ A $는 어떤 토큰의 정보가 어디로 이동하는지 관리합니다.
  • $ W_{O}W_{V} $는 소스 토큰에서 어떤 정보를 읽고 목적지 토큰에 어떻게 기록하는지를 관리합니다
• $ A $는 이 방정식에서 유일하게 비선형적인 부분입니다 (softmax에서 계산됨). 즉, attention pattern을 고정하면
  attention heads는 선형 연산을 수행합니다.
  이는 $ A $를 고정하지 않으면, 토큰당 선형 연산이 일정하기 때문에 어떤 의미에서 attention heads는 "half-linear" 합니다.
• $ W_{Q} $와$ W_{K} $는 항상 함께 작동합니다. 결코 독립적이지 않습니다.
  마찬가지로 $ W_{O} $와$ W_{V} $ 항상 함께 작동합니다.
  • 비록 별도의 행렬로 매개변수화되어 있지만, $ W_{O}W_{V} $와$ W_{Q}^{T}W_{K} $는 항상 개별적이고 low-rank 행렬로 생각할 수 있습니다.
  • key, query, value 벡터는 어떤 의미에서는 피상적인 개념입니다. 이들은 하위 행렬을 계산할 때 발생하는 중간 부산물 입니다.
    하위 행렬의 두 요소를 쉽게 재변수화하여 다른 벡터를 만들 수 있지만 여전히 동일하게 작동합니다.
  • 왜냐하면 $ W_{O}W_{V} $와$ W_{Q}W_{K} $는 항상 함께 작동하므로 이러한 결합된 행렬을 나타내는 변수를 정의하는 것이 좋습니다. $ W_{OV} = W_{O}W_{V} $ 와 $ W_{QK} = W_{Q}^{T}W_{K} $
• attention heads의 products는 attention heads 자체와 매우 유사하게 작동합니다.
  분배 속성에 의해, $ (A^{h2}\otimes W_{OV}^{h2})\bullet(A^{h1}\otimes W_{OV}^{h1}) =(A^{h2}A^{h1})\otimes (W_{OV}^{h2}W_{OV}^{h1}) $.
  이 product의 결과는 기능적으로 attention heads와 동등한 것으로 볼 수 있으며, attention pattern은 두 개의 heads $ A^{h2}A^{h1} $와 출력값 행렬 $ W_{OV}^{h2}W_{OV}^{h1} $의 구성입니다.
  이를 "virtual attention heads"라고 하며, 나중에 자세히 설명합니다.