논문/Interpretable AI

[논문] Toy Models of Superposition

guungyul 2025. 3. 27. 14:18

Neural network들은 많은 연관되지 않은 개념들을 하나의 neuron에 mapping한다. 이 현상은 "polysemanticity"로 알려져 있고 모델의 해석을 매우 어렵게 하는 요소 중 하나이다. 본 논문에서는 이 polysemanticity를 완전히 이해할 수 있는 toy model을 제시한다. 이는 모델이 추가적인 feature들을 "superposition"으로 저장한 결과이다.

 

 

모든 figure들과 설명은 아래 논문에서 추출됐다.

https://arxiv.org/abs/2209.10652

 

Toy Models of Superposition

Neural networks often pack many unrelated concepts into a single neuron - a puzzling phenomenon known as 'polysemanticity' which makes interpretability much more challenging. This paper provides a toy model where polysemanticity can be fully understood, ar

arxiv.org

 

 

 

Introduction

Neural nework들의 neuron들이 이해 가능한 feature 하나에 대응된다면 매우 편리할 것이다.

  • 예를 들어 각 neuron이 빨간색, 강아지 코 등 이미지의 특정 feature에만 반응하는 경우이다.
  • 하지만 neuron들은 항상 이런 행동을 보이지 않는다. 왜일까?

본 논문에서는 멀리 떨어진 input feature들을 가지는 합성 데이터로 학습한 작은 ReLU network로 모델이 언제 그리고 어떻게 dimension보다 많은 feature를 나타내는지 조사한다.

  • 우리는 이러한 현상을 "superposition"이라 명명한다.
  • Superposition은 linear model이 할 수 있는 compression을 보다 늘려주고 대신 nonlinear filtering을 필요로하는 "간섭"을 만들게 된다.

5개의 다른 중요도를 가진 feature들을 2D로 embedding하는 toy model을 고려해보자.

  • Dense feature(entry가 0인 feature가 거의 없음)들을 사용하면 가장 중요한 두가지 feature에 대한 orthogonal basis를 표현하도록 학습한다.
  • 하지만 sparse feature(entry가 0인 feature가 많음)들을 사용할수록 이는 아래 그림과 같이 바뀌게 된다

  • 즉, 모델이 약간의 interference을 허용하면서 추가적인 feature들을 저장할 뿐 아니라
  • 가끔 superposition 상태에서도 계산을 하기도 한다

이는 neural network들이 사실은 약간의 오류가 있는 훨씬 더 크고 sparse한 network를 흉내내고 있다는 가설을 세우게 한다.

 

본 논문의 기여점은 다음과 같다.

  • Superposition이 단순히 post-hoc interpretation에 쓰이는 것이 아니라 neural network의 "ground truth"라는 사실을 입증한다.
  • Superposition이 언제, 왜 일어나는지 가설들을 제시하고 superposition의 phase diagram을 보여준다.
  • Toy model을 통해 superposition이 복잡한 기하학 구조를 가진다는 것을 보여준다.

추가적으로 toy model을 통해

  • phase change, 정다면체에 기반한 기하학 구조, 학습 과정에서 "energy level" 같은 jump, fractional quantum Hall effect와 비슷한 현상 등을 조사할 수 있었다.

 

하지만 toy model들이 단순한 ReLU network 구조를 가지기 때문에 다른 큰 network들로 일반화 된다고 하기 어렵다.

 

 

 

Definitions and Motivation: Features, Directions, and Superposition

Emperical Phenomena

"Feature"가 무엇이고 이들이 어떻게 표현되는지를 논의할 때, 관측된 여러 경험적인 현상에 기반하여 이론을 구축한다.

  • Word Embeddings
    • 하나의 semantic 속성에 대응하는 방향이 존재한다
    • V(king) - V(man) + V(woman) = V(queen)
  • Latent Spaces
    • 비슷한 "vector arithmetic"이 generative adversarial networks들에서 관측된다
  • Interpretable Neurons
    • RNNs, CNNs, GANs 모델들에서 특정 속성에만 반응하는 해석 가능한 neuron들을 찾을 수 있다
  • University
    • 같은 속성에 반응하는 비슷한 neuron들이 network 전역적으로 관측된다.
  • Polysemantic Neurons
    • 반대로 전혀 상관 없는 input의 혼합에 반응하는 polysemantic neuron들도 존재한다.

따라서 우리는 neural network의 representation들이 방향으로 표현되는 feature들로 구성되어있다 생각한다.

다음 section들에서 이 아이디어를 더 자세히 탐구한다.

 

 

What are Features?

Feature을 정의하기 위해 세 가지 가능한 정의들을 제시한다.

  • Features as arbitrary functions
    • Feature를 input에 따른 함수로 정의할 수 있다.
    • 하지만 이는 본 논문의 동기와는 거리가 있다.
  • Features as interpretable properties
    • Feature를 사람이 이해 가능한 개념들로 정의할 수 있다.
    • 하지만 이해할 수 없는 개념도 feature가 될 수 있어야 한다.
  • Neurons in Sufficiently Large Models
    • 최종적인 접근법으로 feature를 충분히 큰 모델이 하나의 neuron을 할당해 표현할 input의 속성으로 정의한다.
    • Polysemantic neuron에서 관측되는 이해 가능한 속성들이 충분히 큰 모델에서는 각각 하나의 neuron에 할당될 것이라 기대한다.

 

 

Features as Directions

위 section에서 언급되었듯이 우리는 feature들이 방향으로 표현된다고 생각한다.

 

