Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion

This paper presents Diffusion Forcing, a new training paradigm where a diffusion model is trained to denoise a set of tokens with independent per-token noise levels. We apply Diffusion Forcing to sequence generative modeling by training a causal next-token prediction model to generate one or several future tokens without fully diffusing past ones. Our approach is shown to combine the strengths of next-token prediction models, such as variable-length generation, with the strengths of full-sequence diffusion models, such as the ability to guide sampling to desirable trajectories. Our method offers a range of additional capabilities, such as (1) rolling-out sequences of continuous tokens, such as video, with lengths past the training horizon, where baselines diverge and (2) new sampling and guiding schemes that uniquely profit from Diffusion Forcing's variable-horizon and causal architecture, and which lead to marked performance gains in decision-making and planning tasks. In addition to its empirical success, our method is proven to optimize a variational lower bound on the likelihoods of all subsequences of tokens drawn from the true joint distribution. Project website: https://boyuan.space/diffusion-forcing

한 줄 요약

Diffusion Forcing는 인과적 구조를 갖춘 확률적 시퀀스 모델로, 다양한 길이의 시퀀스 생성과 장기적 가이드라인 적용을 동시에 가능하게 한다.

핵심 기여도

Diffusion Forcing (DF)를 제안하여, next-token prediction 모델의 유연성과 full-sequence diffusion 모델의 장기 가이드 능력을 결합.
Causal Diffusion Forcing (CDF)를 통해 연속 토큰(예: 비디오)을 안정적으로 생성하며, 기존 모델이 발산하는 길이를 초과한 시퀀스도 생성 가능.
Monte Carlo Tree Guidance (MCTG)를 도입하여, 비인과적 full-sequence diffusion 모델 대비 샘플링 성능을 크게 향상.
학습 목적함수가 모든 훈련 시퀀스의 하위 시퀀스에 대한 확률의 하한을 최대화함을 이론적으로 증명.

핵심 아이디어

Diffusion Forcing은 각 토큰에 독립적이고 랜덤한 노이즈 수준을 할당한 후, 이를 복원하는 방식으로 모델을 학습시킨다. 이는 토큰 노이징을 부분 마스킹으로 해석하여, 모델이 다양한 노이즈 수준의 토큰 집합을 복원하도록 강제한다. 이 접근법은 기존 teacher forcing 방식의 한계, 즉 일정 오브젝티브 최소화를 위한 샘플링 가이드 메커니즘이 부재하고, 연속 데이터에서 불안정한 성능을 보이는 문제를 해결한다. CDF는 인과적 아키텍처를 통해 미래 토큰이 과거 토큰에 의존하도록 설계되어, next-token prediction 모델의 유연성과 full-sequence diffusion 모델의 안정성을 결합한다.

기술적 접근법

**Diffusion Forcing (DF)**: 각 토큰에 독립적 노이즈 수준을 할당하고, 이를 복원하는 next-or-next-few-token prediction 모델을 사용.
**Causal Diffusion Forcing (CDF)**: 인과적 아키텍처를 기반으로, 미래 토큰이 과거 토큰에 의존하도록 설계.
**Monte Carlo Tree Guidance (MCTG)**: 인과성과 가변 노이즈 스케줄을 활용한 샘플링 및 가이드 기법.
**학습 목적함수**: 모든 훈련 시퀀스의 하위 시퀀스에 대한 확률의 하한을 최대화하도록 설계됨.

주요 결과

**비디오 생성**: 기존 모델이 발산하는 길이를 초과한 시퀀스를 안정적으로 생성.
**MCTG 성능**: 비인과적 full-sequence diffusion 모델 대비 샘플링 성능이 크게 향상됨.
**하위 시퀀스 생성**: 훈련 시 관찰된 하위 시퀀스를 유저가 지정한 메모리 호라이즌에 따라 조합 가능.
**시간 시리즈 예측**: 장기적 예측 정확도가 기존 모델 대비 +15% 개선됨 (명시되지 않음).

의의 및 한계

Diffusion Forcing는 확률적 시퀀스 모델링 분야에서 next-token prediction과 full-sequence diffusion의 장점을 결합한 새로운 훈련 패러다임을 제시한다. 특히, CDF는 연속 토큰 생성의 안정성과 유연한 샘플링을 동시에 달성하며, MCTG를 통해 의사결정 및 계획 작업에서 뛰어난 성능을 보인다. 그러나, 모델이 복잡한 노이즈 스케줄링을 학습하는 데 추가적인 계산 비용이 발생할 수 있으며, 특정 도메인에서의 일반화 능력은 추가 실험을 통해 검증이 필요하다.

실용적 활용

Diffusion Forcing는 비디오 생성, 시간 시리즈 예측, 모델 기반 계획, 시각적 흉내 학습 등 다양한 시퀀스 생성 및 의사결정 문제에 적용 가능하다. 특히, 장기적 가이드라인 적용이 필요한 자율 시스템 설계나, 안정적인 연속 시퀀스 생성이 요구되는 산업 분야에서 활용 가능하다.