기술적 요약: 차기 잠재 예측 트랜스포머는 컴팩트한 세계 모델을 학습한다

문제 정의

표준 트랜스포머 아키텍처는 순환 구조를 자기 주의 (self-attention) 와 시퀀스 길이에 비례하여 확장되는 메모리 메커니즘으로 대체합니다. 이는 뛰어난 병렬 처리 능력과 예측 능력을 가능하게 하지만, 일관된 전이 규칙을 가진 컴팩트한 잠재 상태로 역사를 압축하려는 선제적 귀납적 편향 (inductive bias) 을 제거합니다. 결과적으로 트랜스포머는 종종 훈련 데이터에는 잘 적합하지만 일반화 성능이 낮은 복잡하고 작업별 단축키 (shortcuts) 를 학습하여, 장기 계획 및 추론에 필요한 일관된 "세계 모델"이나 신념 상태 (belief states) 를 형성하지 못합니다. 표준 차기 토큰 예측 목적 함수는 모델의 내부 은닉 상태가 과거 역사의 충분 통계량 (sufficient statistics) 으로 진화함을 보장하지 않으며, 이러한 상태가 미래 관측치를 재귀적으로 예측할 수 있음을 보장하지도 않습니다.

방법론: 차기 잠재 예측 (NextLat)

저자들은 차기 토큰 훈련에 잠재 공간에서의 자기지도식 예측을 보강하는 차기 잠재 예측 (Next-Latent Prediction, NextLat) 프레임워크를 소개합니다. 핵심 아이디어는 다음 토큰이 주어졌을 때 다음 잠재 상태를 예측할 수 있도록 트랜스포머가 잠재 표현 (특히 최종 레이어 은닉 상태 $h_t$) 을 학습하도록 훈련하는 것입니다.

이론적 기반

이 논문은 모델이 차기 토큰 일관성($h_t$에서 $X_{t+1}$ 예측) 과 전이 일관성($h_t$와 $X_{t+1}$에서 $h_{t+1}$ 예측) 을 모두 최적화할 경우, 은닉 상태 $h_t$가 이론적으로 **신념 상태 (belief states)**로 수렴함이 보장됨을 입증합니다. 신념 상태는 미래 $X_{t+1:T}$를 예측하기 위한 역사 $X_{1:t}$의 충분 통계량으로 정의됩니다. 과거 토큰의 정보를 압축하지 않고 임의의 조회 (ad-hoc lookups) 를 허용하는 표준 자기회귀 훈련과 달리, NextLat 은 잠재 상태가 다음 상태로 전이하고 다음 토큰을 예측하는 데 필요한 모든 정보를 포함해야 하는 순환적 귀납적 편향을 강제합니다.

구현

NextLat 은 표준 트랜스포머 훈련 루프에 경량 보조 목적 함수를 추가하여 작동합니다:

1. 잠재 역학 모델 ($p_\psi$): 현재 은닉 상태 $h_t$와 다음 토큰 $X_{t+1}$을 입력받아 다음 은닉 상태 $\hat{h}_{t+1}$를 예측하도록 훈련된 간단한 다층 퍼셉트론 (MLP) 입니다.
2. 손실 함수: 총 손실은 다음 요소들을 결합합니다:
   • 차기 토큰 손실 ($L_{next-token}$): 토큰 예측을 위한 표준 교차 엔트로피.
   • 차기 은닉 상태 손실 ($L_{next-h}$): 표현 붕괴를 방지하기 위해 목표값에 stop-gradient 를 적용하여 예측된 상태 $\hat{h}_{t+1}$를 실제 다음 상태 $h_{t+1}$에 대해 회귀하는 Smooth L1 손실.
   • KL 발산 손실 ($L_{KL}$): 예측된 잠재 상태에서 예측된 토큰 분포가 실제 잠재 상태에서의 분포와 일치하도록 보장하는 보조 목적 함수로, 부드러운 감독 신호 역할을 합니다.
3. 다단계 감독: 모델은 잠재 역학을 $d$단계에 걸쳐 재귀적으로 풀어서 훈련할 수 있으며, 이는 장기 의존성에 대한 더 풍부한 기울기 신호를 제공합니다.

중요하게도 NextLat 은 트랜스포머의 아키텍처, 추론 절차, 그리고 병렬 훈련 효율성을 변경하지 않습니다. 잠재 역학 모델은 표현을 형성하기 위해 훈련 중에만 사용되며, 추론 중에는 모델이 평소와 같이 자기회귀적으로 디코딩합니다.

주요 기여

