[Paper Review] Algorithmic Capabilities of Random Transformers (NeurIPS, 2024)

TL;DR: Transformer는 전혀 학습하지 않아도 이미 어느정도의 Algorithmic Capability를 가지고 있다

 

Introduction

Transformer는 왜이렇게 강력한걸까?

Transformer기반 Language Model은 problem solving, string manipulation, in-context learning 등 고도의 reasoning이 필요한 작업을 잘 수행한다.

그런데, 어떻게 이게 가능한걸까? 이에 대해 저자들은 두 가지 가설을 세운다.

1. Transformer architecture makes these behaviors easy to learn
2. Transformer’s capabilities are already implemented in some fashion even in randomly initialized models

이 중 두번째 가설, 즉 Randomly initialized transformer의 성능을 조사하기 위해 Embedding layer만 학습된 모델을 검증한다.

만약 이렇게만 학습해도 괜찮다면, 이는 subspace에서 randomly initialized model의 행동이 이미 input-output mapping과 잘 대응한다는 뜻이고, 바꿔 말하면 target behaviour로의 encoding에 random initialize가 이미 충분하다는 뜻이다.

 

7가지 task에서의 실험 결과 Embedding-only training이 arithmetic, associative recall, and sequence generation같은 문제를 잘 풀 수 있다는 것을 발견했다. 이를 통해 저자들은 다음과 같이 주장한다

At least in transformers, some capabilities are available as soon as models are initialized, requiring only a learned encoding of inputs

 

이에대한 해석으로 저자들은 Embedding-only training이 low-dimensional subspace에서 input과 internal representation을 조정하는, 이른바 Subspace Selection으로 설명한다.

즉, Embedding-only training은 low-dimension에서 수행되는 algorithmic computation이라는 것이다.

 

본 연구는 기존에 수행되었던 initialized 딥러닝 네트워크의 행동 (혹은 성능)에 대해 조사한 연구들과 구분된다.

그들은 random network를 feature extractor로 보거나, lottery ticket*을 가진 sub-network를 찾아가는 과정으로 접근하였지만, 본 논문에서는 some useful capability는 transformer의 inductive bias덕분에 pruning이나 optimization없이 보유하고 있음을 주장하는 것에 가깝다

 

* Lottery Ticket Hypothesis 설명

The Lottery Ticket Hypothesis와 그 후속 연구들 리뷰

 

Setup

Models

모델은 decoder-only transformer language model을 가정한다.

 

Input $x$가 Embedding layer $E$를 통해 vector assigned되었을 때 ($h^{(0)} = E(x)$),

이 Embedding은 m개의 중간 layer들 $F^{(m)}$을 거친다.

$h^{(i)} = F^{(i)}(h^{(i-1)}) = \text{FFN}^{(i)}(\text{SelfAtt}^{(i)}(h^{(i-1)}))$

마지막 activation $h^{(m)}$은 **unembedding layer $U$**를 거쳐 next token의 distribution에 매핑된다.

 

Embedding layer와 Unembedding layer를 조금 더 자세히 살펴보면,

Embedding layer $E$는 token embedding matrix $E_{token}$와 positional embedding matrix $E_{pos}$로 구성된다

$h^{(0)} = E([x_1, x_2, ..., x_n]) = [E_{token}[x_1]+E_{pos}[1], \cdots, E_{token}[x_n] + E_{pos}[n]]$

 

Unembedding layer $U$는 그냥 single matrix로 parameterize된다

$p(x_{n+1}|x_n;E, F, U) \triangleq \text{softmax}(Uh_n^{(m)})[x_{n+1}]$

 

Initialization

FFN의 파라미터들은 평균 0, 표준편차가 $0.02/\sqrt{2n}$인 isotropic Gaussian에서 샘플링한다

나머지 파라미터들은 평균 0, 표준편차가 $0.02$인 Gaussian에서 샘플링한다

LayerNorm에서의 affine transformation은 Identity Matrix로 초기화한다

 

Training

중간 레이어 $F^{(1)} \sim F^{(m)}$을 모두 고정하고, $E$와 $U$만 학습한다 (이를 random transformers라 부른다)

 

Fully Training: $\underset{E, F, U} {\text{arg min}} \sum_{x, n \ge 0} -\log p(x_{n+1}|x_{1\cdots n};E, F, U)$

Embedding-Only Training: $\underset{E, U}{\text{arg min}} \sum_{x, n \ge 0} -\log p(x_{n+1}|x_{1\cdots n};E, F, U)$

