Taming Diffusion Transformer for Real-Time Mobile Video Generation

著者 Yushu Wu, Yanyu Li, Anil Kag, Ivan Skorokhodov, Willi Menapace, Ke Ma, Arpit Sahni, Ju Hu, Aliaksandr Siarohin, Dhritiman Sagar, Yanzhi Wang, Sergey Tulyakov
所属 Snap Inc., Northeastern University
投稿日 2025年07月18日
カテゴリ cs.CV, cs.LG

Taming Diffusion Transformer for Real-Time Mobile Video Generation

基本情報

  • arXiv ID: 2507.13343v1 (https://arxiv.org/abs/2507.13343v1)
  • 著者: Yushu Wu, Yanyu Li, Anil Kag, Ivan Skorokhodov, Willi Menapace, Ke Ma, Arpit Sahni, Ju Hu, Aliaksandr Siarohin, Dhritiman Sagar, Yanzhi Wang, Sergey Tulyakov
  • 所属: Snap Inc., Northeastern University
  • 投稿日: 2025年07月18日
  • カテゴリ: cs.CV, cs.LG

簡単に説明すると

この論文は、モバイルデバイス上でリアルタイムに高品質な動画を生成できる効率的なDiffusion Transformer(DiT)モデルを提案しています。高圧縮VAE、感度を考慮した3段階プルーニング、敵対的ステップ蒸留という3つの最適化技術を組み合わせることで、iPhone 16 Pro Max上で10FPS以上の動画生成を実現しています。

プロジェクトページでは、デモ動画や詳細な情報が公開されています:
https://snap-research.github.io/mobile_video_dit/

1. 研究概要

1.1 背景と動機

Diffusion Transformerは、時間的に一貫性があり視覚的に魅力的な動画コンテンツを生成する能力で、動画生成タスクにおいて優れた性能を示してきました。しかし、これらのモデルは計算コストが非常に高く、特に高解像度動画生成において顕著です。3Dフルアテンションの計算量とメモリ消費は、トークン総数(t×H×W)に対して二次的にスケールします。

この制約により、リアルタイムやインタラクティブな設定、特に処理能力とエネルギー予算が限られたモバイルデバイスでの展開が困難となっています。既存の拡散ベースモデルの最適化努力は主にUNetベースのデノイザーに焦点を当てており、DiTの効率性を調査する研究は非常に少なく、多くの場合、知覚品質や時間的一貫性の損失に苦しんでいます。

1.2 主要な貢献

本研究の主要な貢献は次の3点です。

第一に、オンデバイス動画生成のための潜在圧縮比、生成品質、速度のトレードオフを体系的に調査した最初の研究です。Diffusion Transformerでは8×32×32 VAEが生成速度と品質の良いトレードオフを達成することを発見しました。
第二に、効率的なDiTバックボーンを得るための新しい手法を提案しました。小さなネットワークをゼロから訓練すると劣った生成結果をもたらすことを発見しました。代わりに大規模な事前訓練されたスーパーネットワークから開始して、蒸留ガイド付きの感度を考慮したプルーニングを実行します。
第三に、敵対的ステップ蒸留のために、DiT用に特別に設計された新しいディスクリミネーター設計を提案しました。従来の手法を50%以上上回る性能を達成し、CFGなしで4ステップ推論を実現しました。

2. 提案手法

2.1 手法の概要

本研究では、動画Diffusion Transformerをモバイル展開向けに加速するための包括的な最適化パイプラインを提案しています。このアプローチは3つの主要な戦略を組み合わせています。

第一に、高圧縮動画VAEにより潜在表現のトークン数を80%削減します。第二に、効率的なモバイルDiTを実現するため、感度を考慮した3段階プルーニングとKDガイド付きフレームワークを使用します。第三に、敵対的ステップ蒸留により推論ステップ数を4ステップまで削減します。

2.2 技術的詳細

高圧縮動画VAEについて、VAEの圧縮率を調査しました。より高い圧縮率を持つ動画VAEが潜在表現のトークン数を大幅に削減し、DiT推論を高速化することを示しました。4×16×16から8×64×64までの様々な圧縮率の動画VAEを構築し、生成速度と品質を比較しました。その結果、8×32×32の構成が速度と品質のバランスを良好に保つことを発見しました。

KDガイド付き3段階プルーニングについて、3つのレベルでプルーニングを実行します。ブロックプルーニング、アテンションヘッドプルーニング、線形層のチャネルプルーニングを含みます。学習可能なバイナリマスクを使用して実装し、各粒度の重要性をエンコードします。28個のDiTブロック、32個のアテンションヘッド、8192のFFN次元を持つ2Bパラメータのベースモデルから開始します。8ブロックをプルーニング、12個のアテンションヘッドを削減、FFN次元を25%削減して、915Mパラメータの効率的なモデルを実現しました。

敵対的ステップ蒸留について、DiT用に特別に設計された新しいディスクリミネーターアーキテクチャを提案しました。ジェネレーターから最初のKブロックを継承し、時間条件付き特徴パーサーとして凍結します。学習可能なディスクリミネーター層として完全な3Dアテンションとクロスアテンションを追加します。この設計により、CFGなしで4ステップ生成が可能になります。典型的な40ステップのCFG付きレシピと比較して20倍高速になります。

2.3 新規性

既存の効率的な動画生成手法は主にUNetベースのアーキテクチャに焦点を当てていましたが、本研究はDiTベースのモデルを初めてモバイルデバイス上でリアルタイム動画生成を可能にしました。

また、VAEの圧縮率と生成品質のトレードオフを体系的に調査し、DiTに特化した敵対的蒸留手法を設計した点も新規性があります。特に、ディスクリミネーターの設計において、ジェネレーターの初期ブロックを再利用し、時間条件付けを保持しながら効率的な学習を可能にした点は革新的です。

3. 実験結果

3.1 実験設定

すべての訓練は内部で収集した画像と動画データ(実データと合成データの両方を含む)で行われました。DiT訓練には128個のNVIDIA A100 80GB GPUを使用し、AdamWオプティマイザーで学習率5e-5を使用しました。

評価はVBenchの標準設定に従い、各プロンプトに対して5つの動画を生成し、1Kプロンプトセットでスコアをテストしました。敵対的蒸留モデルを使用して、576×1024の解像度で121フレームの水平動画を4デノイジングステップで生成しました(CFGなし)。

3.2 主要な結果

提案手法は、モバイル向けにコンパクトで高速推論用に設計されているにもかかわらず、VBenchで既存の動画生成モデルと比較して高いスコアを達成しました。

VBench総合スコアは81.45(モバイルモデル)、83.09(サーバーモデル)を記録しました。生成速度はiPhone 16 Pro Maxで12.2 FPS(49フレームを4秒で生成)を達成しました。モデルサイズは0.9B(モバイル)、2.0B(サーバー)に削減しました。

最近のDiTベースのOpenSora-V1.2、CogVideoX-2B、UNetベースのVideoCrafter-2.0を含む既存手法よりも高い総合スコアを達成しました。

3.3 既存手法との比較

4ステップ蒸留されたT2V-TurboやAnimateLCMと比較して、50%以上のサイズ削減でより良い性能を達成しました。

VAE圧縮率のスケーリング実験では次の結果を得ました。
4×16×16 VAEは最高のPSNR(33.1)を達成しましたが、速度が非常に遅く(7900ms/step)、モバイルでOOMが発生しました。8×32×32 VAEはバランスの取れた性能(PSNR 30.6、380ms/step、VBench 79.80)を示しました。8×64×64 VAEは最速(90ms/step)でしたが、品質が15%低下しました(VBench 78.40)。

4. 実用性評価

4.1 実装の容易性

提案手法は、Diffusersライブラリを使用して構築され、QK正規化とRotary Positional Embeddings(RoPE)を組み込んでいます。モバイル展開のメモリ効率をさらに確保するため、ネットワーク全体でRMSNormを採用しています。

CoreMLToolsを使用してiPhone 16 Pro Maxへの展開が可能で、半精度で動作します。プルーニングされたモデルは密で圧縮された形式に変換できるため、追加のコンパイルや特殊なハードウェアサポートなしでモバイルデバイス上で実行できます。

4.2 計算効率

最適化により、次の3つの効率化を達成しました。
推論ステップ数は通常の40ステップから4ステップへ削減し、10倍高速化を実現しました。VAE圧縮による高速化では、8×32×32 VAEにより、4×16×16と比較して約20倍高速を実現しました。モデルサイズは2.0Bから0.9Bパラメータへ削減し、55%削減を達成しました。

iPhone 16 Pro Max上で49×384×512の動画を10FPS以上で生成可能となり、消費者向けモバイルハードウェア上で初めてリアルタイムの拡散ベース動画合成を実現しました。

4.3 応用可能性

この技術により、次の4つの応用が可能になります。
第一に、モバイルアプリでのリアルタイム動画生成と編集が可能になります。第二に、AR/VRアプリケーションでのインタラクティブなコンテンツ生成が実現できます。第三に、エッジデバイスでのクリエイティブツールの開発が可能になります。第四に、低遅延が要求されるライブストリーミングアプリケーションへの応用が期待できます。

5. まとめと所感

5.1 論文の意義

この研究は、DiTベースの動画生成モデルを初めてモバイルデバイス上でリアルタイム実行可能にした画期的な成果です。高度な最適化技術を組み合わせることで、品質を大幅に犠牲にすることなく、計算コストを劇的に削減することに成功しています。

特に、VAE圧縮率の体系的な調査、感度を考慮したプルーニング手法、DiT専用の敵対的蒸留設計という3つの革新的なアプローチを統合した点は高く評価できます。これらの技術は個別でも価値がありますが、組み合わせることで相乗効果を生み出しています。

ただし、高圧縮の潜在空間とDiTプルーニングにより、高速な動きや複雑なテクスチャのシーンでは細かい詳細が時折劣化するという制限もあります。

5.2 今後の展望

現在の実装では121フレーム、576×1024解像度の動画生成に最適化されていますが、今後は次の拡張が期待されます。
第一に、より長い動画シーケンスや高解像度への対応が必要です。第二に、細かいディテールの保持を改善する新しい圧縮とプルーニング手法の開発が重要です。第三に、他のモバイルプラットフォーム(Android、エッジデバイス)への展開が期待されます。第四に、さらなる推論ステップの削減(2ステップ生成など)の実現が望まれます。

この研究は、モバイルAIアプリケーションの新しい時代を切り開く重要な一歩であり、今後のエッジデバイス向け生成AIの発展に影響を与えることが期待されます。特に、計算資源が限られた環境でも高品質なコンテンツ生成を可能にすることで、AIの民主化に貢献する技術と言えるでしょう。