Sora / DiT

OpenAI 2024 · 論文

一行要約 — Soraは、動画のスペースタイムパッチに対してDiffusion Transformer(DiT)アーキテクチャをスケールさせることで、テキストから長く時間的に一貫した動画を生成する。その出力には創発的な3D一貫性が見られ、大規模な動画生成が暗黙的にシーンの幾何構造を学習していることを示唆している。

問題

ここでは2つの問題が交差している。第一に、画像拡散モデルは畳み込みU-Netのバックボーン上に構築されており、そのスケーリング挙動は十分に理解されていなかった。DiT(Peebles and Xie、“Scalable Diffusion Models with Transformers”、arXiv:2212.09748)は、潜在パッチに対する単純なTransformerがデノイザーとして機能し、他の領域のTransformerと同様にスケールするかを問うた。第二に、動画生成器は固定された解像度と長さの短いクリップに限定されていた。OpenAIのテキスト-動画生成システムであるSoraは、DiT型のバックボーンとインターネット規模で学習された統一的な動画表現を組み合わせている — OpenAIはこれを明示的に「世界シミュレータ」への一歩として位置づけている。

手法とアーキテクチャ

DiTの論文はLatent Diffusion Model(LDM)の枠組みで動作する。凍結されたVAEエンコーダ $E$ が $256\times 256\times 3$ の画像を形状 $32\times 32\times 4$ の潜在表現 $z = E(x)$ に写し、拡散モデルはその潜在空間上で学習される。パイプラインは以下の通りである。

• パッチ化: 形状 $I\times I\times C$ のノイズ付き潜在表現を、サイズ $p\times p$ のパッチに分割し、幅 $d$ の $T=(I/p)^2$ 個のトークン列に線形埋め込みし、ViT型のサイン・コサイン位置埋め込みを付与する。$p$ を半分にするとトークン数(とGflops)は4倍になる一方、パラメータ数は実質的に変化しない。
• DiTブロック: 4種類の構成(DiT-S/B/L/XL、0.3から118.6 Gflopsに及ぶ)からなる、標準的なTransformerブロックを$N$個積み重ねたもの。拡散のタイムステップ $t$ とクラス $c$ への条件付けについては、4つの設計 — インコンテキストトークン、クロスアテンション(約15%のGflopsオーバーヘッド)、adaptive layer norm(adaLN)、そしてadaLN-Zero — が比較される。adaLN-Zeroでは、スケール/シフトパラメータ $\gamma, \beta$ とブロックごとの残差ゲート $\alpha$ が $t$ と $c$ の埋め込みから回帰され、$\alpha$ はゼロに初期化されるため、すべてのブロックは恒等関数として学習を開始する。
• 学習目的関数: 標準的なDDPMのノイズ予測。前向き過程を $x_t = \sqrt{\bar{\alpha}_t}\, x_0 + \sqrt{1-\bar{\alpha}_t}\, \epsilon_t$、$\epsilon_t \sim \mathcal{N}(0,\mathbf{I})$ とすると、デノイザーは次を最小化する。

$$\mathcal{L}_{simple}(\theta) = \lVert \epsilon_\theta(x_t) - \epsilon_t \rVert_2^2$$

一方で、逆過程の共分散 $\Sigma_\theta$ は完全なKL項で学習される。線形デコーダが各トークンをノイズと共分散の予測値 $p\times p\times 2C$ に写し戻す。

• Classifier-free guidance: $\hat{\epsilon}_\theta(x_t, c) = \epsilon_\theta(x_t, \emptyset) + s\cdot\big(\epsilon_\theta(x_t, c) - \epsilon_\theta(x_t, \emptyset)\big)$、ガイダンススケール $s > 1$。
• Sora(OpenAIの技術レポートによる。数式を伴う論文は存在しない): 動画は視覚エンコーダによってスペースタイム的な潜在表現に圧縮され、スペースタイムパッチトークン(DiTのパッチ化の動画版)に分割される。これにより単一の統一されたトークン形式が得られ、単一のモデルが可変の長さ・解像度・アスペクト比の動画に対して学習・生成を行える。

実験結果

DiTの中心的な発見は、モデルのGflopsとFID-50Kとの間に強い負の相関があることである。12種類のモデル(構成S/B/L/XL、パッチサイズ8/4/2、クラス条件付きImageNet)にわたって、深さ・幅・トークン数を増やすと一貫してFIDが改善し、総Gflopsが同程度のモデル(例えばDiT-S/2とDiT-B/4)は同程度のFIDに達する — パラメータ数ではなく計算量が重要であることを意味する。adaLN-Zeroはすべての学習段階でクロスアテンションおよびインコンテキスト条件付けを上回る。40万イテレーションではそのFIDはインコンテキスト版のほぼ半分である。700万ステップ学習したDiT-XL/2(118.6 Gflops)は、classifier-free guidance付きでクラス条件付きImageNet 256×256においてFID 2.27を達成し、それまでの拡散モデルの最良値3.60(LDM)や、当時の最先端であったStyleGAN-XLを上回った。しかもピクセル空間のU-Net(ADM: 1120 Gflops)に比べてフォワードパスあたりのコストは遥かに低い。512×512でも最先端であり、235万ステップの時点で既にFID 2.55に達している。Sora自体は技術レポートとプロダクトとしてリリースされ、定量的な結果はなく質的な結果のみが示された — 長く、高忠実度で、時間的に一貫した動画に、創発的な3D一貫性と物体の永続性が見られる。全ての実演例についてはOpenAIのレポートを参照のこと。

SLAMにおける意義

Soraは、インターネット規模の動画学習が3D教師データなしに暗黙的な3Dシーン理解をもたらすことを示す、最も強力な公開の証拠である — これはGAIA-1が運転領域で観察したものと同じ創発特性である。これはSLAMにとって2つの方向で重要である。第一に、DiT型の生成モデルは、身体性を持つシステムの学習データを合成する世界基盤モデル(Cosmosおよびその後継)の中核エンジンになりつつある。第二に、「動画生成器の内部にある3D知識を、明示的で計量的に一貫したマップとして抽出できるか」という問いは、SLAMと生成AIの境界における現在進行中の研究フロンティアである。NVIDIA Cosmosは、その拡散世界モデルにDiTのadaLN層(LoRA付き)を採用している。

関連ノート

• World model
• GAIA-1
• NVIDIA Cosmos
• DreamFusion