Sparse Autoencoders Find Highly Interpretable Features in Language Models - 논문 리뷰

https://arxiv.org/abs/2309.08600

Sparse Autoencoders Find Highly Interpretable Features in Language Models

 

기존 Transformer 구조에서 뉴런 수 보다 feature 수가 더 많아 feature가 한 뉴런에서 여러 개가 발현되는 superposition 현상이 발생한다. 

그럼 뉴런 수를 늘려 한 feature에 한 개의 뉴런만 발현하게 되면 LLM을 해석할 수 있지 않을까 -> 여기서 나온 방향이다. 

 

 

 

 

 

 

1. 문제 정의 (Problem Definition)
   • 핵심 문제: AI 시스템의 내부 동작을 해석하는 데 큰 어려움을 주는 요소 중 하나는 신경망에서 여러 문맥에서 활성화되는 다의성(Polysemanticity)입니다. 이 다의성으로 인해 신경망의 내부 작동을 인간이 이해할 수 있는 방식으로 명확하게 설명하기 어려워집니다. 논문에서는 이러한 다의성의 주요 원인 중 하나로 '중첩(Superposition)' 현상을 지목합니다. 이는 신경망이 뉴런의 개수보다 많은 피처(특징)를 학습하고, 이를 독립적인 뉴런이 아닌 비정교적인 방식으로 여러 방향으로 표현하는 문제입니다.
   • 연구 목표: 중첩 현상을 해결하고 신경망 내부에서 학습된 피처들이 더욱 단일 의미를 가지도록 하는 해석 가능성을 높이는 방법을 찾는 것. 이를 통해 신경망의 행동을 더 잘 해석할 수 있는 길을 마련하고자 함.
2. 해결하려는 문제 (Objective)
   • 중첩 해결: 신경망 내부의 중첩된 피처를 분해하여 더 해석 가능한 형태로 만들기 위해 sparse autoencoder(희소 자동 인코더)를 활용해 언어 모델의 내부 활성화(activation)를 재구성하려고 시도합니다.
   • 구체적 목표: 희소성을 띠는 피처를 학습하여 그 피처들이 더 해석 가능하고 단일 의미(모노세맨틱)를 가질 수 있도록 하고, 이를 통해 더 세밀하게 신경망의 특정 행동을 유발하는 피처를 식별하는 것.
3. 사용된 방법 (Methods)
   • Sparse Autoencoder 사용: 언어 모델의 내부 활성화를 재구성하기 위해 희소 자동 인코더를 사용합니다. 이는 sparse dictionary learning(희소 사전 학습) 기법을 사용하여 언어 모델의 내부 활성화 벡터를 희소 선형 결합으로 나타낼 수 있는 방향을 학습하는 방법입니다. 여기서 희소성을 적용한 피처들은 기존의 비정교적인 방식보다 더 해석 가능성이 높습니다.
   • 실험 대상: Pythia-70M과 같은 대형 언어 모델의 내부 활성화 데이터를 사용하여 실험을 진행하였으며, residual stream, MLP(다층 퍼셉트론) 하위 계층 및 attention head 계층에서 내부 활성화를 샘플링하여 sparse autoencoder를 훈련시킴.
4. 결과 (Results)
   • 피처 해석 가능성 증가: 논문은 sparse autoencoder가 학습한 피처들이 기존의 뉴런이나 다른 행렬 분해 기법으로 학습한 피처들보다 더 해석 가능성이 높다는 것을 자동 해석 가능성 점수(autointerpretability score)를 통해 증명했습니다.
   • 피처 분해 정확도 향상: indirect object identification task(간접 목적어 식별 작업)에서 sparse autoencoder로 학습된 피처들이 다른 분해 방법보다 더 세밀하게 신경망의 특정 행동을 유발하는 피처를 식별할 수 있음을 확인했습니다.
5. 한계점 (Limitations)
   • 완전한 복원 불가능: 희소 자동 인코더를 사용하여 학습된 피처들로 신경망의 모든 정보를 복원하지는 못했습니다. 재구성 손실이 남아 있으며, 이로 인해 복원된 활성화가 원래의 활성화를 완벽하게 대체하지 못했습니다.
   • MLP 계층에서의 문제: residual stream 계층에서는 희소 자동 인코더가 잘 작동했으나, MLP 계층에서는 비활성화된(dead) 피처들이 다수 존재해 모든 피처를 효과적으로 학습하지 못했습니다.
   • 추가 연구 필요성: 중첩 현상을 완전히 해결하고 모든 계층에서 신경망의 피처를 해석할 수 있는 방법을 찾기 위해서는 더 많은 실험과 개선이 필요합니다.

