기술 요약: 세계 모델을 위한 식별 가능한 토큰 대응

1. 문제 제기

트랜스포머 기반 세계 모델은 과거 상태와 행동의 시퀀스를 토큰 스트림으로 처리하여 미래 상태를 예측함으로써 시각적 강화 학습 (RL) 에서 강력한 성능을 입증해 왔습니다. 그러나 이러한 모델은 장기적 롤아웃 (rollouts) 중에 시간적 불일치로 인해 자주 실패합니다. 일반적인 실패 양상은 객체의 중복, 소멸, 그리고 다른 객체로의 변형 등을 포함합니다.

저자들은 이러한 오류의 주된 원인을 다음과 같이 규명했습니다: 기존 접근법들은 다음 프레임 예측을 순수한 토큰 생성 문제로만 취급합니다. 그들은 시간 간격을 거친 토큰 간의 대응 관계를 명시적으로 모델링하지 못합니다. 현실적인 시각 환경에서 연속된 프레임의 토큰들은 종종 유지되고 이동하는 동일한 근본적인 실체 (entities) 에 대응합니다. 식별성을 추적하지 않은 채 모든 토큰을 처음부터 다시 생성하도록 강요받을 때, 모델은 지속적인 구조를 유지하는 데 어려움을 겪으며, 이는 시간이 지남에 따라 취약한 정체성 유지와 오류의 누적 (compounding errors) 로 이어집니다.

2. 방법론: 식별 가능한 토큰 대응 (ITC)

본 논문은 잠재 토큰 대응 변수를 가진 구조화된 확률적 추론 문제로 다음 프레임 예측을 공식화하는 트랜스포머 세계 모델인 **식별 가능한 토큰 대응 (ITC)**을 제안합니다.

핵심 메커니즘

ITC 는 모든 다음 프레임 토큰을 처음부터 새로 생성하는 대신, 각 다음 프레임 토큰에 대해 다음 중 하나를 결정합니다:

1. 이전 프레임의 토큰에서 복사 (정체성 유지).
2. 새로운 토큰으로 생성 (변화 또는 새로운 콘텐츠 처리).

이 할당은 최적 수송 (Optimal Transport, OT) 문제로 모델링됩니다.

기술 아키텍처

1. 토큰화 및 임베딩:

   • 시각적 관측치는 nearest-neighbor 패치 룩업 토크나이저를 사용하여 이산 토큰으로 변환됩니다.
   • 토큰 임베딩은 **3D 회전 위치 임베딩 (RoPE)**으로 보강됩니다. 표준 1D RoPE 와 달리, 이 공식은 시간, 수직, 수평 좌표를 독립적으로 인코딩하여 모델이 공간과 시간 양쪽에 걸쳐 국소화된 관계 구조를 포착할 수 있게 합니다.

2. 최적 수송 솔버:

   • 모델은 소스 정점이 이전 상태 토큰 ($u$) 과 트랜스포머가 예측한 후보 다음 상태 토큰 ($\tilde{p}$) 을 나타내고, 목적지 정점이 최종 다음 상태 토큰 ($\hat{u}$) 을 나타내는 그래프를 구성합니다.
   • **친화도 행렬 (affinity matrix)**이 계산되어 이전 토큰과 예측된 토큰 간의 관계를 점수화합니다. 이 행렬은 다음을 포함합니다:
     • $A^{(prev)}$: 이전 토큰과 예측 간의 친화도를 점수화하며, 물리적 타당성을 강제하기 위해 공간적 거리로 페널티를 부과합니다 (예: 객체가 임의로 매우 멀리 순간 이동할 수 없음).
     • $A^{(gen)}$: 이전 토큰을 재사용하는 대신 새로 생성된 토큰 (와일드카드) 을 허용하는 "보너스"를 점수화하며, 일정한 페널티를 따릅니다.
   • 솔버 (Sinkhorn 알고리즘 사용) 는 소스에서 목적지로 질량을 이동시키는 비용을 최소화하는 수송 계획을 계산합니다.

3. 이진화 및 디코딩:

   • 연속적인 수송 계획은 이진 할당 행렬로 변환됩니다 (그리디 이진화 절차, 알고리즘 2). 이는 엄격한 1 대 1 매핑을 보장합니다.
   • 디코딩 규칙: 다음 프레임의 각 위치에서, 수송 계획이 이전 토큰을 할당하면 해당 토큰이 복사됩니다. 계획이 "생성" 토큰을 할당하면, 모델은 트랜스포머의 예측 분포에서 샘플링합니다.

