기술 요약: 양방향 유도(Induction)

문제 정의

자기회귀(Autoregressive, AR) 트랜스포머의 내부 메커니즘은 기계론적 해석 가능성(mechanistic interpretability)의 관점에서 광범위하게 연구되어 왔으나, 확산 언어 모델(Diffusion Language Models, DLMs)의 작동 방식은 여전히 미개척 영역으로 남아 있다. DLM은 순차적인 다음 토큰 예측이 아닌 반복적인 노이즈 제거(denoising)를 통해 텍스트를 생성하며, 양방향 어텐션과 다른 학습 목적 함수를 사용한다. 여기서 핵심적인 공백은 DLM이 AR 모델과 동일한 방식으로 유도(induction)—모델이 반복되는 컨텍스트를 식별하고 그 뒤에 오는 토큰을 복사하는 전형적인 인컨텍스트 학습 메커니즘—를 구현하는지 여부이다. 특히, AR 모델이 사용할 수 있는 엄격한 인과적(causal) 컨텍스트와 달리, 타겟 토큰이 마스크 토큰 [M] 뒤에 숨겨져 있고 양방향 컨텍스트에 둘러싸여 있을 때 DLM이 유도를 어떻게 처리하는지는 불분명하다.

방법론

저자들은 매칭된 어텐션 전용(attention-only) AR 트랜스포머와 흡수 마스크(absorbing-mask) DLM을 비교하는 통제된 기계론적 연구를 수행하였다.

• 모델 아키텍처: 연구에는 FineWeb-Edu 시퀀스로 학습된 소규모 모델(깊이 $L \in \{1, 2, 3\}$, 은닉 차원 $d_{model}=512$, 16개 어텐션 헤드)을 활용하였다. 두 모델 군은 어텐션 마스크(AR은 인과적, DLM은 양방향)와 학습 목적 함수(다음 토큰 예측 vs. 흡수 마스크 확산)만 다를 뿐, 동일한 아키텍처와 학습 레시피를 공유한다.
• 평가 태스크: 유도 능력은 첫 번째 절반이 두 번째 절반에서 반복되는 합성 시퀀스를 사용하여 측정되었다. 이 태스크는 두 번째 절반의 마스크된 토큰을 첫 번째 절반의 마스크되지 않은 소스에 기반하여 예측하는 작업(순방향 유도)과 그 반대(역방향 유도)를 포함한다.
• 분석 기법:
  • 인과적 절제(Causal Ablations): 개별 어텐션 헤드가 유도에 미치는 인과적 필요성을 결정하기 위해 평균 절제(mean ablations)를 수행하였다.
  • 가중치 분해(Weight Decomposition): 어텐션 로짓이 어떻게 형성되고 정보가 소스에서 타겟으로 어떻게 흐르는지 분석하기 위해 QK(Query-Key) 및 OV(Output-Value) 분해를 사용하였다.
  • 선형 프로빙(Linear Probing): 잔차 스트림(residual stream)에 선형 프로브를 학습시켜 글로벌 마스크 비율 정보의 존재를 탐지하였다.
  • 패칭 실험(Patch Experiments): 글로벌 마스크율의 표현에 인과적으로 개입하여 예측 엔트로피의 변화를 관찰하기 위해 "패칭(patching)" 기법을 사용하였다.
• 모델 변형: 기계론적 분석을 위해, 저자들은 학습된 LayerNorm 모델의 성능이 원래 모델과 유사함을 확인하면서, 잔차 스트림의 가산적 분해(additive decomposition)를 용이하게 하기 위해 훈련된 LayerNorm 모델을 선형 맵으로 변환한 "폴디드 노-레이어놈(folded no-LayerNorm)" DLM을 사용하였다.

주요 기여 및 결과

1. 양방향 유도 회로 (Bidirectional Induction Circuit)

본 연구는 DLM이 방향 대칭적인 양방향 유도 회로를 학습한다는 것을 밝혀냈다. 단일 인과적 경로에 의존하는 AR 모델과 달리, DLM은 두 개의 구별된 경로를 활용한다:

• 이전 토큰 및 다음 토큰 헤드: 첫 번째 레이어(Layer 0)에서 특정 헤드들은 마스크된 토큰의 즉각적인 왼쪽 및 오른쪽 이웃에 어텐션을 가하며, 이러한 로컬 컨텍스트 특징을 잔차 스트림에 기록한다.
• 유도 헤드: 후속 레이어(Layer 1+)에서 유도 헤드는 QK 가중치를 사용하여 마스크 위치에 기록된 레이어 0 헤드의 특징을 소스 답변 위치의 대응하는 특징과 일치시킨다.
• 메커니즘: QK 분해 결과, 유도 헤드는 마스크의 로컬 컨텍스트로부터 구축된 Q-composition(쿼리)과 소스의 로컬 컨텍스트로부터 구축된 K-composition(키)을 사용하여 일치하는 소스를 찾는다. 이후 OV 경로는 소스 토큰을 마스크된 위치로 복사한다.

