로그it 렌즈는 머신 러닝 해석 가능성에서 트랜스포머 기반 대규모 언어 모델의 내부 표현을 검사하는 데 사용되는 기법입니다. 이 기법은 특정 레이어의 은닉 상태를 가져와 모델의 최종 언임베딩 행렬(은닉 상태를 어휘에 대한 로짓으로 매핑하는 행렬)을 적용한 다음, 소프트맥스를 적용하여 확률 분포를 얻는 방식으로 작동합니다. 이 분포는 모델이 해당 레이어에서 멈췄을 때 예측했을 결과를 나타내며, 예측이 레이어를 거치며 어떻게 진화하는지 보여줍니다. 이 방법은 2020년 nostalgebraist에 의해 소개되었으며, 이후 GPT-2 및 GPT-3과 같은 모델의 내부 계산을 연구하는 표준 도구가 되었습니다.
프로빙 분류기와 달리 추가적인 훈련이 필요한 경우가 많은데, 로짓 렌즈는 모델 자체의 출력 헤드를 사용하므로 매개변수 추가나 훈련이 필요 없는 접근 방식입니다. 잔차 스트림 모델에서 특히 효과적인데, 각 레이어의 은닉 상태가 더해지는 구조 덕분에 직접 투영하여 해석할 수 있기 때문입니다. 이 기법은 모델이 특정 예측을 확정하는 시점을 식별하고, "늦은" 디코딩과 "이른" 디코딩을 구분하며, 사실 회상이나 산술 추론과 같은 현상을 분석하는 데 사용됩니다. 그러나 가중치가 묶인(공유된) 임베딩이 없는 모델이나 복잡한 레이어 정규화 및 비선형 변환이 있는 모델에서는 적용이 제한적일 수 있습니다.
메커니즘 및 구현
로짓 렌즈는 트랜스포머의 잔차 스트림을 기반으로 작동합니다. 일반적인 트랜스포머에서 각 레이어는 출력을 잔차 스트림에 추가하므로, 레이어 \( l \)에서의 은닉 상태는 초기 임베딩과 그 이전 레이어 출력들의 합입니다. 최종 로짓은 (레이어 정규화가 있다면 적용 후) 언임베딩 행렬 \( W_U \)를 곱하여 계산됩니다. 로짓 렌즈는 이 동일한 언임베딩을 중간 레이어의 은닉 상태에도 적용하며, 분포 이동을 보정하기 위해 최종 레이어 정규화(또는 학습된 아핀 변환)를 적용하기도 합니다.
가중치가 묶인 모델(입력 임베딩과 출력 언임베딩이 동일한 행렬을 공유)에서는 투영이 간단합니다. 별도의 언임베딩 행렬을 가진 모델에서도 언임베딩을 사용할 수 있으면 작동합니다. 결과적으로 각 레이어에 대한 어휘 확률 분포 시퀀스를 얻을 수 있으며, 이를 히트맵으로 시각화하거나 후보 토큰 간의 "로짓 차이"와 같은 지표를 계산하는 데 사용할 수 있습니다.
해석 가능성 응용
로짓 렌즈는 다양한 해석 가능성 작업에 적용됩니다. 일반적인 용도 중 하나는 레이어를 거치며 예측이 어떻게 발전하는지 추적하는 것입니다. 예를 들어, "프랑스의 수도는 __"라는 문장을 완성하는 모델에서 로짓 렌즈는 초기 레이어에서는 "파리"에 대한 확률이 낮다가 특정 레이어에서 급격히 높아지는 것을 보여줄 수 있으며, 이를 통해 모델이 관련 사실을 검색하는 시점을 파악할 수 있습니다. 연구자들은 이를 사용하여 "유도 헤드"와 같은 특정 행동에 기여하는 주의 패턴을 식별하기도 합니다.
또 다른 응용으로는 "늦은" 디코딩을 감지하는 것입니다. 모델이 초기 레이어에서 특정 토큰을 예측했다가 나중에 변경하는 경우, 로짓 렌즈를 통해 최종 예측이 "고정"되는 시점을 확인할 수 있습니다. 이는 모델의 확신도와 개입 설계에 대한 통찰을 제공합니다. 또한 산술 추론과 같은 작업에서 중간 레이어가 중간 합계나 자리올림 값을 표현하는지 보여줌으로써 다단계 계산 과정을 연구하는 데도 사용됩니다. 이는 Anthropic과 같은 연구 그룹의 "유도 헤드" 및 "회로" 프레임워크와 같은 메커니즘 해석 가능성 연구로 이어졌습니다.
한계 및 비판
로짓 렌즈는 유용하지만 몇 가지 한계가 있습니다. 첫째, 언임베딩 행렬이 모든 레이어의 은닉 상태에 대해 의미 있는 투영이라고 가정하지만, 레이어 정규화와 같은 변환이 적용된 모델에서는 이 가정이 성립하지 않을 수 있습니다. 사전 정규화(pre-norm) 아키텍처를 가진 모델에서는 중간 상태에 최종 레이어 정규화를 적용해야 하지만, 항상 충분하지는 않습니다.
둘째, 로짓 렌즈는 어휘가 매우 크거나 다중 토큰 예측이 필요한 작업에서는 효과가 떨어질 수 있습니다. 단일 다음 토큰 확률만 제공하므로, 모델의 전체 상태를 포착하지 못합니다. 또한 주의 메커니즘과 피드포워드 하위 레이어의 기여를 개별적으로 분리하지 못하며, 단지 더해진 결과만 보여줍니다.
셋째, 별도의 출력 헤드를 사용하는 모델이나 잔차 스트림 구조가 다른 모델에서는 잘못된 결과를 초래할 수 있습니다. 이러한 경우 투영된 로짓이 노이즈가 많거나 정보를 제공하지 못할 수 있습니다. 이를 개선하기 위해 "튜닝된 렌즈"와 같은 변형이 제안되었는데, 이는 레이어별 아핀 변환을 학습하여 정렬을 개선하지만 추가적인 훈련 데이터가 필요합니다.
다른 해석 가능성 방법과의 관계
로짓 렌즈는 프로빙 분류기, 활성 패칭, 인과 추적과 같은 더 넓은 해석 가능성 기법 패밀리의 일부입니다. 프로빙 분류기는 은닉 상태에서 특정 특징을 예측하도록 선형 또는 비선형 모델을 훈련시키지만, 이는 지도 학습이 필요하고 모델의 실제 계산을 반영하지 않을 수 있습니다. 반면 로짓 렌즈는 모델 자체의 출력 헤드를 사용하므로 더 직접적입니다.
활성 패칭은 특정 레이어의 은닉 상태를 변경하여 출력에 미치는 영향을 측정함으로써 인과 관계를 분석합니다. 로짓 렌즈는 패칭 결과를 해석하는 데 보완적으로 사용될 수 있습니다. 인과 추적은 David Kaplan 등에 의해 대중화된 기법으로, 손상된 실행에서 특정 레이어가 특정 사실을 담당하는지 식별합니다. 로짓 렌즈는 각 레이어에서 예측된 토큰을 보여줌으로써 이러한 분석을 보완할 수 있습니다.
또한 "잔차 스트림 관점"은 트랜스포머를 일련의 덧셈 업데이트로 간주하며, 로짓 렌즈는 이 관점과 자연스럽게 맞아떨어집니다. 이는 "Transformer Debugger"와 같은 시각화 도구와 해석 가능성 라이브러리에 영감을 주었습니다.
확장 및 변형
로짓 렌즈의 여러 확장이 제안되었습니다. "튜닝된 렌즈"는 표준 로짓 렌즈가 실패하는 모델에서 성능을 개선하기 위해 레이어별 아핀 변환을 학습합니다. "소프트맥스 렌즈"는 로짓에 온도를 적용하여 확률 분포를 더 선명하거나 평평하게 만들어 저확률 예측을 더 쉽게 관찰할 수 있게 합니다.
다중 출력 헤드나 디코더를 가진 모델에서는 각 헤드에 로짓 렌즈를 개별적으로 적용할 수 있습니다. 인코더-디코더 모델에서는 디코더의 은닉 상태에 적용하는 것이 일반적이지만, 인코더 상태에 적용하는 방법은 덜 명확합니다. 최근 연구에서는 이미지 토큰이나 다른 양식을 포함하는 다중 모달 모델에도 로짓 렌즈를 적용하려는 시도가 있습니다.
계산 고려 사항
로짓 렌즈는 계산적으로 효율적입니다. 각 레이어에서 단일 행렬 곱셈(언임베딩)과 소프트맥스만 필요하므로, 모델의 순전파 비용에 비해 무시할 수 있는 수준입니다. 모델에 \( L \)개의 레이어와 어휘 크기 \( V \), 은닉 차원 \( d \)가 있을 때, 추가 비용은 \( O(L \cdot V \cdot d) \)입니다.
그러나 모든 레이어의 로짓을 저장하는 것은 메모리 집약적일 수 있습니다. 특히 어휘가 큰 모델에서는 각 레이어마다 \( V \)-차원 벡터를 저장해야 하므로, 실제로는 각 레이어에서 상위 \( k \)개 토큰만 계산하여 메모리 사용을 줄이는 방법이 사용됩니다. 로짓 렌즈는 TransformerLens 라이브러리와 Hugging Face의 해석 가능성 도구 등 여러 오픈소스 라이브러리에 구현되어 있습니다.
향후 방향
딥 러닝 모델이 계속 확장됨에 따라 로짓 렌즈와 같은 해석 가능성 도구의 중요성도 커지고 있습니다. 향후 연구는 혼합 전문가나 상태 공간 모델과 같은 새로운 아키텍처에 로짓 렌즈를 적용하는 방법, 자동화된 해석 가능성 파이프라인에 통합하는 방법에 초점을 맞출 수 있습니다. 또한 로짓 렌즈는 모델이 유해한 출력을 "생각"하는 시점을 감지하는 데 사용될 수 있어, 정렬 전략을 개발하는 데 기여할 수 있습니다.
전반적으로 로짓 렌즈는 신경망의 블랙박스를 들여다보는 간단하면서도 강력한 도구로 남아 있으며, 메커니즘 해석 가능성 분야가 성숙해짐에 따라 그 응용 범위는 더욱 확장될 것으로 보입니다.