---

위 논문에서는 sparse autoencoder를 사용하여 언어 모델의 내부 피처를 분해하고, 그 피처들이 더욱 해석 가능하고 단일 의미를 가지도록 하는 방법을 제시했습니다. 이 연구는 AI 시스템의 해석 가능성을 높이는 데 기여할 수 있으며, 신경망의 특정 행동을 더 정확히 이해하는 데 유용할 수 있습니다.

 

그럼 여기선 feature 수보다 뉴런 수가 너무 적어서 중첩이 일어나니 뉴런 수를 많이 늘려서 중첩을 줄이고, 단일적으로 뉴런이 반응하도록 만드는거네 
그럼 여기서 autointerpretability score은 또 뭐야? 

 

맞습니다! 이 논문에서는 뉴런 수보다 더 많은 피처가 학습되기 때문에 "중첩"이 발생한다고 설명하고, 이를 해결하기 위해 뉴런의 수를 많이 늘려서 각 뉴런이 단일 피처에만 반응하도록 만들려는 목표를 가지고 있습니다. 이는 희소 자동 인코더(Sparse Autoencoder)를 통해 이루어지며, 이를 통해 신경망 내부의 복잡한 피처들이 더 단일 의미를 가지게 하고, 해석 가능성을 높이려는 것입니다.

Autointerpretability Score(자동 해석 가능성 점수)

Autointerpretability Score는 학습된 피처의 해석 가능성을 평가하는 자동화된 방식입니다. 구체적으로는, 피처가 얼마나 잘 해석될 수 있는지를 수치화한 것입니다. 이 점수는 다음과 같은 과정으로 계산됩니다:

1. 피처가 활성화되는 텍스트 샘플 수집: 우선 피처가 언어 모델에서 활성화되는 텍스트 샘플들을 수집합니다. 이때 피처가 강하게 활성화된 텍스트 조각(top-activation)과 무작위로 선택된 텍스트 조각(random-activation)을 함께 사용합니다.
2. GPT 모델을 통한 설명 생성: 활성화된 텍스트를 바탕으로 GPT-4와 같은 모델을 사용해, 그 피처가 어떤 의미를 가지는지에 대한 설명을 자동으로 생성합니다. 예를 들어, 특정 피처가 활성화된 텍스트를 보고 "이 피처는 주로 사람 이름에 반응한다"와 같은 설명을 만들 수 있습니다.
3. 설명과 예측 일치도 계산: GPT 모델이 생성한 설명을 사용해 다른 텍스트 샘플에서 해당 피처의 활성화를 예측합니다. 이때 실제 피처 활성화와 예측된 활성화 사이의 상관관계를 계산해, 그 피처의 해석 가능성을 점수화한 것이 Autointerpretability Score입니다.

이 점수가 의미하는 것

이 점수는 피처가 얼마나 명확하게 해석될 수 있는지를 평가하는 지표로, 점수가 높을수록 피처가 더 단일적인 의미(모노세맨틱)를 갖고 있다는 것을 의미합니다. 반대로 점수가 낮다면 해당 피처는 여러 가지 서로 관련 없는 문맥에서 활성화되는 경향이 있으며, 이는 다의성을 가진 피처임을 나타냅니다.

결론적으로, Autointerpretability Score는 신경망에서 학습된 피처들이 얼마나 명확하게 해석될 수 있는지를 평가하는 도구로, sparse autoencoder로 학습한 피처들이 더 높은 해석 가능성을 가질 수 있음을 입증하기 위해 사용됩니다.

 

MLP 계층에서는 왜 문제가 발생했을까?
MLP Layer에서는 뉴런수가 residual stream layer 보다 많잖아

MLP 계층에서 문제가 발생한 이유는 여러 요인들이 결합되어 발생한 것으로 보입니다. 뉴런의 수가 많음에도 불구하고, MLP 계층에서 기대한 만큼의 해석 가능성을 얻지 못한 이유는 단순히 뉴런 수와는 다릅니다. 이를 몇 가지 요소로 나눠 설명할 수 있습니다:

1. 희소성(sparsity)의 적용 문제