Feature들이 activation space에서의 방향과 일치하다면 이를 neural network의 representation이 linear하다고 부른다고 하자.

  • 즉, linear representation에서 각 feature $f_i$는 대응하는 방향 $W_i$를 가진다.
    • Linear representation은 간단하게 각 feature들의 방향의 linear combination으로 나타내진 representation이라 생각하면 된다.
    • 명확히 해야될 것은 feature들은 input의 nonlinear 함수로 표현된다. Feature들을 activation vector로 맵핑하는 것이 linear 한 것이다.

Neural network의 대부분은 linear 함수들로 이루어져 있다.

  • 즉, linear representation이 neural network에서 정보를 저장하는 자연스로운 형식인 것이다.

 

여러 layer를 사용하다면 non-linear representation을 구축할 수는 있다.

  • 하지만 우리의 직관은 non-linear representation이 일반적으로 더 효율적이지 못하다는 것이다.

Linear representation을 사용하면 차원 수 만큼의 feature들만 저장할 수 있다고 생각할 수 있다.

  • 하지만 superposition을 통해 더 많은 feature들을 linear하게 표현 가능하다.

 

 

Privileged vs Non-previleged Bases

Feature들이 방향으로 encoded 되어있다면 어떤 방향일까? 또 어떤 경우에는 basis 방향을 고려하는게 도움이 될 때도 있고 아닐 때도 있다. 왜 그럴까?

 

Word embedding의 경우 basis 방향이 의미를 가지지 않는다. 하지만 neural network layer들의 경우 그렇지 않다.

  • Activation 함수를 적용하게 되면 "대칭이 어긋나게"된다.
  • 이는 basis 방향들을 특별하게 만들고 feature들이 basis 차원들과 일치하도록 장려한다.

우리는 basis dimension과 일치하는 basis를 privileaged basis라 부르고 basis 방향들을 "neurons"라 명명한다.

  • 일반적인 경우 이런 neuron들은 해석 가능한 feature들과 대응한다.

 

 

The Superposition Hypothesis

Privileged basis가 존재한다고 해도 neuron들은 보통 "polysemantic"하다.

  • 즉, 다양한 연관되지 않은 featuer들에 대응된다.

이는 superposition hypothesis로 설명될 수 있다.

  • 간단하게 "이미 존재하는 neuron들보다 더 많은 feature들을 나타내고 싶기" 때문에 high-dimensional space의 속성을 이용해 많은 neuron들을 가진 model을 흉내내는 것이다.

 

다음과 같은 수학적 배경이 superposition을 설명할 수 있다.

  • Almost-orthogonal한 vector들
    • N 차원 공간에 N개의 orthogonal한 vector들을 가질 수 있다.
    • 하지만 "거의 orthogonal"한 vector를 사용하면 N 차원에 exp(N)개의 vector를 나타낼 수 있다.
      • Johnson-Lindenstrauss lemma
  • Compressed sensing
    • 일반적으로 vector를 저차원 공간으로 project하면 기존 vector를 복원할 수 없다.
    • 하지만 기존 vector가 sparse하다는 성질을 이용하면 기존 vector를 복원할 수 있다.

 

구체적으로 superposition hypothesis에서는 feature들이 almost-orthogonal한 방향들로 나타내진다.

  • 따라서 하나의 feature가 활성화 될 때 다른 feature들도 함께 약간 활성화 되는 것처럼 보일 수 있다.
    • 이런 "noise"나 "interference"는 network가 견뎌야하는 하나의 cost이다.
  • Neural network가 매우 sparse한 feature들로 이루어져 있다면 더 많은 feature들을 나타내는 이점이 이런 cost보다 커질 수 있다.

 

이를 이해하는 다른 방법은 작은 neural network가 noise와 함께 spare한 더 큰 모델을 "흉내낼" 수 있다는 것이다.

 

 

 

Summary: A Hierarchy of Feature Properties

이 섹션에서의 핵심 아이디어는 neural network representation들이 가지는 네 가지 속성으로 생각될 수 있다.

  • Decomposability
    • Decomposable한 neural network activation들은 feature들로 decompose 될 수 있다.
    • 이때 feature들은 서로 독립적인 의미를 가진다.
  • Linearity
    • Feature들은 방향에 대응된다.
    • 각 feature $f_i$는 방향 $W_i$에 대응한다.
    • 여러 feature가 각각 activate 값 $x_f_1, x_f_2, ...$를 가지고 있다면 이는 $x_f_1W_f_1+x_f_2W_f_2...$로 나타내진다.
  • Superposition vs Non-Superposition
    • $W^TW$가 invertible 하지 않을 때 linear representation은 superposition을 보인다.
    • Invertible하다면 superposition을 가지지 않는다.
      • Invertible -> W vectors가 linearly independent -> 각 input feature가 unique direction을 가질 수 있음
  • Basis-Aligned
    • 모든 $W_i$가 one-hot basis vector들이라면 representation은 basis aligned하다.
    • 만약 모든 $W_i$가 sparse하다면 representation은 부분적으로 basis aligned 하다.
    • 이는 privileaged basis를 필요로 한다. (하나의 기준점을 통해 sparse한지 dense한지 one-hot한지 정의한다)

첫번째 두 속성들은 널리 적용되지만 나머지 두 속성들은 항상 관측되지 않는다.

 

 

 

Demonstrating Superposition

Superposition hypothesis를 검증하기 위해서는 neural network들이 실제로 neuron 수보다 많은 feature들을 나타낼 수 있는지 확인해야 한다.

  • Linear 모델들만으로는 불가능하다.
  • 하지만 작은 nonlinearity를 추가하면 완전히 다른 방식으로 작동하는 것을 알 수 있다.

 

 