1. 이론적 보장: 이 논문은 차기 토큰과 차기 잠재 일관성을 최적화하는 것이 트랜스포머 표현을 신념 상태로 확실히 형성함을 증명하며, 이는 차기 토큰 예측만으로는 보장되지 않는 특성입니다.
2. 실용적 프레임워크: NextLat 은 트랜스포머의 병렬 훈련 또는 추론 구조를 변경하지 않고 순환적 귀납적 편향을 주입하는 계산 효율적인 보조 목적 함수를 도입합니다.
3. 가변 길이 자기 추측 디코딩: 일관된 잠재 역학을 학습함으로써 NextLat 은 잠재 공간에서 재귀적 예측을 가능하게 합니다. 이는 고정된 훈련 지평을 넘어 토큰을 초안 (draft) 으로 작성할 수 있는 가변 길이 자기 추측 디코딩을 허용하여 추론을 크게 가속화합니다.

실험 결과

저자들은 세계 모델링, 추론, 계획, 언어 모델링이라는 네 가지 도메인에서 NextLat 을 평가했습니다.

• 세계 모델링 (맨해튼 택시 탑승): NextLat 은 표준 GPT, MTP, JTP 베이스라인에 비해 실제 세계 모델 (거리 토폴로지) 과 훨씬 더 일관된 내부 지도를 학습합니다. 시퀀스 압축률 (GPT 의 0.65 대비 0.71) 이 더 높고 유효 잠재 차원 (GPT 의 160.1 대비 52.7) 이 훨씬 낮아 더 컴팩트하고 효율적인 표현을 나타냅니다.
• 추론 (카운트다운): 수학 추론 작업에서 NextLat 은 베이스라인을 크게 능가합니다 ($d=1$에서 GPT 의 33.1% 대비 54.8% 정확도). 이는 모델을 목표와 일치시키기 위해 유효하지 않은 최종 방정식을 강제로 맞추는 "후회스러운 타협"을 줄이는 우수한 선제적 계획 능력을 보여줍니다.
• 계획 (Path-Star 그래프): NextLat 은 다양한 그래프 토폴로지에서 거의 100% 해결률을 유지하는 반면, BST 및 JTP 와 같은 베이스라인은 더 큰 그래프 ($G_{7,7}$) 에서 실패합니다. 이는 NextLat 이 토큰 공간 예측에서 흔히 발생하는 "Clever Hans" 단축키 학습을 피함을 시사합니다.
• 언어 모델링 (FineWeb-Edu):
  • 하류 정확도: NextLat 은 MTP 및 JTP 보다 차기 토큰 퍼플렉시티를 더 잘 유지하면서 하류 제로샷 정확도에서 modest 한 개선을 달성합니다 (GPT 의 58.82% 대비 평균 59.21%).
  • 추론 속도: 가변 길이 자기 추측 디코딩을 활용하여 NextLat 은 표준 자기회귀 샘플링에 비해 최대 3.3 배 빠른 추론을 달성합니다. 이 속도 향상은 훈련 지평을 넘어 초안을 작성할 때에도 지속되며, 이는 고정된 토큰 공간 지평을 가진 MTP 및 JTP 가 결여된 기능입니다.

중요성 및 주장

이 논문은 NextLat 을 트랜스포머에서 컴팩트하고 예측 가능하며 일반화 가능한 표현을 학습하기 위한 간단하면서도 효과적인 패러다임으로 위치시킵니다. 그 주요 중요성은 컴팩트한 세계 모델의 순환적 귀납적 편향을 트랜스포머의 확장성과 병렬성과 통합하는 데 있습니다.

저자들은 NextLat 이 차기 토큰 예측의 "단시적 (myopic)"인 성격을 해결하여 모델이 얕은 토큰 상관관계에 의존하기보다 구조화된 예측 내부 상태를 형성하도록 장려한다고 주장합니다. 아키텍처 변경 없이 가변 길이 추측 디코딩을 가능하게 함으로써 NextLat 은 더 빠른 추론과 더 나은 장기 계획 추론을 위한 실용적인 경로를 제공합니다. 논문은 NextLat 이 복잡한 시스템에 대한 더 단순하고 컴팩트하며 따라서 더 일반화 가능한 표현을 자기회귀 모델에 부여하는 훈련 목적 함수를 향한 한 걸음이라고 결론지었습니다.

저자들은 향후 방향에 대해 겸손하게 언급하며, 복잡한 잠재 역학 아키텍처 (단순 MLP 만 사용) 를 탐구하지 않았으며, 목적 함수의 최적 설계 (예: 더 큰 규모에서 KL 손실의 필요성) 는 추가적인 체계적인 연구가 필요하다고 지적합니다. 또한 NextLat 이 표현을 개선하지만 고정 깊이 트랜스포머의 근본적인 회로 복잡성은 TC0 클래스에 의해 여전히 제한되지만, 공동 훈련된 잠재 역학 모델 (RNN) 은 이러한 한계를 넘어 일반화할 수 있는 능력을 가진 것으로 보인다고 인정합니다.