argmin은 mini-batch stochastic gradient descent로 계산된다

 

Random Transformers Can Perform Simple Algorithmic Tasks

1) Tasks

Modular Addition

일정 modulus $p=199$를 두고 $(a+b)\; \text{mod} \;p$를 계산하라고 시키는 것.

일반적으로 이 task로 학습된 over-parameterized model은 Grokking* 현상이 발생되는 것이 관측된다

* Grokking 설명

딥러닝은 암기(Memorize)하지 말고 일반화(Generalize)되어야 한다는 인식에 반하는 논문.

소규모 알고리즘 데이터셋으로 학습하면 오버피팅을 일으키지만, 오버피팅인 상태로 compute를 계속 투입시키면 어느 순간 원리를 “깨닫는” 것처럼 Generalize 성능을 획득한다

Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets

Needle-in-a-Haystack

Markers $m_i$와 values $c_i$를 번갈아가며 계속 입력한 후 Qurey $m_u$를 입력했을 때 정답 $c_u$를 traking할 수 있는지를 보는 task.

예를 들어

$[m_1, c_1, m_2, c_2, \cdots, m_k, c_k, m_u]$이고, $m_u = m_2$라면 정답은 $c_2$이다

Decimal Addition

$39 + 71 = 110$ 처럼 텍스트를 입력으로 덧셈을 수행하는 Task이다.

이 때 two equal-length numbers만 사용하였으며, Input을 $9\; 3 \; 1 \; 7$으로, Output을 $0 \; 1 \; 1$으로 reverse해서 입력했다. 이 때 십진법만을 고려하였다

Parenthesis Balancing (Dyck Recognition)

괄호가 balancing한지 판단하는 Task.

이 때는 입력이 괄호와 Label밖에 없기 때문에 Vocabulary size가 4이다. 즉, 학습되는 파라미터가 엄청나게 적다

 

2) Results

Random transformers learn to perform all four tasks

• 1024차원의 random transformer는 4개의 task를 완벽하게 수행한다
• 이는 transformer-specific inductive bias가 embedding-only training에 효과적이기 때문이다

 

In general, embedding and unembedding parameters must both be trained

• Ablation 같은 느낌이다.
• token embeddings & positional embeddings & unembedding이 모두 optimized 되어야 한다

 

Random transformers exhibit interpretable attention patterns

• Attention weight를 전혀 학습하지 않았는데도, fully-trained가 가지고 있는 induction heads 패턴을 보인다

 

Random transformers use structured embeddings

• (In the modualr addition task) fully-trained가 low-dimensional subspace에서 circular embedding을 가지는 것처럼, random transformer도 그렇다

 

Random Transformers Can Memorize and Generate Structured Sequences

더 어려운 task인 memorization이나 free-form text generation에도 적용 가능할까?

 

1) Memorization

Two-integer key에서 one-integer value를 매핑하는 Task이다. 이전 Task와 다르게 현 Task에서 learned function은 반드시 embedding 단위에서 이루어져야 하며, 이는 embeddings에서 얼마나 효율적으로 정보가 저장되는지를 가늠할 수 있다

 

Memorization 능력은 Randomly initialized blocks에서는 발현되어있지 않다

(이 논문의 제목이 Algorithmic capabilities인 이유이기도 하다)

 

2) Language Modeling

LLM의 핵심 능력으로써 Natural language modeling은 memorization of arbitrary associations과 structured sequence generation procedures를 모두 요구하는 어려운 task이다.

 

Random transformer는 Normal과 달리 layer를 추가한다고 Loss가 줄어들지 않는다 (transformer block의 능력이 아니다)

 

 

Perplexity는 당연히 낮지만, 저자들은 몇가지 sub-task를 수행하는 것이 보인다고 주장한다.

1. Generated sentences are generally grammatical (e.g. Max didn't know what to do)
2. Topically coherent with the prompt
3. Correctly resolve some long-distance dependencies (e.g. references to the name Max)
4. (perheps) High-level narrative structure

 

Random Transformers Operate in Low-Dimensional Subspaces

이 Section에서는 앞서 수행한 task를 성공한 이유를 설명한다.

저자들은 Embedding-only training이 이미 target function이 구현되어있는 low-dimensional subspaces에서의 hidden computation을 최적화한다고 주장한다.

이는 sparsification과는 구분되는 Subspace Selection이라 부른다

 

• algorithmic task에서는, Normal transformer와 Random transformer 모두 low-dimensional subspaces에서 work 하기 때문에 task를 풀기 충분하다고 설명한다
• 그러나 memorization과 language modeling의 경우,  Random transformer는 더 많은 subspace selection을 보이기 때문에 성능 저하를 보인다고 한다.
• 마지막 task는 반드시 high-dimensional space에서 이루어져야 하는 circuit imitation이라는 task로 검증한다. 이는 Normal과 Random의 차이를 분명하게 보여줄 것이다

 

1) Low-Dimensional Hidden Representations in Algorithmic and Memorization Tasks

