[논문 리뷰] Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion

Boyuan Chen, Diego Martí Monsó외 4명|arXiv (Cornell University)|2024. 07. 01.

Forensic and Genetic ResearchBiochemistry, Genetics and Molecular Biology인용 수 3

한 줄 요약

이 논문은 각 토큰에 독립적인 노이즈 수준를 부여함으로써 다음 토큰 예측과 전순서 확산을 통합하는 새로운 훈련 철학인 Diffusion Forcing를 소개한다. 이는 안정적인 가변 길이 생성과 장기 시계열 모델링에서 효과적인 가이드링을 가능하게 하며, 영상 생성, 계획, 시계열 예측 분야에서 최고 성능을 달성한다. 이는 인과적 모델링과 유연한 토큰 기반 복원 일정을 결합함으로써 달성된다.

ABSTRACT

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는 각 토큰이 고유한 노이즈 수준로 독립적으로 노이즈 처리된 시퀀스를 복원하도록 단일 인과적 다음 토큰 예측 모델을 훈련시킨다. 이는 토큰 기반 복원 일정을 가능하게 한다.
• 모델은 노이즈 처리된 시퀀스에서 임의의 토큰 부분집합을 재구성하도록 훈련되어, 각 토큰의 개별 노이즈 수준에 기반해 토큰을 '해독'하는 데 효과적으로 학습한다.
• 추론 시, 모델은 임의의 토큰 기반 노이즈 수준를 갖는 전체 시퀀스를 점진적으로 복원하며, 즉각적인 다음 토큰부터 가변 길이 생성을 가능하게 한다.
• 새로운 샘플링 기법인 몬테카를로 트리 가이드링(Monte Carlo Tree Guidance, MCTG)은 인과 아키텍처와 가변 수평선 능력을 활용하여 계획 과제에서 고보상 경로로의 샘플링을 유도한다.
• 공유된 인과 트랜스포머 아키텍처를 사용하여 미래 토큰 예측을 과거 토큰에 기반하게 하여 인과성을 유지하면서도 전체 시퀀스 복원을 가능하게 한다.
• 훈련 목표는 훈련 데이터에서 관측된 모든 부분 시퀀스의 가능도에 대한 변동 하한을 최적화하여 연합 분포와의 일致성을 보장한다.

Figure 2 : Method Overview. Diffusion Forcing trains causal sequence neural networks (such as an RNN or a masked transformer) to denoise flexible-length sequences where each frame of the sequence can have a different noise level. In contrast, next-token prediction models, common in language modeling

실험 결과

연구 질문

• RQ1 단일 모델이 다음 토큰 예측의 가변 길이 생성 능력과 전순서 확산의 장기 수평 가이드링을 동시에 통합할 수 있는가?
• RQ2 토큰 기반 노이즈 일정이 영상과 같은 연속 시퀀스의 훈련 수명을 초월한 안정적 생성을 가능하게 하는가?
• RQ3 다음 토큰 모델의 인과 아키텍처를 확산 프레임워크 내에서 활용하여 몬테카를로 트리 가이드링과 같은 새로운 가이드링 기법을 가능하게 할 수 있는가?
• RQ4 제안된 훈련 목표가 부분 시퀀스 가능도에 대한 하한 최적화를 통해 이론적으로 타당한가?
• RQ5 Diffusion Forcing는 영상, 계획, 시계열 예측 등 다양한 시퀀스 모델링 과제에서 어떻게 성능을 발휘하는가?

주요 결과

• Diffusion Forcing는 영상과 같은 연속 시퀀스의 안정적인 장기 수평 생성을 가능하게 하며, 훈련 수명을 수천 토큰 초월해 확장되며 자동회귀 기반 모델이 발산하는 영역에서도 성능을 유지한다.
• 영상 생성 및 시계열 예측 분야에서 최고 성능을 달성하며, 여러 개의 GluonTS 데이터셋에서 기준 모델 대비 CRPS 점수가 뚜렷이 낮아졌다.
• Diffusion Forcing를 사용한 몬테카를로 트리 가이드링(MCTG)은 의사결정 및 계획 과제에서 뚜렷한 성능 향상을 이끌었으며, 비인과적 전순서 확산 기반 모델을 능가했다.
• 모델은 새로운 경로로의 구성적 일반화를 가능하게 하며, 효율적인 트리 탐색과 온라인 피드백 제어를 지원하여 자동회귀 모델과 확산 모델의 장점을 결합한다.
• 실험 결과는 훈련 목표가 부분 시퀀스 가능도에 대한 변동 하한을 효과적으로 최대화함을 보여주며, 이론적 기반의 타당성을 검증한다.
• 시각적 모방 학습 및 기반 모델 기반 계획 과제에서, Diffusion Forcing는 개선된 샘플 효율성과 안정성을 바탕으로 고보상 경로 생성을 가능하게 했다.