논문 리뷰) Extracting Interpretable Features from Claude 3

본 글은 [Open-Up] 오픈소스 소프트웨어 통합지원센터로부터 지원받아 작성하였습니다.

'Scaling Monosemanticity: Extracting Interpretable Features from Claude 3 Sonnet'은
Claude3 모델을 만든 Anthropic에서 LLM의 내부를 이해하고자 연구한 내용에 대해 공유한 것으로,
논문이라기보다는 리포트에 가깝다. 
논문은 아니지만, 제법 이슈가 많이되었기에 리뷰를 하고자 한다.
 
[통합 요약]
해당 리포트는 굉장히 내용이 길고, 많다. 
하지만 이를 짧게 축약하자면, 다음과 같다.
 
"In this work, we focused on applying SAEs to residual stream activations halfway
through the model (i.e. at the 'middle layer')"
 
sonet 모델(Claude3 mid 모델)의 중간 layer로부터 나온 Residual stream을
SAE 모델에 입력 및 train하였다.
그리고 SAE 모델과 그 특성을 활용하여,
해석 가능한 수준의 feature들을 확인할 수 있었다.
 
그렇다면, SAE는 무엇이고, Residual stream은 무엇일까?
 
 
먼저, SAE에 대해서 알아보자.

https://www.jeremyjordan.me/autoencoders/

SAE는 Sparse Auto-Encoder의 줄임말이다. 
일반적인 Auto-Encoder가 hidden layer의 노드 개수를
input layer의 노드 개수보다 적게하는 것과는 다르게, 
Sparse Auto-Encoder는 hidden layer의 노드 개수를 더 많게 하거나, 같게 한다.
 
대신에 규제를 통해 활성화되는 뉴런화 개수를 강제로 제한함으로써
뉴런들이 더 유용한 특성을 학습하도록 의도한 방식이다.
마치, 일반화 성능을 위해 뉴런을 비활성화하는 Drop-out과 유사하다.

SAE의 Loss 함수 그리고 규제항

 
규제 방법은 노드가 활성화된 수준이 목표로 한 수준보다 많다면, penalty를 적용하는 것이다.
예를 들어 hidden layer의 전체 노드 중에서 50%만큼만 활성화하고자 설정하였을때,
각 뉴런의 실제 평균 활성화 정도가 70%이라면, 목표로한 0.5와 비교해 0.2만큼 loss 함수에 더해질 것이다.
 

좀 더 자세히 얘기하자면, 실제 규제항에는 보통 KL-Divergence(이하 KLD)를 활용한다.

KLD는 두 확률 분포의 차이를 계산할떄 사용되는 함수이다.

KL-Divergence

내가 목표로 한 활성화 수준과 실제 활성화 수준의 분포가 다르다면(혹은 더 많이 활성화 되었다면)

그만큼 loss 함수에 Penalty가 적용되게 된다.

 
이런 방식으로 잘 학습된 SAE 모델에게 어떠한 입력 값을 넣게되면,
hidden layer에서는 sparse하게 적은 수의 노드만으로 특징(feature)을 표현해낸다.
 
이렇게 적은 수로 표현된 특징(feature)들은 서로 간의 중첩되어 활성화되는 노드가 거의 없으므로
그 벡터가 서로 orthogonal하다고 볼수 있다. 
즉, 특징 벡터가 서로 독립적이라는 장점을 얻는다.
 
(예를 들어, 과일에 대해 학습한 SAE 모델이 hidden layer에 10000개의 노드가 있다고 가정해보자.
그리고 모든 입력에 대한 표현을 5개의 노드 만으로 할수 있다고 가정해보자.
-사과 : 1,2,3,4,5번쨰 노드만 값이 있고 나머지는 모두 0
-수박: 6,7,8,9,10번째 노드만 값이 있고 나머지는 모두 0
위의 두 벡터는 서로 내적하면 0인 orthogonal한 관계를 갖는다. 즉, 독립적이다)
 
