Logit Lens는 24개 레이어로 구성된 모델의 12번 레이어에서 잔차 스트림 (residual stream)을 가져와, 최종 레이어 정규화 (layer norm)와 언임베딩 행렬 (unembedding matrix)을 적용하여 토큰 분포를 읽어냅니다. 즉, 모델이 그 지점에서 답변해야 한다면 무엇이라고 말할지를 보여줍니다. 이 코드는 약 30줄 정도이며, 모델의 내부 구조를 들여다볼 수 있는 가장 저렴한 창(window)입니다.

아이디어 (The idea)

트랜스포머 (transformer)의 잔차 스트림 (residual stream)은 네트워크의 전체 길이를 따라 흐릅니다. 각 레이어는 이를 읽고, 무언가를 계산한 뒤, 그 결과를 다시 더합니다. 마지막에는 하나의 레이어 정규화 (layer norm)와 하나의 행렬 곱셈을 통해 최종 잔차 상태를 어휘 사전 상의 로짓 (logits over the vocabulary)으로 변환합니다.

Logit Lens는 동일한 최종 연산을 중간 단계에서 조기에 적용합니다. 아키텍처 상에서 이를 수행할 권한이 있다고 명시되어 있지는 않습니다. 즉, 중간 상태는 최종 상태가 아닙니다. 하지만 많은 모델에서 그 결과는 판독 가능합니다. 8번 레이어의 상위 토큰들은 막연하게 연관되어 있고, 16번 레이어에서는 주제에 부합하며, 22번 레이어에 이르면 정답에 도달합니다. 이 기술은 2020년 nostalgebraist라는 필명의 연구자에 의해 발표되었으며 그곳에서 이름이 붙여졌습니다.

이것이 가능한 이유 (Why this is allowed at all)

아키텍처의 두 가지 특성이 이 기술을 단순한 속임수 이상으로 만듭니다.

첫째, 잔차 스트림 (residual stream)은 가산적 (additive)입니다. $n$번째 레이어의 출력은 $h_n = h_{n-1} + f_n(h_{n-1})$이므로, 모든 깊이의 상태는 동일한 공간에 존재하며 기여분들의 누적 합계입니다. 레이어 사이에 되돌려야 할 기저 변환 (basis change)이 없습니다.

둘째, 언임베딩 (unembedding)은 해당 공간에서 어휘 로짓 (vocabulary logits)으로 가는 고정된 선형 사상 (linear map)입니다. 이를 부분 합에 적용하는 것은 해당 부분 합에 대한 정당한 투영 (projection)입니다. 즉, 지금까지 기록된 기여분들이 어떤 토큰을 가리키고 있는지를 알려줍니다. 해석상의 도약은 이를 "모델이 현재 믿고 있는 것"으로 취급하는 것인데, 이는 중간 상태가 언임베딩이 기대하는 것과 대략적으로 동일한 좌표계에 있다고 가정하는 것입니다. 아래에서 논의할 모델들에서는 바로 이 가정이 실패합니다.

스크립트 (The script)

CPU 환경에서 gpt2를 대상으로 완전하게 실행 가능합니다. 모델별로 지정되는 두 가지 명칭은 최종 정규화 (final norm)와 언임베딩 (unembedding)입니다. GPT-2의 경우 각각 transformer.ln_f와 lm_head이며, Llama 스타일 모델의 경우 model.norm과 lm_head입니다. 확실하지 않다면 model을 출력하여 마지막 두 모듈을 확인하십시오.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

...

엔트로피 (entropy) 열은 토큰만큼이나 가치가 있습니다. 모델이 결정을 내릴수록(commits) 엔트로피는 낮아지며, 엔트로피가 급격히 떨어지는 지점이 바로 결정이 내려진 지점입니다. 동일한 형태의 두 프롬프트에서 결정이 발생하는 깊이가 다르다는 것은 시각화의 오류가 아니라 모델에 대한 진정한 관찰 결과입니다.

여기서 발생하는 표준적인 오류는 'Off-by-one' (하나 차이 오류)입니다. out.hidden_states는 0번 인덱스가 임베딩 출력값이기 때문에 n_layers + 1개의 항목을 가집니다. 위 스크립트의 레이블은 이를 고려하고 있습니다. 직접 코드를 작성한다면, 마지막 항목이 모델의 실제 로짓 (logits)을 재현하는지 확인하십시오. 정확히 일치해야 합니다.

출력 읽기 (Reading the output)

GPT-2에서 사실적 완성 (factual completion) 작업을 수행할 때, 일반적으로 세 가지 양상 (regimes)이 나타납니다. 초기 레이어는 고엔트로피 노이즈를 생성합니다: 빈번하게 등장하는 토큰, 문장 부호 등 프롬프트와 무관한 것들입니다. 중간 레이어는 주제와 인접한 무언가를 보여줍니다. 예를 들어 프랑스(France)에 관한 프롬프트라면 다른 지명들이 나타납니다. 후기 레이어는 정답으로 수렴하며, 보통 끝에 도달하기 몇 레이어 전부터 수렴이 시작됩니다. 그 이후에는 분포가 변하기보다는 주로 날카로워집니다 (sharpens).

