Review

TL; DR

저자소속: UCLA, University of Alberta, UIUC, Harvard

Summary

Background

• 요즘 Post-training의 목적
  • Truthfulness 향상
  • Safety alignment (악의적 공격에 대한 방어)
  • 모델의 confidence 보정

• 요즘 Post-training 기법들
  • Direct Preference Optimization (DPO)
  • Reinforcement Learning from Human Feedback (RLHF)
  • Downstream task improvement

Motivation

• Post-training을 더 잘하기 위해서, 기계적으로 해석하자!
  • 기존 연구들은 제한된 알고리즘, 태스크, 모델, 방법론(SAE)을 쓰고 있음

• 우리는 더 체계적으로 함!

Contribution

• Base모델과 Post 모델을 체계적, 기계적으로 분석
  • Instruct model (모든 post-training 끝난 모델) , SFT model (fine-tuning만 함) 가지고 분석

• 다음의 관점에서 post-training 전후 비교
  • 지식 저장과 표현
  • 내부적 truthfulness
  • Refusal
  • Confidence

• 아래는 요약 figure
  • 지식 → 저장하는 위치, 표현 방식 유지되거나 강화됨
  • 내부적 Truthfulness 방향 유지됨
    • Latent vector가 유지됨
  • Refusal은 방향 바뀜
  • Confidence는 entropy neuron이 아닌 다른데서 강화됨
    • Confidence가 보정된다고 entropy neuron이 달라지는건 아님

Knowledge Storage and Representation

• Post-training 전후로 지식 저장 위치가 바뀌는지, 지식 표현 벡터가 크게 달라지는지 확인하자!
  • 유지되면 좋은 것

실험 세팅

• 지식에 관한 T/F 문제로 테스트
  • E.g. “The city of New York is in the United States. This statement is:”

• 참/거짓 문장끼리의 hidden state 차이를 보고 지식 저장 위치를 판별함
  • E.g. “The city of Seattle is in France.” vs. “The city of Paris is in France.”

• 참 문장 줬을 때의 $l$번째 layer $i$번째 토큰의 hidden state $h^l_i(s)$, 거짓 문장은 $h_i^l(\hat{s})$

• 거짓 문장 주고 hidden state를 교체해서 답변이 False에서 True로 바뀔 때(정확히는 확률 보고 결정) 그 위치에 지식이 저장되어 있던 것!
  • $M_i^l(s,\hat{s}) := \log\Big[\frac{P(\text{“TRUE”})}{P(\text{“FALSE”})}\ \big|\ \text{patching}(h_i^l(s), h_i^l(\hat{s}))\Big]$ 이게 크면 지식 저장이라는 뜻
  • $\tilde{M}_i^l=\frac{1}{|D|}\sum_{(s,\hat{s})\in D} M_i^l(s,\hat{s}),\qquad M=\text{normalize}(\tilde{M}).$
    이렇게 데이터셋에 대해 정규화시켜서 계산

Internal Belief of Truthfulness

• 내부 truthfulness 방향이 유지되는지 확인하자!
  • 유지되면 좋은것

실험 세팅

• 이번엔 참문장하고 거짓문장 사이의 hidden state 차이로 truthfulness 방향성 계산
  • $t^l=\frac{1}{|D_{\text{train}}^{\text{true}}|}\sum_{s\in D_{\text{train}}^{\text{true}}} h_i^l(s)\;-\;\frac{1}{|D_{\text{train}}^{\text{false}}|}\sum_{s\in D_{\text{train}}^{\text{false}}} h_i^l(s)$
  • 여기서 $i$는 마지막 토큰을 씀!
  • $l$은 truthfulness가 가장 강하게 인코딩되는 layer 선택

Refusal

• 내부 refusal 방향이 유지되는지 확인하자!
  • 유지되면 안 좋은 것 (refusal 잘하도록 하는게 post-training이기 때문)

실험 세팅

• $D^{train}_{harmful}$, $D^{train}_{harmless}$를 이용해서, truthfulness와 마찬가지로 이번엔 refusal 방향 $r$계산

• Refusal을 유도할 때는 가장 강하게 작동하는 layer $l$에서 hidden state에 refusal 방향을 더함
  • $\tilde{h}_l \leftarrow h_l + r_l$

• Refusal을 감소할 때는 모든 레이어에서 refusal 방향을 뺌
  • $\tilde{h} \leftarrow h - \hat{r}\hat{r}^{\top}h$

Confidence

• Post-training하면 base모델하고 다른 confidence를 가지게 됨(토큰 생성 확률 스케일이 보정됨)

• 그리고 모델 내 엔트로피 neuron이 confidence를 조정한다고 알려짐
  • https://arxiv.org/abs/2401.12181
  • https://arxiv.org/abs/2406.16254
  • 이 뉴런들은 가중치(weight norm)이 크고, unembedding matrix와의 composition이 낮기에,
    확률 분포의 순위는 바꾸지 않으면서 스케일을 조절함

실험 세팅

• 엔트로피 neuron들을 식별하자!

• 마지막 MLP layer의 각 뉴런에 대해, 출력 가중치 $w_{out}$을 unembedding matrix $W_U$로 vocab space에 projection해 logit attribution을 계산함
  • 이 projection은 해당 뉴런이 최종 예측 logit에 미치는 영향을 근사함

• 그리고 projection의 분산 LogitVar 계산
  • $\text{LogitVar}(w_{out})=\text{Var}\left(\frac{W_U w_{out}}{\|W_U\|_{dim=1}\,\|w_{out}\|}\right)$

• LogitVar가 낮다는 것은 특정 토큰을 밀어주는게 아니라 어휘 전체에 대해 가중치를 부여함

• 엔트로피 neuron은 보통 LogitVar는 낮고 weight norm은 큼
  • Weight norm이 가장 큰 neuron 상위 25%를 선택하고, 그 부분집합에서 LogitVar가 가장 낮은 10개
    neuron을 마지막 MLP 층의 엔트로피 neuron으로 식별