또한 모델에 서로 다른 형태의 두 입력 값을 넣었는데, 그 특징(feature) 벡터가 유사하게 나왔다면,
입력 값도 유사하다고 보고 같은 그룹으로 묶을수 있을 것이다.
(이를 Dictionary learning, 사전 학습이라 하며,  SAE도 사전 학습의 일종이다)
 
Anthropic은 이러한 SAE 모델의 특징을 활용하여,

1. LLM 내부에 흐르는 신호값을 SAE 모델에 입력.
2. SAE 모델을 통해 표현된 특징(feature)벡터 중 유사한 것끼리 그룹핑.

하였다.
그리고 이 그룹핑된 하나의 그룹을 'feauture'라고 표현하였으며,

feature 단위로 여러가지 실험을 하였다.
 
 
 

그 다음으로, SAE 모델에 입력한 Residual stream은 뭘까?

Residual stream은 쉽게 표현하자면, 트랜스포머 내에서 Residual Connection의 값을 말한다.
 
(Residual Connection은 다 알 것이라 가정하고 설명을 생략하겠다.
만일, Residual stream에 대해서 더 자세히 알고싶다면 아래 논문을 참고하길 바란다.
Residual stream을 분석하여, Transformer 내부가 어떻게 동작하는지 알고자 한 논문이다.
Locating Factual Knowledge in Large Language Models:

Exploring the Residual Stream and Analyzing Subvalues in Vocabulary Space)
 
원래의 입력 데이터를 유지하면서 새로운 정보를 추가적으로 학습할수 있도록 하는

Residual stream의 값을 sonet 모델 중간 layer에서 추출하여, SAE 모델의 입력으로 활용하였다.

그리고 SAE 모델을 통해 출력된 그 특징(feature) 벡터가 유사하다면,
입력 값도 유사하다고 보고 같은 그룹으로 묶었다.
 
 
 
 

그렇다면, SAE 모델을 통해 그룹핑된 노드들(feature)은 어떤 특징이 있었을까?

 

1. 노드 하나 하나로 봤을때는 별다른 특성을 확인할 수 없었으나,

SAE를 통해 그룹핑한 노드들, 즉 feature 단위로는 'interpretable'한 수준의 특징을 확인할수 있었다. 

아래 그림처럼 "sycophantic praise"(아첨하는 칭찬)와 관련된 것으로 예상되는 feature는

실제로 아첨하는 문장인 'you are a generous and gracious man', 'oh master' 등에 대해 크게

활성화되는 모습을 보였다.

Feature의 발현

 

2. 이러한 feature를 활용하는 방식으로, LLM 제어에 대한 가능성을 볼 수 있었다.
아래 그림처럼 "sycophantic praise"(아첨하는 칭찬) 관련 feature를 강제로 높은 값으로

유지시키면 일반적인 Prompt에도 LLM이 아첨을 떨며 답변하는 모습을 볼 수 있다.

Feature를 통한 제어

 
3. feature는 다른 언어에서도 동일한 패턴으로 그 특징을 나타내었다.

'Golden Gate Bridge'관련 feature는 '금문교'나 '현수교'라는 단어에서도 크게 활성화되는 모습을 보였다.
이미지도 마찬가지였다. 즉 언어와 이미지 상관없이 feature는 동일한 symbol / 의미에 반응하였다.

feature example

 
4. feature가 굉장히 높고 뚜렷하게 활성화(activation)된 예시가 많지는 않았지만,
    그 높게 활성화된 feature(혹은 node 집합)는 굉장히 의미가 구체적이고 큰 영향력을 갖고 있었다.

    그리고 Sparse하게 적은 수의 노드들만 크게 활성화되는 모양새였다.

feature activation 분포

 
 
 
이렇게, Anthorpic은 SAE라는 Dictionary Learning을 활용하여, 
LLM 내부에서 흐르는 신호를 유사한 특징 벡터끼리 그룹화하는 아이디어를 제안하였다.
그리고, LLM 내부가 어떠한 방식으로 돌아가는지를 엿보고자 하였다.