テキストから画像を生成するモデルを3.6倍高速にトレーニング

TL;DR
Linum v2は、その巨大なアテンションコンテキストウィンドウのサイズによってボトルネックとなっていました。720p、5秒のクリップで110Kトークンという膨大なコストがかかりました。比較のために言うと、LLMは事前トレーニングの97%で8Kトークン未満のサンプルを見ています。アテンションはコストに対して二次関数的に増加するため、モデルトレーニングを加速するための最大のレバーは、コンテキストウィンドウを狭めることです。
ほとんどの生成画像およびビデオシステムはLatent Diffusion Models（LDMs）です。これらは、独立してトレーニングされたモジュールであるVariational Autoencoder（VAE）とDiT（Diffusion Transformer）に圧縮と生成を分割します。最近、JiTのようなピクセル空間モデルが有望な代替手段として示されています。これは2つのモデルを1つに減らし、拡散モデルが再構築用に構築された潜在空間に依存するのではなく、生成専用の潜在空間を構築できるようにします。
（画像、キャプション）データセットでトレーニングすると、JiTは細かいディテールを生成するのに苦労するようです。私たちは、このディテールを回復し、LDMの競合相手よりもはるかに効率的にトレーニングできる、新しいエンコーダー・デコーダーアーキテクチャ（JiT-DDT）を提案します。Linum v2のベースラインと比較して、JiT-DDTは、4倍のピクセルを持つ画像を生成するにもかかわらず、3.6倍少ないGPU時間でテキストから画像を生成するモデルをトレーニングします。トレーニングが3.6倍高速、ピクセルは4倍
Linum v2（当社、以前）* · 256×256
2.0B ラテント空間 DiT + VAE
256 ラテントトークン

• 画像のみのチェックポイント
  JiT-DDT（当社、新規） · 512×512
  2.5B アクティブ ピクセル空間 DiT
  320 ピクセル トークン = 64 エンコーダー + 256 デコーダー
  GPU時間
  0
  3.6倍少ない
  0
  サンプル数
  0M
  4.2倍少ない
  0M
  詳細な比較については、付録を参照してください。
  研究リリース
  モデルコード
  モデルウェイト
  JiT-DDTのコードとモデルウェイトは、Apache 2.0ライセンスの下で利用可能です。私たちの発見をより広いコミュニティと共有することで、他の人々もより効率的なトレーニング方法を探求することを奨励できることを願っています。これは、完全なモデルリリースではなく、研究成果として扱われるべきです。Linum v3に向けた、このようなさらなる研究チェックポイントにご期待ください。
  VAE圧縮の壁にぶつかる
  ほぼすべての生成画像およびビデオモデルはLatent Diffusion Models（LDMs）です。これらには2つの主要コンポーネントがあります。圧縮用のVariational Auto Encoder（VAE）と生成用のDiffusion Transformer（DiT）です。生のピクセルで操作するのは（特にビデオの場合）高価すぎるため、まずRGBピクセルをDiT用のより少ないトークン数に削減する方法を見つける必要があります。ここでVAEが登場します。これは圧縮と再構築のためにトレーニングされます。具体的には、ピクセル空間サンプルを確率的エンコーダーに通し、k次元トークンを出力し、これらのラテントトークンを確率的デコーダーに通してピクセル空間に戻します。
  VAEは圧縮と再構築のためにトレーニングされます。
  replay
  input x
  Encoder
  encoder
  μ = [?, ?]
  σ = [?, ?]
  z ∈ ℝᵏ
  sample z
  Decoder
  decoder
  output x̂
  ‖x − x̂‖²+ β · KL(q‖N)
  loss
  ‹ready›
  gradient (purple) reaches every weight
