VAEなしでのテキスト画像生成モデルのトレーニング

TL;DR 前回の投稿では、ピクセル空間のエンコーダー・デコーダーアーキテクチャ（JiT-DDT）を開発しました。VAEを廃止することで、トークンをより積極的に圧縮できるようになり、結果としてアテンションのコストが削減され、トレーニングと推論が高速化しました。 しかし、固定されたパラメータ予算の下でエンコーダー・デコーダーアーキテクチャをどのようにスケールさせるかは不明瞭でした。そこで今回、DiTのトランクに沿って複数の解像度でターゲット画像を予測する、新しいデコーダーのみのピクセル空間アーキテクチャであるPyramid-JiT（P-JiT）を提案します。これは、Linum v2のFD-DINOv2を11.3倍少ない学習サンプルで達成し、4倍のピクセル数で4.3倍少ないGPU時間でトレーニングできます。P-JiTは、11.3倍少ないサンプルでLinum v2の最終FD-DINOv2に匹敵します。 Linum v2は、当社の以前の潜在拡散モデルです。Linum v2の画像チェックポイントは4つしか残っていないため、その曲線は117Mサンプルから始まります。チェックポイント間の線はログログで補間されています。 研究リリース モデルコード モデルウェイト P-JiTのコードとモデルウェイトはApache 2.0ライセンスで公開されています。この発見をより広いコミュニティと共有することで、他の人々もより効率的なトレーニング方法を探求することを奨励できると期待しています。これは、完全なモデルリリースではなく、研究成果物として扱われるべきです。Linum v3に向けた、このような研究チェックポイントにご期待ください。 パラメータはどこに配置すべきか？ 私たちの目標は、誰もが自分の物語を画面にもたらすことができる、手頃な価格のアニメーションツールを構築することです。現在、世界最高のビデオモデル（例: Seedance）は、非常に大きなトークンシーケンスでアテンションを計算するため高価です。Linum v2モデルでは、720p、5秒の短いビデオで110Kトークンが必要でした。これに対し、LLMはトレーニングの97%で8K以下のコンテキストサイズでトレーニングされます。このビジョンを実現するには、ビデオ生成のコストを1桁削減する必要があります。私たちは、コンテキストウィンドウのトークン数に対する操作であるアテンションによってボトルネックになっています。トークン数を削減できれば、アテンションのコストを削減し、トレーニングと推論全体で大幅な節約を達成できます。Linum v3ビデオモデルに向けて、コンテキストウィンドウを圧縮するのに役立つアイデアを試すためのテストベッドとしてテキスト画像生成を使用しています。 直感的には、ビデオモデルは動詞（アクション、動きなど）を学習する前に名詞（静的な人、場所、物）を学習する必要があります。その結果、ビデオモデルは画像モデルの拡張としてトレーニングされることがほとんどです。以前、LDM（Latent Diffusion Models）は、オートエンコーダーを大幅に調整しない限り、16x16トークン圧縮の経験的な天井があるように見えると書きました。そこで、ピクセル空間モデル（JiT）に移行し、32x32トークン圧縮を達成し、圧縮と生成を単一の合理化されたモデルに統合しました。そこから、以前のLDMよりも優れた結果を生み出し、3.6倍高速にトレーニングされた新しいピクセル空間アーキテクチャ（JiT-DDT）を開発しました。 しかし、エンコーダー・デコーダーアーキテクチャをスケールさせる方法が不明瞭でした。エンコーダーのサイズを大きくして構造をより良く学習させるべきか？そうすると、低解像度の画像を予測するために多くのパラメータを費やすのは無駄に思えます。あるいは、デコーダーのサイズを大きくして、人間にとって本当に重要な細部を学習するのに時間を費やすべきか？そうすることで、シーンのフレーミングや構造を学習するモデルの能力を制限してしまう可能性があります（それなしでは、詳細は無意味です）。 私たちのアーキテクチャは、エンコーダーにデコーダーよりも4倍多くのパラメータを割り当てるのが最適であると見出したDecoupled Diffusion Transformer（DDT）に触発されました。彼らは、デコーダーからクラス条件を奪うことで、フル解像度のv-spaceで動作しながら構造学習を促進します。私たちは、エンコーダーがターゲット画像の低解像度バージョンを予測させることで、x-spaceでの構造学習を強制します。これらの大きな違いを考えると、最適なエンコーダー・デコーダーサイズがJiT-DDTに引き継がれるとは限りませんでした。そして、初期の実験では、DDTによって提案された4:1のエンコーダー：デコーダーパラメータ比は、1:1のエンコーダー：デコーダーパラメータ比よりも劣っていました。 約10年前に、言語も同様のジレンマに直面しました。アーカイブを調べると、トランスフォーマーは当初、言語翻訳でトレーニングされたエンコーダー・デコーダーのビルディングブロックとして提案されていたことがわかります。その後、フィールドはエンコーダーのみのモデル（BERTスタイルの分類器）とデコーダーのみのモデル（GPTスタイルのLLM）に分かれました。デコーダーのみのモデルの美しさは、明示的なアーキテクチャ上のボトルネックがないため、スケーリングが非常に容易であることです。したがって、推進力のある質問は次のとおりです。エンコーダー・デコーダーからの進歩を、通常のJiTにポートして、同様に高品質な結果を得ることはできますか？ JiTの再訪 以前の研究では、トレーニング中にノイズスケジュールを広げ、別々のDiTブロックを追加して、画像ストリームとテキストストリームを独立して洗練させてから、単一の共有ストリームにマージすることで、生成の品質を劇的に向上させることができることがわかりました。そこで、これらのイノベーションをJiTにコピー＆ペーストしました。さらに、線形パッチ化のボトルネック次元を256から1024に広げ、ネットワーク内の実効圧縮率を12倍から3倍に削減しました。 オリジナルのJiT論文では、著者らはボトルネック次元を128以上に増やしても、ImageNet FIDのスコアが悪化することを発見しました。これは直感的ではありません。ボトルネック次元が大きいほど、実効圧縮率が低下するため、詳細をより簡単にキャプチャできるはずです。 32x32ピクセルでは、32*32*3（RGB）=3072の値が256次元のトークン（3072/256 = 12倍圧縮）に統合されます。1024のボトルネックを使用すると、これは3倍の実効圧縮（3072/1024）になります。256ボトルネックと1024ボトルネックの削減を試しました。実際には、結果はほぼ同じでした（これもそれ自体で奇妙です）。そのため、より広く、より表現力豊かなボトルネックを採用しました。そして、Z-Imageのイノベーションを取り入れました。RMS-Norm（LayerNormの代わりに）、tanhゲート（post-norm付き）、およびtruncated AdaLN。これらの他の機能はJiT-DDTには役立たなかったようですが、ポートしても害はありません。残差ストリームへの保守的な書き込みは、トレーニングをより安定させるのに役立つはずです。 JiT、JiT-DDT、および改訂版JiTを同じプロンプトで比較 JiT（前回の投稿） 2.0B アクティブピクセル空間DiT 256 画像トークン 100M トレーニングサンプル JiT-DDT（前回の投稿） 2.5B アクティブピクセル空間DiT 320 画像トークン 138M トレーニングサンプル 改訂版 JiT（新規） 2.18B アクティブピクセル空間DiT 256 画像トークン 75M トレーニングサンプル これらのアップデートにより、多くのパフォーマンスを回復できましたが、JiT-DDTの2つの大きな違いがまだ欠けていました。それはマルチ解像度予測とデュアルパッチ化です。 エンコーダーをリードアウトに置き換える JiT-DDTでは、エンコーダーに512x512入力画像の64x64バージョンを予測させていました。この仮説は、モデルが初期の低解像度の回答を作成するのを助け、デコーダーが詳細を埋めるために容量を割り当てられるようにすることでした。RAE v2やSelf-Flowのような隣接する作業を広範囲に見ると、エンコーダー予測にはさらに2つの具体的な目的があるように思われました。第一に、適応がはるかに遅いDiTの初期レイヤーで勾配信号をブーストしました。第二に、低解像度画像を回帰することで、DiTの初期部分が明示的に構造を学習するのを助けました（REPAのようなハックが学習を加速するために利用するもの）。エンコーダーを破棄しても、これらの改善にアクセスできます。私たちのモデルはx-predictionでトレーニングされているため、ever