기술 요약: 인컨텍스트 회상을 위한 예측 메커니즘의 분해

문제 정의

본 논문은 연속적인 인컨텍스트 학습(ICL) 작업(선형 회귀와 같은)과 이산적 연상 회상(MQAR, Multi-Query Associative Recall과 같은) 작업 사이의 간극을 메우기 위해 설계된 새로운 유형의 토이 문제군을 소개한다. 핵심 과제는 이산적으로 인터리빙된(interleaved) 시계열 시스템의 상태를 예측하는 것이다.

기저 데이터는 무작위 직교 선형 동역학 시스템($x_{i+1} = U x_i$, 여기서 $U$는 무작위 직교 행렬)에 의해 생성된 노이즈 없는 궤적들로 구성된다. 이 궤적들은 세그먼트 단위로 나뉘어 "심볼릭 구두점 레이블"(열기 및 닫기 토큰)에 의해 구분되어 단일 컨텍스트 윈도우 내에 인터리빙된다. 모델은 이 인터리빙된 트레이스(trace)에 대해 다음 토큰 예측을 수행하도록 훈련된다. 이 과제는 모델에게 다음을 요구한다:

1. 심볼릭 레이블을 기반으로 어떤 특정 시스템의 상태를 회상해야 하는지 **식별(Identify)**할 것.
2. 해당 상태에 올바른 기저 선형 동역학을 **적용(Apply)**하여 후속 관측치를 예측할 것.
3. 재개된 시퀀스를 **지속(Continue)**할 것.

저자들은 성공적인 성능을 위해서는 두 가지 뚜렷한 능력, 즉 시스템이 처음 나타났을 때 그 동역학을 학습하는 능력(ICL)과 컨텍스트 내에서 이전에 보았던 시스템의 동역학을 회상하는 능력(연상 회상)을 구분하는 것이 필요하다고 상정한다.

방법론

데이터 생성

저자들은 40,000개의 직교 행렬($U \in \mathbb{R}^{5 \times 5}$)과 그에 대응하는 초기 상태를 포함하는 훈련 라이브러리를 생성한다. 훈련 트레이스는 이러한 시계열 세그먼트들을 인터리빙하여 구축된다.

• 인터리빙: 세그먼트들은 잘리고 재조립되며, 각 시스템 인스턴스에는 무작위 "열기" 및 "닫기" 심볼릭 레이블이 할당된다. 결정적으로, 시스템과 심볼릭 레이블 사이의 매핑은 모든 훈련 예시마다 무작위로 변경되어, 모델이 정적인 레이블-시스템 연관성을 암기하는 것을 방지한다.
• 분포: 트레이스당 시스템의 수는 Zipf 분포를 따르며, 컷(cut)의 수는 Poisson 분포를 따른다. 이는 모델이 컨텍스트 내의 다른 여러 시스템 중에서 특정 시스템을 찾아내야 하는 "건초더미 속 바늘 찾기(needle-in-a-haystack)" 시나리오를 생성한다.

모델 아키텍처 및 훈련

저자들은 GPT-2 스타일의 트랜스포머(파라미터 수 212K ~ 10.7M 범위)를 평균 제곱 오차(MSE) 손실을 사용하여 다음 토큰 예측에 훈련시킨다. 입력 임베딩 차원은 57(심볼릭 레이블용 50, 시작용 1, 페이로드 플래그용 1, 상태 벡터용 5)이다. 출력 레이어는 5차원 상태 벡터를 예측한다.

평가 프로토콜

1. 비인터리빙 테스트(Uninterleaved Test): 컨텍스트에서 처음 본 시스템의 동역학을 학습하는 능력을 평가한다.
2. 건초더미 속 바늘 찾기(NiH) 테스트: 연상 회상을 평가한다. $N$개의 시스템이 제시된 후, 어떤 시스템을 재개할지를 나타내는 쿼리 심볼("바늘")이 등장한다. 성능은 쿼리 이후의 첫 번째 토큰(1-after) 및 후속 토큰(2-after 등)에 대해 측정된다.
3. 분포 외(OOD) 실험: 회상 메커니즘을 테스트하기 위해 네 가지 특정 실험이 설계되었다:
   • 잘못된 시퀀스로의 유도(Misdirection to Incorrect Sequence): 쿼리 레이블을 건초더미 내의 다른 시스템을 가리키도록 교체한다.
   • 동기화된 회전(Synchronized Rotations): 건초더미의 모든 시스템이 테스트 세그립의 첫 번째 인덱스에서 동일한 상태를 갖도록 강제하여, 상태 기반 변별을 제거한다.
   • 보이지 않은 시퀀스로의 유도(Misdirection to Unseen Sequence): 건초더미에 존재하지 않는 시스템에 해당하는 쿼리 레이블을 사용한다.
   • 보았던 시퀀스로의 유도(재시작)(Misdirection to Seen Sequence (Restart)): 실제 테스트 세그먼트가 새로운 미지의 시스템인 상황에서, 이전에 보았던 시스템을 가리키는 레이블을 사용한다.
4. 회로 분석(Circuit Analysis): 훈련된 가중치에 에지 프루닝(edge pruning)을 적용하여 "1-after" 토큰과 "2-after" 토큰을 예측하는 데 사용되는 특정 계산 회로를 격리한다.

주요 기여 및 결과

1. ICL과 회상의 구별된 학습 역학