학습 루프

세계 모델과 정책은 동시에 학습됩니다. 세계 모델은 "상상 (imaginations)" (합성 궤적) 을 생성하며, 이는 실제 환경 데이터와 함께 Proximal Policy Optimization (PPO) 을 통해 정책을 업데이트하는 데 사용됩니다.

3. 주요 기여

• 토큰 대응의 공식화: 본 논문은 토큰 대응을 최적 수송 문제로 명시적으로 모델링함으로써 다음 프레임 예측을 재구성하여, "지속적인 실체 재사용"과 "새로운 콘텐츠 생성" 작업을 분리합니다.
• 최적 수송 통합: 이는 (Sinkhorn 을 통해) 미분 가능하면서도 (이진화를 통해) 이산적인 메커니즘을 도입하여 이전 토큰 복사와 새로운 토큰 생성 사이를 선택함으로써 환각 (hallucinations) 을 효과적으로 줄입니다.
• 3D 위치 인코딩: 3D RoPE 의 구현은 모델이 시간적 차이와 공간적 차이를 구별할 수 있게 하여, 비디오 시퀀스에서 움직이는 객체를 추적하는 데 필수적인 구성 요소가 됩니다.

4. 실험 결과

저자들은 ITC 를 네 가지 까다로운 벤치마크인 Craftax-classic, Craftax, MinAtar, Atari 100K에서 평가했습니다.

• Craftax-classic: ITC 는 72.5% 의 리턴 (return) 과 35.6% 의 점수를 달성하여, 67.4% 리턴과 27.9% 점수를 기록한 이전 최첨단 (Dedieu et al., 2025) 을 크게 능가했습니다. 또한 더 빠른 수렴을 보여주었습니다.
• Craftax: 더 복잡한 환경인 이 환경에서 ITC 는 7.09% 의 리턴과 2.40% 의 점수를 달성하여 Simulus 및 Dedieu et al. 을 포함한 베이스라인을 능가했습니다.
• MinAtar: ITC 는 네 가지 게임 (Asterix, Breakout, Freeway, SpaceInvaders) 에서 모두 모델 기반 RL 의 이전 최첨단 성능을 능가했습니다.
• Atari 100K: 26 개 게임 세트에 적용된 ITC 는 1.092 의 Interquartile Mean (IQM) 과 0.376 의 Optimality Gap 에서 새로운 최첨단 성능을 달성하여, 토큰 기반 세계 모델인 Simulus 를 능가했습니다.
• 예측 정확도: 10,000 개의 전이 (transitions) 데이터셋에서 ITC 는 베이스라인 대비 전체 예측 정확도를 3.39% 향상시켰으며, 움직이는 생물이 포함된 전이의 경우 3.45% 의 주목할 만한 개선을 보였습니다.
• 정성적 개선: 시각적 분석은 ITC 가 DreamerV3 및 Dedieu et al. 과 같은 베이스라인 모델에서 관찰되는 객체 중복 및 소멸과 같은 일반적인 아티팩트를 제거함을 확인했습니다.
• 효율성: 이 방법은 미미한 오버헤드를 도입하여 총 엔드 - 투 - 엔드 학습 시간을 단 2.8% 만 증가시켰습니다.

5. 중요성 및 주장

본 논문은 시간 간격을 거친 토큰 간의 대응 관계를 명시적으로 모델링함으로써, 장기적 기간 동안 객체 정체성을 유지하지 못하는 세계 모델의 광범위한 문제를 해결한다고 주장합니다. 저자들은 이 접근법이 모델이 "이전 프레임과 강한 대응 관계를 가진 토큰을 선택적으로 재사용"할 수 있게 하여, 모든 토큰을 트랜스포머가 다시 생성하는 것에만 의존하는 대신 프레임 간 유사성을 활용한다고 명시합니다.

중요성은 격자 기반 오픈 월드부터 표준 Atari 게임에 이르기까지 다양한 환경에서 계산 효율성을 유지하면서 최첨단 성능을 달성하는 데 있습니다. 저자들은 ITC 가 중복 및 소멸과 같은 시간적 불일치를 효과적으로 방지하여 더 높은 품질의 상상과 향상된 정책 성능으로 이어진다고 결론지었습니다.