Experiment Setup

목표는 neural network가 high dimensional vector $x \in R^n$을 lower dimensional vector $h \in R^m$에 project하고 복원할 수 있는지 탐구하는 것이다.

 

The feature vector (X)

먼저 high-dimensional vector $x$를 묘사한다.

  • 이는 이상적인, disentangled 된 더 큰 모델의 activation들이다.
  • 각 element $x_i$는 "feature"라 부른다.
    • 이는 각 feature가 더 큰 모델의 각 neuron들과 완벽하게 일치한다 가정하기 때문이다.

Feature들에 대한 ground truth가 없기 때문에 $x$를 위한 합성 데이터를 생성해야 한다.

이때 세가지 가정을 따른다.

  • 1) Feature Sparsity
    • 많은 feature들은 가끔 발생하기 때문에 매우 sparse하다.
      • 예를 들어 vision에서 강아지 꼬리나 언어에서 Martin Luther King을 나타내게 되는 경우는 매우 드물다.
    • 따라서 feature들이 sparse한 분포를 가진다 가정한다.
  • 2) More Features than Neurons
    • 모델이 잠재적으로 표현 가능한 feature는 neuron에 비해 훨씬 더 많다.
  • 3) Features Vary in Importance
    • 주어진 작업에서 모든 feature들이 동일하게 중요하지 않다.
      • 예를 들어 강아지 종을 분류하는 작업에서 귀 모양이 다른 feature들에 비해 훨씬 더 중요하다.

 

구체적으로 합성 데이터는 다음과 같이 정의된다.

  • Input vector $x$는 실제 underlying feature를 나타내는 합성 데이터이다.
  • 각 차원 $x_i$는 "feature"라 생각한다.
    • 각 차원은 sparsity $S_i$와 importance $I_i$를 가진다.
    • $S_i$의 확률로 $x_i=0$으로 설정한다. 그 외의 경우 [0, 1] 값으로 균일하게 분포한다.

 

 