이를 통해 수행할 가치가 있는 두 가지 작업이 있습니다. 정답이 처음으로 상위 5위 안에 들어오는 레이어를 찾고, 프롬프트를 변경하며 그것이 어떻게 이동하는지 관찰하십시오. 어려운 프롬프트일수록 결정이 더 늦게 일어납니다. 그리고 역전 현상 (reversals)을 주시하십시오. 레이어 14에서는 1위였던 토큰이 레이어 20에 이르러 사라졌다면, 이는 이후의 컴포넌트가 해당 토큰을 능동적으로 억제 (suppressed)했음을 의미하며, 이는 활성화 패칭 (activation patching)을 통해 추적해 볼 만한 단서입니다. 억제는 실제로 존재합니다. 회로 (circuit) 연구를 통해 모델의 나머지 부분이 제안한 후보를 억제하는 역할을 하는 헤드 (heads)들이 반복적으로 발견되었습니다.

노이즈가 발생하는 경우

Logit lens는 아키텍처 독립적이지 않으며, 이 점이 대부분의 튜토리얼에서 누락되는 부분입니다. Belrose와 동료들은 2023년에 여러 모델에서 일반적인 logit lens가 중간 레이어(intermediate layers)에서 본질적으로 해석 불가능한 출력을 생성한다고 보고했습니다. 이는 GPT-2의 양상이 약화된 버전이 아니라, 아예 노이즈(noise)에 가깝습니다. 이 기술은 특정 모델 제품군에서는 잘 작동하지만, 다른 모델군에서는 실패합니다.

그 이유는 앞 섹션에서 언급한 가정 때문입니다. 중간 표현(intermediate representations)이 unembedding이 읽어들이는 기저(basis)로 표현되어야 한다는 제약 조건은 아무것도 없습니다. 모델은 중간 상태를 회전(rotated)되거나 재조정(rescaled)된 형태로 완벽하게 유지하다가, 마지막 몇 개의 레이어에서만 출력 기저(output basis)로 가져올 수 있습니다. 이런 일이 발생하면 투영(projection)은 무의미해지며 — 이 부분이 위험한 지점입니다 — 그것이 무의미해 보이지도 않습니다. 마치 모델이 20개의 레이어 동안 문장 부호(punctuation)에 대해 생각하고 있는 것처럼 보이기 때문입니다.

실무적 규칙: 이전에 사용해 보지 않은 모델에 대해 logit lens로 결론을 내리기 전에, 정답을 알고 있으며 정답이 대략 어디쯤 나타나야 하는지 아는 프롬프트(prompts)에서 먼저 확인하십시오. 만약 정답 토큰이 마지막 레이어에 도달할 때까지 상위 10위(top ten) 안에 들지 않는다면, 중간 레이어가 무엇을 말하는 것처럼 보이든 상관없이 해당 모델에서 lens는 작동하지 않는 것입니다.

Tuned lens (조정된 렌즈)

2023년 Belrose와 동료들이 제시한 해결책은 다음과 같습니다: unembedding을 직접 적용하는 대신, 각 레이어마다 해당 레이어의 잔차 상태(residual state)를 최종 레이어의 기저로 변환하는 작은 어파인 사상(affine map)을 학습시킨 다음, unembed를 수행합니다. 이 사상은 모델 자체의 최종 출력 분포(final output distribution)와 일치하도록 학습되므로, 가정이 아닌 교정(calibrated)된 상태가 됩니다.

이것이 실제로 유용한 점

• 의사결정이 일어나는 위치 파악 (Locating where a decision is made). 비용이 많이 드는 패칭 (patching) 실험을 하기 전 수행할 수 있는 저렴한 1차 탐색입니다. 만약 정답이 18번 레이어에서 나타난다면, 모든 레이어를 대상으로 하는 대신 14번에서 20번 레이어를 패칭하세요.
• 프롬프트 비교 (Comparing prompts). 서로 다른 깊이 (depth)에서 동일한 정답을 생성하는 두 가지 문구는, 각 문구가 모델의 작업을 얼마나 절약해 주었는지에 대한 정보를 제공합니다.
• 억제 현상 포착 (Spotting suppression). 후보군(candidates)의 수치가 상승했다가 다시 하락하는 현상은, 특정 컴포넌트가 의도적인 작업을 수행하고 있음을 보여주는 가장 명확한 렌즈 가시적 (lens-visible) 신호입니다.
• 빠른 직관 형성 (Building intuition quickly). 이 스크립트를 30분 동안 사용하는 것이 트랜스포머 (transformer)에서 깊이 (depth)가 무엇을 의미하는지 이해하는 데 있어 그 어떤 다이어그램보다 더 큰 도움이 됩니다.

적합하지 않은 용도: 어떠한 인과적 주장 (causal claim)도 할 수 없습니다. 렌즈는 상태 (state)를 읽을 뿐, 그 상태가 무언가를 유발했다는 사실을 입증하지는 않습니다. 전체적인 그림을 보고 싶다면 전체 계측 가이드 (the full-instrumentation walkthrough)에 있는 계측 (instrumentation) 기법과 병행하고, 검토를 견뎌낼 수 있는 주장을 하고 싶다면 패칭 (patching)과 병행하세요.

관련 항목 (Related)

• 모델이 생각하는 과정을 관찰하기 위해 오픈 모델 계측하기 (Instrumenting an Open Model to Watch It Think)
• 활성화 패칭 (Activation Patching) 및 인과적 추적 (Causal Tracing)
• 회로 (Circuits): 하나의 동작을 엔드 투 엔드로 역공학하기 (Reverse-Engineering One Behaviour End to End)