main page
interpreting GPT: the logit lens — LessWrong
This post relates an observation I've made in my work with GPT-2, which I have not seen made elsewhere. …
www.lesswrong.com
code
logit lens on non-gpt2 models + extensions
Colaboratory notebook
colab.research.google.com
Overview
“Logit Lens”란?
- GPT의 마지막 층에서는 hidden activation(은닉 벡터)에 출력 projection 행렬 Wᵀ를 곱해 logits (어떤 토큰이 나올 확률)을 예측함
- “Logit Lens”는 이 같은 projection 연산을 중간층의 activation에도 적용해보는 방법즉, GPT의 각 층이 “지금까지 보고 있는 입력으로 다음 토큰을 어떻게 예측하고 있는지”를 살펴보는 일종의 확률적 해석 도구(lens)
$$ \text{logits}^{(I)} = h^{(I)} W^\top $$
distribution over the layers
- Early layers:
- → 예: “The cat sat on the ___” → early layer는 “mat”, “table”, “floor” 모두 비슷한 확률.
- 분포가 “의미 없는 잡음”처럼 보이거나, 단어의 품사/형태 정도 생각
- Middle layers:→ “mat”의 확률이 높아짐.
- 문맥적 단어 후보가 정리되고 “꽤 그럴듯한 예측”을 함
- Late layers:
- 후보들 중 하나로 확신이 높아지고 최종 logits 분포로 수렴
iterative refinement 한 방식으로 각 layer는 이전 layer의 예측을 조금 더 정교하게 보정
https://arxiv.org/pdf/1807.03819 : 같은 연산을 반복하며 점진적으로 추론하는 구조 유사
thinks in predictive space
background on GPT's structure
Input
vocabulary 크기 (Byte Pair Encoding 기반)
$$ N_v=50,257 $$
Output
- 모델 출력 = 다음 토큰의 확률 분포
- 즉, 입력 시퀀스가
- $$ [t_1, t_2, t_3, ..., t_n] $$
- 이라면, 모델은 각 위치 i에 대해 “다음 토큰 t_{i+1}”의 확률 분포를 예측
- $$ P(t_{i+1} \mid t_1, ..., t_i) $$
- 따라서 GPT는 “입력 → 한 칸 오른쪽으로 이동한 출력” 관계로 훈련
- 이걸 흔히 shifted prediction이라고 함
Vocabulary space ↔ Embedding space 변환
Vocabulary space
- 50,257개의 토큰 각각은 one-hot 벡터로 표현하기엔 너무 큼
Embedding space
- GPT는 훨씬 작은 차원의 벡터 공간(예: N_e=1600)에서 연산
- W: “Embedding matrix”
- 토큰을 벡터로 변환:x_i = W_{t_i}
- 모델의 최종 출력 logits을 다시 vocab으로 변환할 때도 같은 행렬의 전치 W^⊤ 사용
- 이게 바로 “weight tying”
$$ W \in \mathbb{R}^{N_v \times N_e} $$
Transformer block 구조
- Input projection:
- $$ x_i^{(0)} = W_{t_i} $$
- 각 토큰은 W를 통해 1600차원 벡터로 매핑
- Stack of Transformer blocks:
- 입력 벡터는 여러 블록(layer)을 통과
- 각 블록은 self-attention + feed-forward 구조이며, 입력과 출력의 차원은 동일(예: GPT-2 1558M은 L=48 layers)
- 즉, 각 블록은 “정보를 조금 더 가공하는 동일 차원 변환기”
- $$ x_i^{(l)} = \text{Block}^{(l)}(x_i^{(l-1)}), \quad l=1,\dots,L $$
- Output projection:
- 마지막 블록의 출력 x_i^{(L)}은 W^⊤로 vocab 공간에 다시 투사
- $$ \text{logits}_i = x_i^{(L)} W^\top $$
- Softmax:
- 각 토큰 위치에서 logits에 softmax를 취하면, 전체 어휘에 대한 다음 단어 확률 분포로 변환
- $$ P(t_{i+1} \mid t_{\le i}) = \text{softmax}(\text{logits}_i) $$
Logit Lens
Logit Lens는 바로 이 마지막 projection 행렬 W^⊤ 를 중간층 activation x_i^{(l)}에 적용해보는 것
$$ \text{Intermediate logits}^{(l)} = x_i^{(l)} W^\top $$
the logit lens
logits
- “We train GPT-3…” → 초반엔 “000?” (아직 추측 단계)
- “GPT-3, an…” → “enormous?”, “massive?” (대체로 문맥상 맞는 추측)
- “We train GPT-3, an aut…” → “oreceptor?” (이후 “autoregressive”로 수렴)
초기 층의 예측은 틀리지만 의미적으로 합리적이며,깊어질수록 점점 정답에 근접
ranks
Top-1만 보는 것은 지나치게 단순하므로, 각 층의 분포 전체를 살펴
- Rank 1: 이미 그 단어를 top-1로 예측
- Rank 3: 상위 후보 중 하나
- Rank 10~100: 여전히 고려 대상
- Rank 50,000: 전혀 후보 아님
결과
- 중간층에서도 이미 최종 예측 토큰이 상위권에 포함되는 경우가 많음
- 즉, 모델의 불확실성(uncertainty) 은 중간층까지만 가도 크게 줄어듦
- 마지막 층들은 주로 상위 후보들의 순서만 미세 조정
KL divergence and input discarding
두 확률 분포 P와 Q의 차이를 측정하는 지표
$$ D_{KL}(P \Vert Q) = \sum_i P(i) \log \frac{P(i)}{Q(i)} $$
- P: 중간 layer의 확률 분포
- Q: 최종 layer의 확률 분포
각 층의 분포가 최종 분포와 얼마나 비슷한지를 연속적으로 확인
결과
- 첫 번째 층 이후 급격한 변화:→ 즉, 입력 정보는 거의 즉시 “예측 공간”으로 변환됨.
- 입력 토큰의 embedding은 “최종 분포”와 완전히 다르며, 단 한 층(첫 transformer block)만 지나도 이미 “출력 쪽”에 가까운 형태로 바뀜.
- 이후에는 점진적 수렴:
- 첫 급격한 변환 이후, 각 층의 분포는 부드럽게 최종 분포로 수렴함.
why? / is this surprising?
Residual Network 구조
$$ x_{l+1}=x_l+f(x_l) $$
구조 자체가 기저(basis)를 유지한 채 점진적으로 수정 하는 형태
따라서 layer마다 큰 회전이나 변환 없이, 전체 네트워크가 하나의 일관된 표현 공간(embedding basis) 을 공유
Weight Decay (L2 정규화)
가중치들의 L2 노름(크기)을 작게 유지하도록 유도
결과적으로, 한 층이 갑자기 큰 변화를 주기보다 여러 층에 걸쳐 작은 변화들을 누적하는 게 더 유리
$$ (\sum \Delta x)^2 > \sum (\Delta x^2) $$
왜 input representation은 바로 사라질까?
모델은 입력을 “보관하지” 않고, 입력 → 예측 공간으로 즉시 변환 하여 그 위에서만 생각(thinking)함
모름