The Model (X -> X')

크게 두가지 모델을 고려한다.

  • Linear model
    • Superposition을 보이지 않는 baseline이다.
  • ReLU output model
    • Superposition을 보이는 간단한 모델이다.

 

왜 이런 모델을 쓰는 것일까?

 

Superposition hypothesis는 high-dimensional model의 각 feature가 lower-dimensional space의 하나의 방향에 대응한다고 가정한다.

즉, down projection을 $h=Wx$인 linear map으로 나타낼 수 있음을 뜻한다.

  • 이때 $W_i$의 각 column은 lower-dimensional space에서 feature $x_i$를 나타내는 방향에 대응한다.

기존 vector를 복권하기 위해서 같은 matrix의 transpose인 $W^T$를 사용한다.

  • 이는 수학적 원칙을 따르고 실제로 작동한다.

그 후 bias도 추가한다.

  • 이는 모델이 예상 값으로 표현되지 않는 feature들을 고치도록 도와준다.
  • 또 나중 section에서 탐구되듯이 negative bias는 약간의 noise를 없애도록 도와준다.

마지막 단계는 activation을 추가할지 여부이다.

  • 이는 superposition이 나타나는지를 결정하는 주요한 요소이다.

 

 

The Loss

Loss는 feature importance $I_i$로 가중된 mean squared error이다.

  • $L = \sum_{x} \sum_{i} I_i (x_i - x_i')^2$

 

 

Basic Results

첫번째 실험으로 다른 sparsity level을 가지는 여러 ReLU output 모델들을 학습시키고 그 결과를 시각화한다.

 

가장 간단한 시각화 방법은 $W^TW$와 $b$를 사용하는 것이다.

  • 시각화는 feature의 중요도에 따라 정렬한다
  • 아래는 작은 함수 (n = 20; m = 5;)의 시각화 결과이다.
    • Dimension 수 만큼의 feature를 표현하는 linear model처럼 행동한다.

 

Superposition hypothesis를 검증하기 위해 확인해야 하는 부분은 모델이 정말 추가적인 feature를 almost-orthogonal하게 표현하는 지이다.

  • 먼저 간단한 질문으로 얼마나 많은 feature를 모델이 학습하는지 질문할 수 있다.
    • 이는 embedding vector의 길이 $||W_i||$로 결정된다. (아래 그림에서 x축)

또한 주어진 feature가 다른 feature들과 차원을 공유하는지 확인하고 싶다.

  • 이는 다른 모든 feature들을 vector $W_i$의 방향으로 project한다.
    • $\sum_{j \ne i} (\hat{W}_i \cdot W_j)^2$
    • 다른 feature들과 모두 orthogonal 하다면 0
    • 1보다 큰 값은 다른 feature들이 $W_i$를 feature $i$ 만큼 강하게 활성화한다는 뜻이다. (아래 그림에서 색)

위 모델의 경우 아래 처럼 시각화할 수 있다.

 

적은 수의 feature들 ($n = 20; m = 5; I_i = 0.7^i$)를 가진 모델을 고려해보자.

  • 아래는 서로 다른 feature sparsity level에 따른 시각화 결과이다.

 

결과 분석

  • Linear model은 항상 top-m 가장 중요한 feature만을 학습힌다.
  • ReLU 모델의 경우 dense feature들은 linear model과 비슷하게 동작하지만 sparsity가 증가할수록 superposition이 생기는 것을 볼 수 있다.
    • 즉, 더 많은 feature들을 나타내기 위해서 모델이 feature들을 서로 orthogonal하게 두지 않는다.
    • 가장 중요도가 낮은 feature부터 천천히 가장 중요한 feature까지 적용된다.
    • 처음에는 feature들을 antipodal pair(반대 방향을 가지는 쌍)로 정렬하다가
    • 다른 기하학 구조를 가지게 변화한다.

위 결과는 더 많은 feature과 hidden dimension에도 동일하게 적용된다. ($n = 80; m = 20; I_i = 0.9^i$)

 

 

Mathmatical Understanding

Superposition이 일어나는 이유는 무엇인가? 또 왜 간단한 non-linearity를 추가하는 것 만으로 linear model과 큰 차이를 만드는 걸까?

그 답은 모델이 두 가지 경쟁하는 요소들 - feature benefit과 interference - 로 좌우되기 때문이다

 

Linear neural network의 weight들은 간단한 closed-form solution에 최적화한다 생각할 수 있다.

  • 이를 우리의 linear case로 바꾸어 식으로 나타내면

  • Feature benefit: 모델은 더 많은 feature를 나타냄으로써 더 낮은 loss를 얻을 수 있다.
  • Interference: 하지만 orthogonal하게 나타낼 수 있는 feature들보다 더 많은 feature를 나타냄으로써 feature 간 interference가 증가한다.
    • 이는 linear model이 dimension보다 더 많은 feature들을 나타내는 것을 가치없게 만든다.

ReLU 모델에도 비슷한 이해방식을 적용할 수 있다.

  • 구체적으로 이해하고 싶은 식은 확률 S로 $x_i = 0$이 분포되어있는 $L = \int_{x} |I(x - \text{ReLU}(W^T W x + b))|^2 , dp(x)$ 이다.
    • x로 적분을 하는 것은 sparsity에 기반해 각 term으로 분해할 수 있다. 이는 binomial expansion $((1-S)+S)^n$의 형태를 가진다.
    • 따라서 Loss를 다시 쓰면
      • $L = (1 - S)^n L_n + \cdots + (1 - S) S^{n-1} L_1 + S^n L_0$
        • 각 $L_k$는 input이 k개의 term을 가진 sparse vector일 때의 loss를 뜻한다.
        • S가 1에 가까워질수록 $L_1$과 $L_0$이 주가 된다.
        • $L_0$은 단순히 postive bias의 loss이기 때문에 $L_1$에 집중한다.

위 식은 화학에서 사용되는 Thomson problem과 매우 유사하다.

  • 전자들을 반지름이 1인 구 위에 서로 거리가 가장 멀게 배치하는 방법
  • 특히 고정된 수의 feature들의 $||W_i|| = 1$이고 나머지는 $||W_i|| = 0$을 가진다 가정하면 feature benefit term과 interference eterm은 Thomson problem의 일반화 term과 동일하다.

또 다른 흥미로운 점은 1-sparse case에서 ReLU가 negative interference를 없애준다는 사실이다.

  • 이는 관측된 solution들에서 가능할 때 왜 항상 negative inference만이 존재했는지 설명해준다.
  • 추가적으로 negative bias를 사용하면 positive inference도 negative로 바꿔줄 수 있다.

다른 less sparse vector들은 독자들이 탐구할 주제이다.

  • 핵심 아이디어는 합성된 여러 interferences들이 존재하고 "active feature"들은 interference를 경험한다는 사실이다.
    • 실제로는 아주 적은 feature들만 서로 interfere한다. 때문에 이는 합성되는 interferences의 확률을 낮추고 1-sparse loss term을 보다 중요하게 한다.

 

 

 

Superposition as a Phase Change

위 결과는 모델이 학습될 때 feature들은 세 가지의 결과를 가진다고 제시한다.

  • 1) Feature가 단순히 학습되지 않는다.
  • 2) Feature가 학습되고 superposition으로 나타내진다.
  • 3) Feature가 학습되고 하나의 정해진 차원으로 나타내진다.

위 세 가지 결과 사이 transition은 분명하다. 즉, 그 사이 어떤 형태의 phase change가 존재할 수 있다.

 

물리에서의 "phase diagram" 같은 분명한 경계를 찾기 위해서 각 영향들을 분리한 실험을 진행한다.

  • 초기 실험처럼 2개의 feature과 하나의 hidden layer 차원을 사용한다.
  • ReLU output 모델 $ReLU(W^TWx-b)$를 사용한다.
  • 첫번째 feature의 중요도는 1로 설정한다.
    • 두번째 "extra" feautre의 중요도를 0.1 ~ 10으로 조절하고
    • Sparsity 또한 1 ~ 0.01로 조절한다.
  • 실험을 통해 "extra" feature가 세 가지 결과 중 어느 결과를 가지는지 실험한다.
  • Noise를 줄이기 위해 각 point마다 10개의 모델을 학습하고 가장 높은 loss값을 뺀 값들의 평균 값을 결과로 사용한다.

 

2개의 feature를 1 차원에 저장하는 방법은 세 가지가 있다.

  • 1) $W = [1, 0]$: [0, 1]을 무시한다. 즉, extra feature를 학습하지 않고 첫번째 feature에 차원을 할당해준다.
  • 2) $W = [0, 1]$: [1, 0]을 무시한다. 즉, 첫번째 feature를 학습하지 않고 두번째 feature에 차원을 할당해준다.
  • 3) $W = [1, -1]$: 두 가지 feature를 superposition으로 학습한다. 대신 [1, 1]을 나타낼 수 없다.
    • 이 마지막 3번째 solution을 "antipodal"이라 부른다. 두 basis vector [1, 0], [0, 1]이 서로 다른 방향으로 mapping되어있기 때문이다.