2. 컨텍스트 접근성 vs. 메커니즘 강도

저자들은 DLM의 우수한 유도 성능이 근본적으로 더 강력한 유도 메커니즘 때문이 아니라, 양방향 컨텍스트에 대한 접근성 때문임을 입증하였다.

• 왼쪽 컨텍스트로 제한될 경우(AR 설정을 모사), DLM의 성능은 AR 모델 수준과 비슷하거나 그보다 낮아진다.
• DLM은 마스크된 토큰의 왼쪽과 오른쪽 이웃이 모두 보일 때만 AR 모델보다 뛰어난 성능을 보인다.
• 전체 양방향 설정에서의 유도 점수는 왼쪽 전용 및 오른쪽 전용 설정의 점수 합에 근접하며, 이는 DLM의 이점이 양방향으로부터 정보를 활용하는 데서 기인함을 확인시켜 준다.

3. 글로벌 마스크율을 통한 암시적 타임스텝 (Implicit Timestep via Global Mask Rate)

본 논문은 DLM이 명시적인 타임스텝 임베딩이 없음에도 불구하고, 글로벌 마스크 토큰의 비율을 암시적 타임스텝으로 계산하고 활용한다는 인과적 증거를 제시한다.

• 표현: 선형 프로브는 첫 번째 어텐션 레이어 이후 잔차 스트림에서 글로벌 마스크 비율을 높은 정확도($R^2 \approx 0.82-0.91$)로 복구할 수 있다.
• 인과적 유용성: 고부패(high-corruption) 상태에서 저부패(low-corruption) 상태로 "마스크-율 방향(mask-rate direction)"을 패칭하면 마스크 위치의 예측 엔트로피가 인과적으로 증가한다. 이는 모델이 글로벌 부패 수준을 사용하여 불확실성을 조절하며, 효과적으로 마스크율을 타임스텝으로 취급함을 확인시켜 준다.

4. 회로의 견고성 및 출현 (Circuit Robustness and Emergence)

• 출현: 유도는 훈련 과정 중(약 8~24k 스텝) 갑작스럽게 출현하며, 기능을 수행하기 위해 최소 두 개의 레이어가 필요하다. 이는 특징을 쓰는 레이어와 이를 추출하는 후속 레이어가 필요하다는 점과 일치한다.
• 로컬 윈도우: 유도 메커니즘은 약 5개의 토큰(마스크된 토큰과 양옆의 이웃 각각 2개씩)으로 구성된 로컬 윈도우에 크게 의존한다. 글로벌 부패율이 낮더라도 이 즉각적인 이웃들이 마스크되면 성능이 크게 떨어진다.

의의 및 주장

이 연구는 DLM이 단순히 AR 동작을 복제한다는 가정을 넘어, DLM의 내부 메커니즘을 이해하는 기초적인 단계가 된다. 저자들은 다음과 같이 주장한다:

1. 기계론적 차이: DLM은 쿼리와 키 특징이 마스크된 위치로부터 구축되는 방식에서 AR 모델의 인과적 유도 회로와 구조적으로 다른, 대칭적이고 양방향적인 회로를 통해 유도를 구현한다.
2. 컨텍스트적 이점: 인컨텍스트 학습에서 DLM과 AR 모델 간의 성능 격차는 유도 알고리즘 자체의 내재적 우월성이 아니라 양방향 컨텍스트의 가용성에서 비롯된다.
3. 암시적 상태 관리: DLM은 명시적인 임베딩 없이 시퀀스의 글로벌 상태(마스크 비율)로부터 시간적 정보(타임스텝)를 내부적으로 도출하며, 이 메커니즘은 예측 엔트로피와 인과적으로 연결되어 있다.

결론적으로, 본 논문은 이러한 발견이 비자기회귀 모델의 인컨텍스트 학습을 이해하는 토대를 제공한다고 밝히면서도, 소규모 어텐션 전용 모델 및 합성 태스크 사용에 따른 한계를 인정하며, 이러한 메커니즘이 더 큰 MLP 포함 DLM 및 자연어 태스크로 확장되는지 확인하기 위한 추가 연구가 필요함을 언급하였다.