• 희소성을 적용하는 방식에서, MLP 계층의 특성상 일부 피처들이 아예 활성화되지 않는 경우가 많았습니다. 논문에서 언급한 것처럼, MLP 계층에서는 "dead feature"라고 부르는 비활성 피처들이 다수 발생하여, 실제로 해석할 수 있는 피처들이 적어지는 문제가 있었습니다.
• 이는 특히 비선형 활성화 함수(ReLU와 같은)가 도입되면서 발생하는 문제일 수 있습니다. MLP 계층에서는 뉴런들이 비선형 활성화를 거치며 정보가 보다 제한적으로 처리되기 때문에, 희소성을 부여하는 과정에서 많은 피처가 전혀 활성화되지 않을 가능성이 높습니다. 결국 이로 인해 해석 가능한 피처의 수가 줄어드는 것입니다.

2. 중첩과 비활성 피처

• MLP 계층에서 뉴런의 수가 많다 하더라도, 중첩(Superposition) 문제는 여전히 발생할 수 있습니다. 즉, 피처가 뉴런 수보다 더 많고, 서로 비정교적인 방식으로 결합되면 여전히 여러 피처가 동일한 뉴런에서 활성화되는 현상이 일어날 수 있습니다.
• 하지만 이 문제는 residual stream 계층보다 MLP 계층에서 더 두드러집니다. MLP 계층에서는 특히 뉴런들이 특정한 정보에 집중하기보다 다양한 정보를 통합적으로 처리하려는 경향이 강하기 때문에, 이러한 중첩이 더 복잡하게 나타나게 됩니다. 이것이 결과적으로 다의성(Polysemanticity)을 증가시키고, 희소성 기반 학습이 어려워지게 만듭니다.

3. 피처 활성화의 비일관성

• MLP 계층에서는 피처들이 고르게 활성화되지 않는 경우가 많습니다. residual stream에서는 피처들이 비교적 일관되게 활성화되지만, MLP 계층에서는 피처가 일부 샘플에서만 활성화되고 나머지 샘플에서는 거의 활성화되지 않는 경우가 있습니다.
• 이는 학습된 피처들이 충분히 다양한 문맥에서 일관되게 활성화되지 않기 때문에 해석 가능성을 높이기 어려운 이유가 될 수 있습니다. 예를 들어, 특정 피처가 특정 조건 하에서만 활성화된다면, 그 피처를 해석하는 것이 훨씬 더 복잡해집니다.

4. Residual Stream과 MLP의 차이

• Residual Stream은 직선적인 정보 흐름을 다루기 위한 계층으로, 신경망의 각 층에서 발생하는 정보 손실을 보충하고, 연속적인 정보를 저장하는 역할을 합니다. 반면, MLP는 더 복잡한 비선형 변환을 수행하는 역할을 맡습니다.
• 이 차이 때문에 Residual Stream에서는 비교적 단순한 피처 분해가 가능하지만, MLP에서는 더 복잡한 정보 처리가 일어나기 때문에 동일한 방법으로 희소성을 적용하는 것이 어려울 수 있습니다.

결론적으로:

MLP 계층에서는 뉴런 수가 많음에도 불구하고 희소성을 적용하는 과정에서 비활성 피처(dead feature)가 다수 발생하고, 정보가 더 복잡한 방식으로 결합되기 때문에 해석 가능성에 한계가 발생합니다. 따라서 MLP 계층에서 더 나은 성능을 얻기 위해서는 피처 학습 방법의 수정이나 활성화 방식의 조정이 필요할 수 있습니다.

 

그럼 여기서 Sparse Autoencoder를 어떤 방식으로 학습했고, 발전시켰어?

논문에서는 Sparse Autoencoder(희소 자동 인코더)를 통해 언어 모델의 내부 활성화 데이터를 해석 가능한 방식으로 분해하는 방법을 제안했습니다. 이 과정을 자세히 설명하자면, sparse dictionary learning(희소 사전 학습) 방식을 사용하여 언어 모델의 내부 활성화를 재구성하고, 각 피처를 더 해석 가능하게 만들려는 목적이 있습니다. 논문에서 Sparse Autoencoder를 학습하고 발전시킨 과정은 다음과 같이 정리할 수 있습니다.

1. Sparse Autoencoder의 구조

• Autoencoder 구조: 희소 자동 인코더는 하나의 숨겨진 레이어를 가진 신경망입니다. 이 신경망은 입력 데이터를 잠재 공간으로 압축한 뒤 다시 원래 데이터로 복원하는 과정을 통해 피처를 학습합니다. 이때 잠재 공간에서의 피처들이 희소성을 띠도록 제약을 추가합니다.
• 희소성 부여: 희소 자동 인코더는 출력이 복원 오류(Reconstruction Loss)를 최소화하면서도, 잠재 공간에서의 활성화가 희소하도록 학습됩니다. 이를 위해 L1 손실을 추가하여 많은 뉴런이 활성화되지 않도록 강제하는 희소성 제약을 적용합니다.
• 수식:
  L(x)=∣∣x−x^∣∣22+α∣∣c∣∣1
  ​여기서 x^는 복원된 입력 벡터, c는 잠재 공간에서의 활성화 벡터이며, ∣∣x−x^∣∣22는 복원 손실, ∣∣c∣∣1​은 희소성 손실을 의미합니다. α는 희소성 제약의 강도를 조절하는 하이퍼파라미터입니다.

