[論文レビュー] Diffusion Forcing: Next-token Prediction Meets Full-Sequence Diffusion

Boyuan Chen, Diego Martí Monsó+4 more|arXiv (Cornell University)|Jul 1, 2024

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 は、各トークンが個別にノイズ処理されたシーケンスを、1つの因果的次トークン予測モデルがノイズ除去するように訓練する。これにより、個別トークンのノイズ除去スケジュールが可能になる。
• モデルは、ノイズがかけられたシーケンスから任意のトークンサブセットを再構築するように訓練され、個々のノイズレベルに基づいてトークンを「アンマスク」する能力を学習する。
• 推論時、任意の個別ノイズレベルを設定した完全なシーケンスに対して段階的なノイズ除去を実行し、直近のトークンから可変長の生成が可能になる。
• 新しいサンプリング方式であるモンテカルロ木ガイド（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 1つのモデルが、次トークン予測の可変長生成能力と完全シーケンス拡散の長時間スパンガイドライン機能を統合できるか？
• RQ2 個別トークンのノイズスケジューリングが、トレーニング範囲を超えて動画のような連続的シーケンスの安定的生成を可能にするか？
• RQ3 次トークンモデルの因果アーキテクチャを拡散フレームワーク内で活用することで、モンテカルロ木ガイド（MCTG）のような新しいガイドライン方式を実現できるか？
• RQ4 提案された訓練目的が、部分シーケンス尤度の下界を最適化することで理論的に正当化されるか？
• RQ5 Diffusion Forcing は、動画、計画、時系列予測など多様なシーケンスモデリングタスクでどのように性能を発揮するか？

主な発見

• Diffusion Forcing は、動画のような連続的シーケンスの安定的かつ長時間スパンの生成を可能にし、トレーニング範囲を数千トークンも超えて生成可能である。一方、自己回帰ベースラインは発散する。
• 本手法は、動画生成および時系列予測の分野で最先端の性能を達成し、GluonTS データセットの複数のセットで、CRPS スコアがベースラインモデルよりも顕著に低い。
• Diffusion Forcing を用いたモンテカルロ木ガイド（MCTG）は、意思決定および計画タスクで顕著な改善をもたらし、非因果的完全シーケンス拡散ベースラインを上回る。
• モデルは新しい軌道に組み合わせ的に一般化可能であり、効率的なツリー探索とオンラインフィードバック制御をサポートする。自己回帰モデルと拡散モデルの長所を統合する。
• 実験的結果から、訓練目的が部分シーケンス尤度の変分下界を効果的に最大化していることが示され、本手法の理論的基盤が裏付けられる。
• 視覚的模倣学習およびモデルベース計画の分野において、Diffusion Forcing は、より高い報酬の軌道を生成可能であり、サンプル効率と安定性が向上する。