[Paper] 일관된 월드 모델을 위한 멀티-토큰 예측 및 잠재 의미 강화

Source: arXiv - 2604.06155v1

개요

이 논문은 대형 언어 모델(LLMs)이 일관된 내부 세계 모델—즉, 그들이 추론하는 환경에 대한 구조화된 표현—을 구축할 수 있는지를 조사한다. 고전적인 다음 토큰 예측(NTP) 패러다임을 넘어 다중 토큰 예측(MTP) 으로 전환함으로써, 저자들은 이론적·실험적으로 MTP가 모델의 은닉 상태가 일관된 “신념 상태(belief states)”로 수렴하도록 장려한다는 것을 보여준다. 또한 주요 실패 모드인 구조적 환각(structural hallucinations) 을 식별하고, 잠재 역학을 현실에 기반하게 유지하기 위해 Latent Semantic Enhancement MTP (LSE‑MTP) 를 도입한다.

주요 기여

• Theoretical analysis of MTP’s gradient bias: MTP가 여러 단계에 걸쳐 그래디언트를 결합하여 수축 효과를 유도하고, 숨겨진 표현을 안정적인 믿음 상태로 밀어낸다는 것을 보여준다.
• Identification of structural hallucinations: 표준 MTP가 잠재 공간에서 지름길을 이용해 알려진 환경 제약을 위반하는 예측을 생성할 수 있음을 보여준다.
• LSE‑MTP algorithm: 예측된 잠재 궤적을 실제 숨겨진 상태 시퀀스와 정렬하는 새로운 학습 목표로, 모델을 실제 세계 역학에 효과적으로 고정한다.
• Empirical validation on synthetic graphs and a real‑world Manhattan taxi‑ride dataset: 표현 정렬, 환각 감소, 입력 교란에 대한 견고성 향상을 정량화한다.
• Practical recipe for developers: 최소한의 코드 변경으로 채택할 수 있는 기존 트랜스포머 학습 파이프라인에 대한 플러그‑앤‑플레이 수정 방안을 제공한다.

방법론

1. Multi‑Token Prediction (MTP)

   • 모델을 단일 다음 토큰에 대해 감독하는 대신, MTP는 모델이 한 번의 순전파에서 (k)개의 미래 토큰 시퀀스를 예측하도록 요구합니다.
   • 손실은 (k)개의 위치에 대해 합산되며, 이는 gradient coupling을 생성합니다: 토큰 (t)에 대한 그래디언트는 토큰 (t+1,\dots,t+k)에 대한 오류에 의존합니다.

2. Gradient Inductive Bias Analysis

   • 저자들은 손실에 대한 은닉 상태의 Jacobian을 유도하고, 완만한 매끄러움 가정 하에 결합된 그래디언트가 은닉 표현 공간에서 contractive map처럼 작용한다는 것을 증명합니다.
   • 직관적으로, 이는 유사한 미래 토큰 시퀀스로 이어지는 은닉 상태들을 서로 가깝게 클러스터링하도록 강제하여 일관된 믿음 상태를 형성합니다.

3. Diagnosing Structural Hallucinations

   • 합성 그래프 탐색 과제에서 은닉 궤적을 시각화함으로써, 모델이 때때로 불법적인 지름길(예: 연결되지 않은 노드 사이를 점프) 을 학습하는 것을 관찰했습니다. 이러한 행동은 토큰 손실이 이산적이고 기본 제약을 무시하기 때문에 여전히 최소화됩니다.

4. Latent Semantic Enhancement MTP (LSE‑MTP)

   • 보조 잠재 손실: 토큰 손실과 함께, 모델은 정답 은닉 상태 궤적 (교사 모델이나 알려진 시뮬레이터에서 얻은)에서 벗어나는 것에 대해서도 페널티를 부과받습니다.
   • Semantic anchoring: 잠재 손실은 적응적으로 가중치가 부여됩니다; 초기 학습 단계에서는 언어 유창성을 위해 토큰 손실에 중점을 두고, 후반 단계에서는 정렬을 강화하기 위해 잠재 항을 증가시킵니다.
   • Implementation: 선택된 은닉 레이어에 간단한 평균 제곱 오차(mean‑squared error) 항을 추가합니다; 구조적 변화는 필요하지 않습니다.