예상했듯이 superposition이 일어나기 위해서는 sparsity를 필요로 한다.

  • 또한 이 phase를 넘어갈 때 weight가 discontinuous하게 바뀐다. (Optimal loss function의 미분 값도 discontinous하다)

 

위와 동일한 방법으로 3개의 feature를 2차원에 embedding하는 방법을 고려할 수 있다.

이론적인 모델로 4가지 방법이 존재한다.

  • 이는 W가 "무시한" feature direction을 통해 할 수 있다.
    • 예를 들어 $W \perp [0, 0, 1]$은 [0, 0, 1]인 마지막 feature를 W가 무시했다는 것을 의미한다.
  • 또 한 가지 흥미로운 점은 모델이 superposition으로 antipodal pair를 만드는 방법이 두 가지 존재한다는 것이다.
    • 예를 들어 $W \perp [1, 1, 0]$는 첫번째와 두번째 feature를 superposition으로 두고 extra feature에 하나의 차원을 배정한다.
    • 또는 $W \perp [0, 1, 1]$처럼 extra feature와 다른 feature 둘 중 하나의 superposition으로 둘 수 있다.
  • 모든 feature들이 superposition으로 있는 $W \perp [0, 1, 1]$의 경우는 고려하지 않는다.

위 다이어그램은 phase change가 정말 존재함을 보여준다. 다음 섹션에서 관측되지 않은 더 복잡한 구조가 존재하는지 탐구한다.

 

 

 

The Geometry of Superposition

위에서 sparsity가 증가할수록 표현되는 extra feature들이 증가한다는 것을 확인했다.

Feature들은 사실 오각형이나 사면체 같은 기하학 구조를 가지게 된다.

이 섹션에서는 먼저

  • uniform superposition을 탐구한다.
    • 모든 feature들이 동일하다: 독립적이고, 똑같은 중요도와 sparsity를 가진다.
    • 그 결과 uniform superposition이 정다면체의 모양을 가지게 된다.
  • 그 후 non-uniform superposition을 탐구한다.
    • 모든 feature들이 동일하지 않다.
    • Uniform superposition의 변형으로 어느 정도 이해할 수 있다.

 

 

Uniform Superposition

Sparsity S의 영향을 탐구하려고 한다.

  • 모든 feature는 중요도 $I_i = 1$을 가진다.
  • n = 400 feature, m = 40 hidden dimension을 사용한다.
    • n이 m보다 충분히 크기만 하면 그 숫자는 크게 중요하지 않다.
    • m이 두배가 되면 학습되는 feature도 2배가 된다.

 

Feature가 학습되었는지 확인하는 간단한 방법은 Frobenius norm $||W||^2_F$를 사용하는 것이다.

  • 1과 유사하거나 같은 값은 학습됨, 0과 비슷하거나 같은 값은 학습되지 않음을 뜻한다.

아래 그래프는 "dimension per feature"인 $D^* = m / ||W||^2_F$를 나타낸다.

 

놀랍게도 그래프가 1과 1/2 근처에서 "sticky"하다는 것을 발견했다.

  • 이는 antipodal pairs가 되는 정확한 기하학 배치와 1/2 "skicky point"가 대응하는 것처럼 보인다.
    • 즉, antipodal pairs가 효과적이기 때문에 모델이 더 많은 sparsity regime에서 사용하는 것으로 추측할 수 있다.

 

Feature Dimensionality

하나의 feature가 가지는 "fraction of dimensionality"를 표현할 방법이 있을까?

i 번째 feature의 dimensionality는 다음과 같이 정의한다.

  • $D_i = \dfrac{|W_i|^2}{\sum\limits_{j} (\hat{W}_i \cdot W_j)^2}$
    • $W_i$는 i 번째 feature와 연관된 weight vector
    • $\hat{W}_i$는 위 vector의 unit vector
  • 분자는 주어진 feature가 표현된 정도, 분모는 얼마나 많은 feature들이 dimension을 공유하는지를 나타낸다.
  • 이론적으로 "packed efficiently"라면 모든 feature들의 dimensionality 합은 embedding dimension의 수가 될 것이라 유추할 수 있다.

이제 위 그래프를 per-feature basis로 분리해 나타낸다.

  • 위 섹션의 line plot으로부터 시작한다
  • 각 feature의 feature dimensionality를 나타내는 scatter plot을 추가한다.
  • 이러한 feature dimensionality가 cluster하는 부분에 선을 추가한다.
  • 몇몇 모델의 weight geometry들을 "feature geometry graph"로 시각화한다.
    • Node는 feature를, edge weight는 feature embedding vector들의 dot product의 절대값이다. 즉, orthogonal하지 않다면 연결되어 있다.

 

이전 section에서는 superposition을 phase change로 보는 이론을 발전시켰다.

하지만 위 그래프에서 나타나듯이 0과 1 사이 모든 결과는 superposition이다.

  • 즉, superposition 안에서도 여러 서로 다른 phase들을 가진다는 것을 의미한다.

 

Why These Geometric Structures?

많은 구조들은 Thomson problem의 solution들이다.

  • 즉, 모델이 feature를 표현할 때 feature는 m-dimensional sphere 위의 점으로 embed된다.

 