• 簡略化: 実際にはKL項は≈0です。以前は、L1再構築損失にLPIPSとGAN損失を加えたσ-VAEをトレーニングしました。VAE投稿を参照してください。
  LDMを構築する際、VAEを別々にトレーニングし、それをフリーズします（つまり、DiTからVAEへの勾配フローはありません）。これにより、DiTトレーニングの過程でラテント空間は静的になります。VAEのエンコーダーを実行してデータを埋め込み、DiTをトレーニングしてVAEのラテント空間を移動させ、DiTが生成したラテントトークンをVAEのデコーダーを使用してピクセル空間に変換します。
  VAEは一度トレーニングされフリーズされます。DiTはそのラテント空間を移動することを学習します。
  Training
  Inference
  replay
  input x
  ❄
  Encoder
  encoder = E(x)
  z(1−t)·z+ t·ε
  INTERPOLATE
  zₜ
  DiT
  trainable
  edit
  v̂pred
  ‖v̂ − v‖²
  v = ε − z
  loss
  ε ~ N(0, I)
  gaussian
  ε
  SAMPLE ε
  t ~ LogitNormalsample t
  ‹ready›
  VAE frozen (dashed) · DiT trainable (purple) · gradient stops at the DiT · t = 0 clean image, t = 1 pure Gaussian noise
  VAEから可能な限り多くのトークン圧縮を引き出し、DiTのアテンションコストを抑制したいと考えています。しかし、FLUX、Ideogram、Z-Imageのような一般的なオープンソースのテキストから画像へのモデルを調査すると、それらがすべて16×16トークン削減で頭打ちになっていることに気づくでしょう。これは、数年前のImage-Video VAEに関する私たちの実験と一致しています。残念ながら、標準的なCNN VAEから再構築を劣化させることなく得られる圧縮量には、経験的な上限があるようです。
  統一モデルによる積極的な圧縮の解除
  昨秋、Tianhong LiとKaiming Heは、VAEを完全に破棄し、圧縮タスクをDiT自体に押し込むことで32×32トークン削減を達成する論文（JiT）を発表しました。
  Patchify: 4×4ピクセルパッチ → 48次元トークン → 12次元への線形ボトルネック
  replay
  16×16ピクセル、各3チャンネル（RGB）
  4×4ピクセルパッチにカット（16パッチ）
  各パッチは独立してトークン化されます：16ピクセル×3チャンネルが1つの48次元トークンになります
  線形層 W ∈ ℝ12×48 が48次元を12次元に削減します
  ‹ready›
  説明用。JiTでは512pxで32×32パッチを使用するため、512×512の画像は256トークンになり、各トークンは32・32・3 = 3,072次元から始まります。ボトルネックはそれを256にマッピングします。
  このトークン数を削減するアプローチは、それほど新しくはありません。5年前にビジョントランスフォーマー（ViT）のために発明され、DiTに入る前にトークンシーケンスをさらに凝縮するためにVAEと組み合わせて使用されることが一般的です。
  Linum v2では、VAEは8×8（高さ×幅）の圧縮と16次元のラテントを提供しました。DiTのベースで、2×2パッチ化を適用して16×16トークン圧縮と64次元のラテントを取得しました。私たちはLinum v2でこれを使用しましたし、FLUXのようなモデルも同様です。では、なぜ誰もこれを以前に試さなかったのでしょうか？これは無料のランチのように感じられます。アテンションコストを削減できる（潜在的に）損失のない方法が得られ、非常にシンプルです。2025年初頭、VA-VAEのような論文は、DiTが高次元入力からの学習に苦労することを示しました。
  VAEトレーニング中に外部モデルを正規化として使用する（例：DINOv3）といった小さなハックがあり、（おそらく）FLUX-2のようなモデルがDiTのラテント次元を64から128に引き上げることを可能にしました。しかし、これらの戦略は、DiT内の明確な学習可能性の問題を先延ばしにするだけです。積極的なパッチ化は明示的に情報をチャンネル次元に押し込むため、この不安定性を引き起こします。しかし、これはアーキテクチャ固有のものではないことが判明しました。むしろ、過去数年間、拡散モデルのトレーニングに使用されてきたv予測、v損失フローマッチング目的関数の下流にあります。
  フローマッチングの簡単な復習
  古いスタイルの2022年頃のノイズ除去拡散（DDPM）では、サンプルを繰り返しノイズ化し、ニューラルネットワークをトレーニングしてノイズを除去します。これにより、推論時にニューラルネットワークを使用して、一連のステップにわたってガウスノイズからデータ分布のサンプルを生成できます。この定式化には多くの問題があります（例：生成時の過飽和、不安定な学習、蒸留崩壊）。そのため、その間にこの分野はフローマッチングへと移行しました。
  フローマッチングでは、データ分布の各サンプルとガウスノイズのサンプルとの間に直線パスを構築します。
  ノイズとサンプルの間のパスは直線である必要はありません。しかし、実際には、私たちは皆それを行います。t=0で、データ分布のサンプルを取得します。t=1で、ガウスノイズのサンプルを取得します。ここで、tはパス上の時間です。
  DDPMの慣習に従って、この投稿全体で：x₀はデータ、x₁はノイズです。一部のフローマッチング論文は逆方向に実行され、x₀をノイズ、x₁をデータとしています。2つの定式化は同等です。
  次に、パスに沿った速度を近似するネットワークをトレーニングします。これをv予測、v損失と呼びます。ニューラルネットワークが明示的に速度を予測し、その速度予測と真の条件付き速度場との間のMSEでトレーニングされるためです。
  v予測と次元の呪い
  フローマッチングモデルをトレーニングしている場合、必ずしもネットワークに速度を予測および回帰させる必要はありません。3つの項は線形に並べ替えることができるため、予測と回帰ターゲット間で混合および一致させることができます。
  1つの目的関数の9通りの書き方
  3つのターゲット、それぞれ線形