5. Evaluation Protocol

   • Synthetic graph task: 에이전트가 방향성 그래프를 탐색합니다; 성공은 그래프 간선을 준수하는 궤적 비율로 측정됩니다.
   • Manhattan Taxi Ride: 원시 호출 데이터에서 위치 토큰(위도/경도 버킷) 시퀀스를 예측합니다; 평가지표에는 토큰 정확도, 궤적 충실도(실제 경로와의 거리), 그리고 잡음이 있는 시작점에 대한 견고성이 포함됩니다.

결과 및 발견

• 수축성 확인: MTP 하에서 숨겨진 표현이 실제 상태 주변에 더 촘촘한 클러스터를 형성하며, LSE‑MTP는 이를 더욱 강화합니다.
• 환각 감소: 잠재 의미 손실을 추가하면 구조적 위반이 크게 감소합니다.
• 견고성: LSE‑MTP로 학습된 모델은 입력이 교란될 때 훨씬 적게 성능이 저하되어, 보다 안정적인 내부 세계 모델을 나타냅니다.

Practical Implications

• More reliable planning and reasoning: For LLM‑based agents (e.g., autonomous assistants, code‑generation bots) that need to maintain a mental map of the world, LSE‑MTP can reduce “impossible” suggestions that violate known constraints.
  보다 신뢰할 수 있는 계획 및 추론: 세계에 대한 정신적 지도를 유지해야 하는 LLM 기반 에이전트(예: 자율 어시스턴트, 코드 생성 봇)에게 LSE‑MTP는 알려진 제약을 위반하는 “불가능한” 제안을 줄일 수 있습니다.

• Improved downstream fine‑tuning: Adding the latent loss is a lightweight regularizer that can be applied when fine‑tuning on domain‑specific sequential data (e.g., navigation logs, transaction streams).
  향상된 다운스트림 파인튜닝: 잠재 손실을 추가하는 것은 가벼운 정규화 기법으로, 도메인 특화 순차 데이터(예: 내비게이션 로그, 거래 스트림)에서 파인튜닝할 때 적용할 수 있습니다.

• Safety & interpretability: By aligning hidden states with interpretable trajectories, developers gain a diagnostic handle to inspect whether the model’s internal belief matches reality.
  안전성 및 해석 가능성: 은닉 상태를 해석 가능한 궤적과 정렬함으로써, 개발자는 모델의 내부 신념이 현실과 일치하는지 검사할 수 있는 진단 도구를 얻게 됩니다.

• Compatibility with existing pipelines: LSE‑MTP only requires extracting hidden activations and computing an auxiliary MSE loss; it works with standard transformer libraries (Hugging Face, DeepSpeed, etc.).
  기존 파이프라인과의 호환성: LSE‑MTP는 은닉 활성화를 추출하고 보조 MSE 손실을 계산하기만 하면 되며, Hugging Face, DeepSpeed 등 표준 트랜스포머 라이브러리와 함께 작동합니다.

• Potential for multimodal world modeling: The same principle can be extended to vision‑language models where latent trajectories correspond to object motion or scene dynamics.
  다중모달 세계 모델링 가능성: 동일한 원리를 시각‑언어 모델에 확장할 수 있으며, 여기서 잠재 궤적은 객체 움직임이나 장면 역학에 대응합니다.

제한 사항 및 향후 작업

• Ground‑truth latent trajectories에 대한 의존성: LSE‑MTP는 숨겨진 상태 “gold” 시퀀스를 제공할 수 있는 교사 또는 시뮬레이터에 접근할 수 있다고 가정하는데, 이는 모든 도메인에서 가능하지 않을 수 있다.
• 매우 긴 horizon에 대한 확장성: 잠재 손실은 고정된 윈도우에서 계산된다; 수천 단계로 확장하면 메모리 오버헤드가 증가할 수 있다.
• 합성 그래프와 단일 실제 데이터셋에 한정된 평가: 일반성을 확인하기 위해 로보틱스, 대화 등 보다 다양한 벤치마크가 필요하다.
• 저자들이 제시한 향후 방향:
  1. 자기 지도(self‑supervision)를 통해 잠재 궤적을 암묵적으로 학습한다.
  2. 추론 시 제약 인식 디코딩을 통합한다.
  3. 서로 다른 토큰 그룹이 세계 모델의 서로 다른 추상화 수준에 해당하는 계층적 MTP를 탐구한다.

저자

• Qimin Zhong
• Hao Liao
• Haiming Qin
• Mingyang Zhou
• Rui Mao
• Wei Chen
• Naipeng Chao

논문 정보

• arXiv ID: 2604.06155v1
• Categories: cs.LG, cs.AI, cs.CL
• Published: 2026년 4월 7일
• PDF: Download PDF