2. 학습 과정

• 데이터 수집: 먼저 언어 모델(Pythia-70M 등)의 내부 활성화 데이터를 수집합니다. 이 활성화 데이터는 residual stream, MLP 계층 등에서 샘플링됩니다.
• Autoencoder 훈련: 수집된 활성화 데이터를 입력으로 사용해 autoencoder를 훈련시킵니다. 훈련 과정에서는 피처들의 희소성을 유지하면서 입력 데이터를 최대한 정확하게 복원하는 것을 목표로 합니다. 이때 희소성을 부여하는 방식으로 L1 정규화를 사용하며, 이는 활성화된 뉴런의 수를 제한하여 더 해석 가능한 피처들이 학습되도록 유도합니다.
• Tied Weights 사용: 희소 자동 인코더의 가중치 행렬에서 encoder와 decoder가 동일한 가중치 행렬을 사용하도록 제한합니다. 이를 통해 메모리 사용량을 절약하고, 해석 가능한 방향으로 가중치가 학습되도록 유도합니다.

3. 발전된 기법

• Autointerpretability 기반 평가: 학습된 피처들이 얼마나 해석 가능한지 평가하기 위해 Autointerpretability Score를 도입했습니다. 이는 학습된 피처가 특정 텍스트에서 활성화되는 패턴을 분석하고, GPT-4와 같은 대형 언어 모델이 그 피처의 의미를 자동으로 해석할 수 있는지 확인하는 방식입니다. 이를 통해 피처의 해석 가능성을 정량적으로 평가할 수 있었습니다.
• MLP 계층에서의 적용: 논문에서는 주로 residual stream 계층에서 Sparse Autoencoder를 성공적으로 적용했지만, MLP 계층에서도 이를 적용하려고 시도했습니다. MLP 계층에서는 더 복잡한 정보 처리가 일어나지만, 여전히 일부 피처들은 해석 가능한 방식으로 학습될 수 있음을 확인했습니다. 다만, MLP 계층에서는 더 많은 "dead feature"(비활성 피처)가 발생하여, 일부 피처는 제대로 학습되지 않는 한계가 존재했습니다.

4. Sparse Autoencoder의 주요 발전점

• 희소성 제어: 하이퍼파라미터 α\alphaα를 조절하여 희소성의 정도를 조정할 수 있었습니다. α\alphaα 값이 커질수록 더 많은 뉴런들이 활성화되지 않도록 억제되었지만, 반대로 복원 정확도는 낮아지는 경향이 있었습니다. 따라서 복원 손실과 희소성 사이의 균형을 맞추는 것이 중요한 과제가 되었습니다.
• Residual Stream에서의 성공적인 적용: Residual stream 계층에서는 sparse autoencoder가 매우 잘 작동하여, 기존의 PCA(주성분 분석), ICA(독립 성분 분석) 등 다른 방법론보다 더 해석 가능한 피처를 도출할 수 있었습니다. 이는 특히 언어 모델의 초기 레이어에서 더 뚜렷한 성과를 보였습니다.
• Causal Feature Identification: sparse autoencoder로 학습된 피처를 통해 causal mediation analysis(원인 중재 분석)를 실행하여, 특정 모델 행동을 유발하는 중요한 피처들을 더 정확하게 식별할 수 있었습니다.

5. 미래 발전 가능성

논문에서는 MLP 계층에서 발생한 문제점을 해결하기 위해 추가적인 연구가 필요하다고 제안합니다. MLP 계층에서 dead feature가 많이 발생하는 문제를 해결하기 위해 새로운 초기화 방법을 적용하거나, 기존 가중치와의 관계를 고려한 학습 방법을 도입하는 방안이 제시되었습니다.

요약:

Sparse Autoencoder는 언어 모델의 내부 피처를 해석 가능하게 만들기 위한 중요한 도구로 사용되었으며, residual stream 계층에서는 매우 성공적인 결과를 보여주었습니다. 다만, MLP 계층에서는 dead feature 문제로 인해 성능이 제한되었고, 이를 해결하기 위한 추가 연구가 필요한 상태입니다.