[논문 리뷰] Emergent Linear Representations in World Models of Self-Supervised Sequence Models

Neel Nanda, Andrew Lee외 1명|arXiv (Cornell University)|2023. 09. 02.

Neural Networks and Reservoir ComputingComputer Science인용 수 3

한 줄 요약

이 논문은 자기지도 학습 시퀀스 모델인 OthelloGPT가 절대 색상이 아닌 상대적 색상—'내 것', '상대의 것', '비어 있음'—을 선형으로 인코딩하여 체스판 상태를 표현함을 보여준다. 이러한 선형 방향을 조사함으로써, 간단한 벡터 산술을 통해 모델의 내부 상태를 정밀하게 제어할 수 있으며, 이는 기존의 기울기 기반 방법보다 더 해석 가능하고 효과적인 간섭을 가능하게 한다.

ABSTRACT

How do sequence models represent their decision-making process? Prior work suggests that Othello-playing neural network learned nonlinear models of the board state (Li et al., 2023). In this work, we provide evidence of a closely related linear representation of the board. In particular, we show that probing for "my colour" vs. "opponent's colour" may be a simple yet powerful way to interpret the model's internal state. This precise understanding of the internal representations allows us to control the model's behaviour with simple vector arithmetic. Linear representations enable significant interpretability progress, which we demonstrate with further exploration of how the world model is computed.

연구 동기 및 목표

• 자기지도 시퀀스 모델인 OthelloGPT와 같은 모델이 내부 세계 상태를 어떻게 표현하는지 조사하는 것, 특히 이러한 표현 방식이 선형인지 여부를 밝혀내는 것.
• 이전 연구에서 비선형 체스판 상태 인코딩을 주장한 바와 현재의 선형 표현 발견 간의 명백한 모순을 해결하는 것.
• 내부 활성화에 대한 벡터 산술을 통해 더 단순하고 해석 가능한 방식으로 모델 행동을 제어하는 방법을 개발하는 것.
• 모델이 유효한 수를 어떻게 계산하고, 빈 칸이나 뒤집힌 타일을 어떻게 감지하는지에 대한 기계적 통찰을 제공하는 것.
• 세계 모델에서 선형 표현이 어떻게 나타나는지와 그 해석 가능성 및 제어 가능성에 대한 영향을 탐구하는 것.

제안 방법

• 저자들은 체스판 상태 표현을 절대 색상(검은색, 흰색, 비어 있음)이 아닌 상대적 플레이어 시점(내 것, 상대의 것, 비어 있음)으로 재정의하여, 턴 순서에 따라 시간에 따라 변하는 방식으로 접근한다.
• 각 토큰 위치에서 상대적 색상 상태를 예측하기 위해 선형 프로브를 훈련하여, 예를 들어 타임스텝 4에서 98.3%의 정확도를 기록한 근사 완벽한 성능을 달성한다.
• 내부 활성화에 대해 학습된 선형 방향(Mine, Yours, Empty)에 투영함으로써 모델의 내부 믿음 상태를 직접 조작할 수 있도록 벡터 산술을 적용한다.
• 수정된 체스판 상태를 시뮬레이션하기 위해 활성화를 편집하고, 수정된 믿음 상태와 일치하는 유효한 수를 생성함으로써 이 방법의 타당성을 검증한다.
• 다른 선형 특징, 예를 들어 각 타임스텝에서 뒤집힌 타일을 탐지하기 위한 프로브를 추가로 수행하고 잔차 스트림 다이내믹스를 분석하여 이동 예측을 위한 다수의 회로를 탐지한다.
• 선형 및 비선형 프로브 성능을 비교하여, 비선형 프로브가 절대 색상 레이블에서 높은 정확도를 달성하는 반면, 상대적 색상에 대한 선형 프로브가 거의 슈퍼리어한 성능을 보임을 확인한다.

Figure 1: The emergent world models of OthelloGPT are linearly represented. We find that the board states are encoded relative to the current player’s colour ( Mine vs. Yours ) as opposed to absolute colours ( Black vs. White ).

실험 결과

연구 질문

• RQ1 OthelloGPT 내부의 Othello 체스판 상태 표현은 선형적인 구조를 보이며, 만약 그렇다면 어떤 형태로 나타나는가?
• RQ2 활성화 공간 내 선형 방향에 대한 벡터 산술을 통해 모델의 내부 세계 상태와 행동을 인과적으로 수정할 수 있는가?
• RQ3 왜 자기지도 세계 모델에서 선형 표현이 나타나며, 이는 이전에 보고된 비선형 표현과 어떻게 관련이 있는가?
• RQ4 선형 프로브를 통해 모델이 유효한 수를 계산하고, 빈 칸이나 뒤집힌 타일을 감지하는 방식에 대해 어떤 기계적 통찰을 얻을 수 있는가?
• RQ5 유효한 수를 예측하기 위한 다수의 별개의 회로가 모델에 존재하는가, 그리고 이들은 선형 특징 표현과 어떻게 관련이 있는가?

주요 결과

• 모델은 절대 색상(검은색, 흰색, 비어 있음)이 아닌 현재 플레이어의 턴에 따라 변하는 상대적 표현—'내 것', '상대의 것', '비어 있음'—을 사용하여 체스판 상태를 인코딩하며, 이는 근사 완벽한 선형 프로브 성능(예: 타임스텝 4에서 98.3% 정확도)을 가능하게 한다.
• 내 것/상대의 것/비어 있음 방향에 대한 선형 투영은 내부 세계 모델을 재구성하는 데 거의 완벽한 성능을 보이며, 절대 색상 레이블에 대한 비선형 프로브보다 뛰어난 성능을 기록한다.
• 내부 활성화에 대한 단순한 벡터 산술을 통해 모델의 행동을 효과적이고 해석 가능한 방식으로 제어할 수 있으며, 예를 들어 수정된 체스판 상태에서 유효한 수를 두도록 유도할 수 있다.
• 기울기 계산 없이도 인과적 간섭을 가능하게 하여, 기존의 기울기 기반 편집 기법에 비해 더 단순하고 투명한 대안을 제공한다.
• 모델는 특히 게임의 말판에서 다수의 이동 예측을 위한 회로를 보이며, 타일 뒤집기 이벤트의 선형 표현도 탐지 가능하다.
• 이러한 발견들은 선형 표현이 이전에 상상한 것보다 더 근본적인 역할을 할 수 있으며, 잔차 연결이 레이어 간 특징의 반복적 정밀화를 가능하게 한다는 것을 시사한다.

Figure 2: Intervening methodology: we intervene by adding either Empty , Mine , or Yours directions into each layer of the residual stream. Red squares in each board indicate the tiles that have been intervened, teal tiles indicate new legal moves post-intervention that the model predicts.