또한 uniform polyhedra인 (ex. tetrahedron) 경우 하나의 선만 존재하지만 non-uniform solution의 경우 두개의 line이 공존한다 (ex. 2/3과 1/2).

  • Uniform의 경우 모든 꼭짓점이 같은 기하학적 특징을 가지고 같은 dimensionality를 공유한다.
  • Non-uniform의 경우 다른 feature들이 서로 더 또는 덜 interfere 한다. (다른 dimensionality를 가진다)

특히, Thomson solution들은 더 작은 uniform polytopes들의 tegum product로 이해될 수 있다.

  • 즉, non-uniform solution의 경우 우리는 실제 dimensionality가 underlying factor uniform polytopes에 대응한다 생각해야 한다.

 

또한 위 설명은 우리가 high dimensional version을 연구하는데 왜 3D Thomson problem solution을 관측할 수 있는지 설명한다.

  • 3D Thomson solution이 2D와 1D의 tegum product인것처럼 몇몇 higher dimensional solution들은 1D, 2D, 3D tegum product로 나타내질 수 있다.

또한 tengum products의 orthogonality는 또 하나의 사실을 말해준다.

  • Superposition 관점에서 tengum-factor들 사이에는 "interference"가 존재할 수 없음을 의미한다.

 

 

Non-Uniform Superposition

위 섹션에서는 모든 feature들이 같은 중요도, sparsity, 독립적인 uniform superposition을 탐구했다. 그리고 이 모델은 Thomson problem의 변형을 해결하려 한다.

이 섹션에서는 non-uniform superposition을 연구한다. Feature들은 중요도와 sparsity가 다르고 서로 의존적일 수 있다.

 

실제 neural network들은 이런 non-uniform한 superposition을 보인다. 하지만 아직 이런 non-uniform geometry를 어우르는 이론을 형성하는 것과 거리가 있다.

따라서 몇가지 중요한 현상들만 제시한다.

  • Feature의 중요도와 sparsity가 변화하면 polytope에 부드러운 deformation이 생기다가 특정 지점을 기점으로 다른 polytope로 변화한다.
  • Correlated feature들은 다른 tegum factor로 형성되며 서로 orthogonal한 것을 선호한다.
    • 그 결과로 corrleated feature들은 orthogonal local basis를 형성한다.
    • Orthogonal 할 수 없다면 side-by-side를 선호한다.
    • 또는 correlated feature들이 하나로 통합되는 경우도 있다.
  • Anti-correlated feature들은 같은 tegum factor에 있는 것을 선호한다.
    • Negative interference를 선호하고 이상적으로는 antipodal하다.

 

Perturbing a Single Feature

가장 간단한 non-uniform superposition은 하나의 feature를 변경하고 나머지는 uniform하게 두는 것이다.

  • n = 5 features, m = 2 차원인 실험을 진행한다.
  • 중요도는 1, activation density 1 - S = 0.05로 정오각형 모양을 가진다.

하나의 점의 sparsity를

  • Dense하게 하면 다른 feature들이 point로부터 밀려난다.
  • Sparse하게 하면 다른 점들이 그 point로 당겨진다.
    • 너무 sparse하게 만들면 sparse한 점을 중심으로 하는 digons로 변경된다.
      • 아래 그래프에서 각 geomtery에 대응하는 loss function이 교차하는 지점을 뜻한다.

위 실험의 결과 non-uniform superposition은 uniform superposition 구조의 deformation이나 조합으로 이루어지고 전혀 다른 구조가 아님을 보여준다.

Pentogonal solution이 원 위에 있지 않는 이유는 positive interference를 줄이기 위해서이다.

  • Negative bias를 이용해 noise를 줄이고
  • Weight를 $||W_i|| = 1/(1-b_i)$를 사용해 이를 보상한다.
  • 즉, 원으로부터의 거리는 positive infernece의 크기라고 이해할 수 있다.

또한 sub-optimal한 solution들을 시각화해 loss curve와 다른 geometries가 어떻게 연관되어있는지 영감을 얻을 수 있다.

 

 

Correlated and Anticorrelated Features

모델이 correlated 된 feature들을 나타낼 때 선호하는 순서가 있는 것 처럼 동작한다.

  • 이상적으로는 correlated feature들은 orthogonal하게 표현된다.
  • 그럴 수 없을 때 두 feature는 가능한 가장 가까이 배치된다.
    • negative interference보다 positive interference를 선호한다.
  • 모든 feature를 나타낼 공간이 없을 때는 두 feature들이 collapse하고 그 주 성분으로 나타내진다.

반대로 anti-correlated되어 있는 feature들의 경우 모델은 negative interference를 선호한다.

 

Setup for Exploring Correlated and Anticorrelated Features

Correlated Feature Sets

  • Correalted feature set에 있는 feature들의 entry가 동시에 0이 되거나 0이 아니게 설정

Anticorrelated Feature Sets

  • 하나의 feature set이 0이면 다른 feature set은 [0, 1] 값을, vice versa

 

Organization of Correlated and Anticorrelated Features

먼저 m = 2인 겨우로 실험을 해 weight들을 2D 공간의 점으로 시각화한다.

 

Local Almost-Orthogonal Bases

Correlated된 feature들을 orthogonal하게 배치하려는 특성이 생각보다 강하다는 것을 발견했다.

  • 큰 모델에서도 모델 전체가 superposition을 가지고 있어도 correlated feature들은 거의 orthgonal하게 배치되고 적은 interference를 가진다.

위 결과가 실제 neural network에도 적용된다면 "local non-superposition" 가정을 만들 수도 있다.

  • 분포의 일부분은 superposition이 아니라고 가정할 수 있다.
  • 이는 superposition이 존재할 때 잘 사용되지 않는 PCA 기법을 사용하 근거를 제공할 수 있다.

 