훈련 역학은 서로 다른 능력이 언제 발현되는지에 대한 명확한 분리를 보여준다:

• 첫 번째 시스템 학습: 컨텍스트에서 처음 보이는 시스템의 동역학을 학습하는 능력은 훈련 전반에 걸쳐 점진적으로 발달하며, 최적의 유사 역행렬(pseudoinverse) 예측자를 향해 수렴한다.
• 재시작 ICL: (다른 시스템을 본 후) 새로운 시스템에 대한 ICL을 재시작하는 능력은 비교적 일찍 나타나지만, 완전한 연상 회상보다는 앞서 발현된다.
• 연상 회상: 심볼릭 레이블에만 의존하여 이전에 보았던 시스템을 회상하는 능력은 훈련 훨씬 후반부에 급격하게(상전이 현상과 함께) 나타난다. 특히 쿼리 이후의 첫 번째 토큰에 대해 그러하다.

2. 다중 메커니즘 가설 (C3)

중심적인 발견은 모델이 전체 작업을 위해 단일 메커니즘을 사용하는 것이 아니라는 점이다. 대신, 모델은 두 가지 뚜렷한 메커니즘을 사용하며, 이들은 서로 다른 시점에 발현되고 서로 다른 정보를 활용한다:

• 메커니즘 1 (레이블 기반): 쿼리 심볼 이후의 첫 번째 토큰을 예측하는 데 독점적으로 사용된다. 이 메커니즘은 올바른 시스템 상태를 검색하기 위해 이산적 심볼릭 레이블에 의존한다. 증거: 레이블이 잘못된 시스템으로 유도되었을 때, 첫 번째 토큰에 대한 오차가 즉각적으로 급증한다.
• 메커니즘 2 (관측 기반/베이지안): 두 번째 토큰 및 그 이후를 예측하는 데 사용된다. 이 메커니즘은 심볼릭 레이블을 거의 무시하고, 대신 관측된 첫 번째 토큰의 상태를 사용하여 기저 시스템 동역학을 추론한다("베이지안 방식"의 예측). 증거: 레이블이 잘못 유도되었을 때, 모델이 관측된 상태로부터 올바른 시스템을 추론하기 때문에 후속 토큰에 대한 오차는 낮게 유지된다.

3. 회로 디텐탱글먼트(Disentanglement)

에지 프루닝은 이 두 메커니즘이 모델의 가중치 내에서 물리적으로 별개임을 확인해 준다. "1-after" 토큰과 "2-after" 토큰을 예측하는 데 책임이 있는 회로는 0%의 에지 중첩을 보이며, 이는 모델이 태스크 에피소드를 시작하는 것과 이를 지속하는 것에 대해 별도의 서브 네트워크를 학습했음을 나타낸다.

4. 자연어로의 일반화

저자들은 OLMo 체크포인트를 사용하여 영어-스페인어 번역이라는 자연어 작업에서 이 현상을 검증한다. 저자들은 유사한 상전이를 관찰했다: 태스크 정의(심볼릭 레이블, 예: "X:")를 기반으로 번역을 시작하는 능력(인컨텍스트 연상 회상)은 번역을 일단 시작한 후 지속하는 능력보다 더 늦고 급격하게 발현된다. 의미론적 레이블(예: "Spanish:")이 심볼릭 레이블 대신 사용될 경우 이 간극이 사라지는데, 이는 지연 현상이 태스크 정의의 인컨텍스트 학습에 특이적임을 확인시켜 준다.

의의 및 주장

본 논문은 트랜스포머가 복잡한 인컨텍스트 작업을 수행하는 방식에 대한 기계론적 설명을 제공한다고 주장하며, ICL을 단일한 능력으로 보는 관점에 도전한다.

• ICL의 분해: 저자들은 "인컨텍스트 학습"이 단일한 기술이 아니라, 순차적으로 발현되는 구별된 능력(태스크 인식 vs. 태스크 지속)의 복합체라고 주장한다.
• 국소적 불충분성(Local Underspecification): 본 논문은 자기회귀적 훈련이 토큰들에 대한 평균 손실을 최소화하기 때문에, 모델이 대다수의 토큰(2번째, 3번째 등)에 대해 "충분히 좋은" 메커니즘(관측 기반 추론)을 사용하여 성공할 수 있으며, 이것이 더 복적인 메커니즘(레이블 기반 검색)을 학습하기 위한 그래디언트 압력을 감소시킨다고 제안한다.
• 상전이: 결과는 학습 역학에서 뚜렷한 상전이가 존재함을 뒷받침한다. 즉, 한 가지 능력(지속)의 발현이 다른 능력(시작)의 즉각적인 발현을 보장하지 않으며, 두 능력에 필요한 정보가 이론적으로 유사하더라도 그러하다.
• 기계론적 해석 가능성: 이 연구는 에지 프루닝이 어떻게 서로 다른 기능을 수행하는 별개의 회로를 격리할 수 있는지 보여주며, 이는 네트워크의 서로 다른 부분이 어떻게 서로 다른 작업을 처리하는지 이해하는 도구를 제공한다.

저자들은 이러한 발견이 I클(ICL)에 대한 이해에 새로운 차원의 미묘함을 제공한다고 결론짓는다. 즉, 멀티 토큰 작업의 경우, 모델은 시퀀스의 위치에 따라 동시에 서로 다른 모드(레이블 기반 vs. 관측 기반)로 작동할 수 있음을 시사한다.