확산 강제: 다음 토큰 예측과 전체 시퀀스 확산의 만남

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

July 1, 2024

저자: Boyuan Chen, Diego Marti Monso, Yilun Du, Max Simchowitz, Russ Tedrake, Vincent Sitzmann

cs.AI

초록

본 논문은 확산 모델이 토큰별 독립적인 노이즈 수준을 가진 토큰 집합을 노이즈 제거하도록 학습하는 새로운 훈련 패러다임인 Diffusion Forcing을 소개합니다. 우리는 Diffusion Forcing을 시퀀스 생성 모델링에 적용하여, 과거 토큰을 완전히 확산시키지 않고 하나 또는 여러 미래 토큰을 생성하도록 인과적 다음 토큰 예측 모델을 학습시킵니다. 우리의 접근 방식은 가변 길이 생성과 같은 다음 토큰 예측 모델의 강점과, 바람직한 경로로 샘플링을 유도할 수 있는 능력과 같은 전체 시퀀스 확산 모델의 강점을 결합한 것으로 나타났습니다. 우리의 방법은 (1) 훈련 범위를 초과하는 길이의 비디오와 같은 연속 토큰 시퀀스를 롤아웃할 수 있는 능력(기존 방법들은 이 경우 발산함)과 (2) Diffusion Forcing의 가변 범위 및 인과적 구조로부터 독특하게 이점을 얻는 새로운 샘플링 및 유도 기법을 제공하며, 이는 의사결정 및 계획 작업에서 뚜렷한 성능 향상을 이끌어냅니다. 실험적 성공 외에도, 우리의 방법은 실제 결합 분포에서 추출된 모든 부분 시퀀스 토큰의 가능도에 대한 변분 하한을 최적화함이 증명되었습니다. 프로젝트 웹사이트: https://boyuan.space/diffusion-forcing/

English

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/