Collapsing of Correlated Features

흥미로운 점은 모델이 Principal Components Analysis (PCA)와 superposition에 trade off를 가지는 것 같다는 사실이다.

  • 두 개의 correlated feature들 a와 b가 있는데 모델이 하나의 feature 만 표현할 수 있다면 모델은 principal component인 $\frac{(a + b)}{\sqrt{2}}$ 만 나타내고 $\frac{(a - b)}{\sqrt{2}}$는 무시한다.

 

Feature들이 더 correlated 될수록 PCA가 더 좋은 방법이 된다.

반대로 feature들이 더 sparse 해질수록 superposition이 더 좋은 방법이 된다.

 

 

 

Superposition and Learning Dynamics

모델의 learning dynamics을 공부하는 것이 흥미로운 여러 이유가 있다.

본 논문에서는 학습 과정에서 발견한 놀라운 현상들을 정리하고 자세한 연구는 미래 연구 주제로 제시한다.

 

Phenomenon 1: Discrete "Energy Level" Jumps

많은 수의 feature들을 사용한 learning dynamic에서 "energy level jump"가 지배적이다.

  • feature들이 다른 feature dimensionality들 사이를 jump한다.

이전 섹션에서 uniform superposition의 geometry를 조사할 때 사용한 모델 중 모든 feature가 digons로 수렴한 모델을 선택해 확인해보자.

  • 각 선은 하나의 feature의 dimensionality를 뜻한다.

흥미로운 점은 몇몇 feature들의 dimensionality는 다른 값들 사이를 "jump"하며 값을 바꾼다는 것이다.

  • 이때 loss 값은 큰 감소를 보인다.

이는 더 큰 모델들에서 부드러운 loss 함수가 사실은 아주 많은 작은 step들로 이루어져 있을 수 있다는 것을 시전한다.

 

Phenomenon 2: Learning as Geometric Transformatinos

특정 기하학 구조를 가지게 되는 learning dynamic은 사실 간단하고 독립적인 기하학 transformation들로 이루어져있다.

  • n = 6, m = 3인 모델에서 3개의 feature 씩 하나의 correlation set으로 구성한다.
  • Loss curve로 알 수 있듯이 학습은 별개의 영역에 대응하는 기하학 transformation로 진행된다.

 

 

 

Superposition in a Privileged Basis

이제까지는 privileged basis가 없는 모델들을 살펴봤다. (word embedding, transformer residual stream 등)

이 섹션에서는 privileged basis인 neuron이 존재하는 경우를 탐구한다. (transformer MLP layer, conv net neuron 등)

 

Privileged basis가 존재하는 toy 모델을 만드는 방법은 두 가지가 있다.

  • hidden layer에 activation function을 추가하거나
  • hidden layer에 L1 regularization을 적용할 수 있다.

본 논문에서는 activation을 추가하는 방법을 탐구한다.

 

아래 식을 "ReLU hidden layer" 모델로 명명한다.

$h = ReLU(Wx)$

$x' = ReLU(W^Th + b)$

 

ReLU를 추가하면 모델을 interpretability 관점에서 완전히 변형시킨다.

  • 아래 그림은 linear hidden layer 모델과 ReLU hidden layer model을 비교한다.

  • 기존에는 input을 features로, middle layer를 neuron으로 생각했다. 따라서 W는 feature들을 neuron들로 mapping한다.
  • 위 그림에서 볼 수 있는 것은 feature들이 neuron들과 구조적인 형태로 정렬된다는 것이다.
    • 즉 거의 모든 neuron들이 하나의 feature를 전담하여 나타낸다.

 

Visualizing Superposition in Terms of Neurons

모델을 시각화하는 또 다른 방법인 per-neuron stacked bar plot을 제시한다.

  • 각 column은 W의 하나의 column을 시각화한다.
  • 각 직사각형은 하나의 weight entry를 나타내고 높이는 그 절대값에 대응한다.
  • 색깔은 나타내는 feature에 대응한다.
  • 음수는 x-axis 아래로 나타내진다.
  • Rectangle의 순서는 중요하지 않다.

 

이제 ReLU hidden layer toy 모델을 시각화한다.

  • n = 10, m = 5, $I^i = 0.75^i$, 그리고 sparsity를 조절한다.
  • Neuron의 색은 monosemantic 또는 polysemantic에 따라 정해진다.

 

중요한 점은 sparsity가 증가할 수록 monosemantic에서 polysemantic으로 바뀐다는 점이다.

  • 하나의 모델에서 monosemantic과 polysemantic neuron이 공존할 수 있다.
  • 또한 neuron-level phase change가 존재하는 것처럼 보인다.

 

Limitations of the ReLU Hidden Layer Toy Model Simulating Identity

이 섹션에서 제시된 toy model은 모델의 적용 범위를 제한시키는 약점이 있다.

  • ReLU hidden layer로부터 모델이 혜택을 받지 않는다.
  • 만약 기회가 있다면 모델은 layer를 사용하지 않으려 한다.
    • 예를 들어 biasa를 모두 양수 값으로 설정해 선형적으로 작동하게 한다.

 

 

 

Computation in Superposition

우리는 모델이 superposition 상태에서의 계산을 할 수 있다 생각한다.

이를 탐구하기 위해서 input과 output은 가상의 disentangled 모델, hidden layer는 관측되는 모델로 설정한다.

더 자세히 우리는 모델이 $y=abs(x)$를 계산하도록 한다.

  • 이는 ReLU를 통해 간단하게 계산할 수 있기 때문이다
  • $abs(x) = ReLU(x) + ReLU(-x)$

 

 