이 부분은 transformer representation의 geometry를 조사하기 위해 각 Layer 단계에서 분석한 부분이다.

이러한 representation에 대해 두 가지 가설을 검증한다.

1. Representation이 low-dimensional subspace에 존재한다면, top Principal components의 Fraction of variance*가 높을 것이다 (Behaviour: Subspace selection)
2. Representation이 sparse sub-network에 의해 결정된다면, Neuron activation의 Fraction of variance가 높을 것이다 (Behaviour: Sparsification)

*Fraction of variance??

한마디로, 회귀 문제에서 종속 변수 Y가 입력 변수 X에 의해서 설명되는 분산의 비율을 의미한다

Explained variation - Wikipedia

• 4가지 algorithmic tasks에서는 Principal Component의 비율이 월등히 높다 (가설 1: Subspace selction)
• Memorization에서 Random은 Normal보다 작은 subspace에 더 집중하는 것으로 보인다.

• Language modeling에서 representation 중 30% 이상의 variance가 10 components로 설명된다

 

정리하자면, Transformer는 선언되는 순간 '일부 능력' (=target function)을 가지고 있는 것이다.

Embedding-only training (Random transformer)은 이미 Initialized weights로인해 정해져있는 
target function을 찾아가는 Subspace selection을 최적화한다.

또한 이 target function은 Low-dimensional subspace에서 동작하며,
Algorithmic task의 경우 이 것으로 target function을 찾아가기 충분하다

하지만 Memorization이나 Language modeling의 경우 Neuron activation도 고려하여 Reproesentation을 
추출하는 것이 필요하지만, Random transformer는 이 부분이 부족하고 따라서 더 작은 Subspace에 집중하는
것 처럼 보인다

 

2) Subspace Selection in Circuit Imitation

이 단계에서는 low-dimensional subspace에서 동작하는 random circuit을 simulate하기 위해 얼마나 큰 Hidden Representation을 가져야 하는가를 분석한 부분이다

 

실험과정은 다음과 같다.

단순하게 분포 $\tilde{p}$를 가지는 Random target tansformer를 구축하고, 또다른 Random transformer를 embedding-only training하여 $\tilde{p}$ 를 모사하게끔 학습한다

$\underset{E, U}{\text{arg min}} \; \mathbb{E}_x[\textbf{KL}(\tilde{p}(\cdot|x)||p(\cdot|x;E,F,U)]$

 

Target model $\tilde{p}$을 초기화한 후 Attention parameter는 10으로, FFN weights는 20으로, Final projection layer는 $100\sqrt{width}$로 나눠줬다 (이는 Attention pattern의 다양성을 증가시킨다고 한다)

• 일반적으로, random transformer는 shallower (3 vs 1) & narrower (512 vs 128) 모델과만 잘 매치된다
•  Target model의 dimension을 12 -> 16 -> 32로 갈수록 급격히 Error가 증가한다, 이는 곧 Random transformer는 lower-dimensional subspace에서의 연산밖에 배우지 못하는 것일지도 모름을 시사한다.

 

Discussion

DNN은 Winning ticket을 갖고 있다는 기존 논문의 주장과는 반대되는 입장의 논문인 듯 하다

 

적어도 Transformer는 Sub-network가 아니라 Subspace selection을 통해 지식을 encapsulate한다는 것이 본 논문의 주장이고, 이는 Transformer의 inductive bias 덕에 embedding-only training만으로 가질 수 있다.

 

그렇기 때문에 Random Transformer는 Algorithmic capabaility를 가지고 있는 것임을 실험으로 증명했고, 이는 Fraction of variance 분석과 Circuit imitation 실험을 통해 low-dimensional subspace에서 수행됨을 확인했다

 

Referecnes

Paper:

Algorithmic Capabilities of Random Transformers

 

Project:

GitHub - fjzzq2002/random_transformers: Official code for "Algorithmic Capabilities of Random Transformers" (NeurIPS 2024)