Experiment Setup

위 설정과는 다르게 input feature vector x의 $x_i$가 [-1, 1]에서 샘플되도록 한다. Target output은 $y = abs(x)$로 정한다.

 

이전 섹션과 동일하게 ReLU hidden layer를 사용하지만 더 이상 weight 2개를 동일하게 사용하지 않는다.

  • $h = ReLU(W_1x)$
  • $y' = ReLU(W_2h+b)$
  • Loss는 이전과 동일하게 중요도 $I_i$를 가중치로 사용한 mean squared error이다.

 

 

Basic Results

두 개의 Weight $W_1$과 $W_2$가 독립적으로 학습되기 때문에 단순히 $W_2^TW_1$을 연구할 수 없다.

이전 section에서 확인했듯이 hidden layer activation function을 추가하면 weight를 neuron으로 직접적으로 시각화 할 수 있다.

 

먼저 n = 3 features가 m = 6 hidden layer 뉴론들로 절대값을 계산하도록 한다.

그 결과 예상했듯이 각 input feature $x_i$에 대해 양수 $ReLU(x_i)$ neuron과 음수 $ReLU(-x_i)$ neuron을 가지게 하고 그 합으로 절대값을 구한다.

 

 

Superposition vs Sparsity

그럼 모델이 superposition을 사용해 더 많은 feature의 절대값도 구할 수 있을까?

  • n = 100, m = 40, $I_i = 0.8^i$과 sparsity를 조절하며 실험한다.
    • 시각화에서 W의 절대값을 사용한다.
    • Neuron 색은 얼마나 polysemantic한지를 알려준다.

위 그래프는 다음과 같은 사실들을 보여준다.

  • Activation function이 privileged basis를 생성하고 feature들이 basis차원에 정렬되도록 한다.
  • Superposition으로 나타내진 데이터에도 계산을 할 수 있다는 것을 보여준다.
  • 많은 neuron들이 pure feature를 encode하지만 일부는 매우 polysemantic하다.
  • 또한 몇몇 neuron들은 큰 weight로 encode된 하나의 "primary" feature과 작은 weight 크기로 encode된 여러 개의 "secondary" feature을 동시에 가진다는 것을 보여준다.

 

 

The Asymmetric Superposition Motif

이런 superposition에서의 계산은 어떻게 가능한가?

이 섹션에서는 weight를 이용해 이를 설명한다. 또 asymmetric superposition을 설명한다.

 

이 두 neuron들은 asymmetric superposition과 inhibition을 구현한다.

 

일반적인 superposition에서는 같은 weight로 feature들을 저장한다.

하지만 asymmetric superposition에서는 feature들을 다른 크기로 저장한다.

  • input: W = [2, -1/2], output: W = [1/2, 2]
  • 이는 하나의 feature가 다른 feature와 크게 interfere하게 하지만 다른 feature는 interfere를 경험하지 않게 해준다.

이런 interference를 피하기 위해서 모델은 positive interference가 있는 경우를 다른 neuron으로 억제한다.

  • Positive interference를 negativei interference로 바꿔 ReLU에 의해 없어지게 한다.

 

 

 

The Strategic Picture of Superposition

Safety, Interpretability, & "Solving Superposition"

"Solving superposition"은 다른 많은 interpretability properties와 연과되어 있다.

  • Decomposing Activation Space
    • Activation space를 독립적으로 이해 가능한 요소들로 나눈다.
  • Describing Activations in Terms of Pure Features
    • 하나의 activation을 여러 요소들로 나누어 설명할 수 있다.
  • Understanding Weights
    • Weight를 이해하기 위해서는 이해할 수 있는 feature들과 연결되어 있어야 한다.

 

 

Three Ways Out

세 가지 superposition을 해결할 방법들을 제시한다.

  • Superposition이 없는 모델 만들기
  • Overcomplete basis 찾기
  • Hybrid approaches

 

Approach 1: Creating Models without Superposition

가장 간단한 방법은 L1 regularization term을 hidden layer activation에 추가하는 것이다.

  • 이는 중요하지 않은 neuron들을 나타내지 않음으로 이해 가능하고 basis aligned된 feature들을 생성한다.
  • 하지만 loss가 증가하게 된다.

또 하나의 방법은 더 큰 모델을 사용하는 것이다. 하지만 이는 높은 계산 비용을 야기한다.

  • 이때 우리는 Mixture of Experts (MoE) 모델에서 영감을 받을 수 있다.

 

Approach 2: Finding an Overcomplete Basis

또 다른 방법은 일단 superposition이 있는 모델을 학습하고 각 feature에 대응하는 overcomplete basis를 찾는 것이다.

  • 이는 기존 sparse coding (dictionary learning) 문제와 연관되어 있다.

이 방법의 challenge들은 다음과 같다.

  • Feature 수를 알 수 없다.
  • Solution은 실제 surface computational structure가 아닌 virtual structure를 가지게 된다.
  • 이미 큰 모델들을 sparse coding으로 접근하면 문제의 크기가 너무 커진다.
  • Superposition에 의한 interference는 모델이 학습 과정에서 더 적은 superposition을 가지게 하고 이를 decode하기 어렵다.

 

Approach 3: Hybrid Approaches

Hyprid approaches

  • Superposition이 적은 모델을 만들면 더 쉽게 decode 할 수 있다.
  • 반대로 구조를 조정해 overcomplete basis를 더 쉽게 찾거나 계산량을 줄일 수 있다.