Apple の「マトリョーシカ」拡散モデルはトレーニング ステップ数を 70% 削減します。

Apple による最近の研究により、高解像度画像における拡散モデルのパフォーマンスが大幅に向上しました。

この方法を使用すると、同じ解像度の画像のトレーニング ステップ数が 70% 以上削減されます。

1024×1024 の解像度では、画像の品質が最大化され、すべての詳細が鮮明に表示されます。

Appleはこの成果をMDMと名付けました。DMはDiffusion Modelの略で、最初のMはMatryoshkaを表しています。

MDM は、本物の入れ子人形のように、高解像度のプロセスの中に低解像度のプロセスを入れ子にし、複数のレイヤーにネストします。

高解像度と低解像度の拡散プロセスが同時に実行されるため、高解像度プロセスにおける従来の拡散モデルのリソース消費が大幅に削減されます。

解像度 256×256 の画像の場合、バッチ サイズが 1024 の環境では、従来の拡散モデルでは 150 万ステップのトレーニングが必要ですが、MDM では 39 万ステップしか必要ありません。これは 70% 以上の削減になります。

さらに、MDM はエンドツーエンドのトレーニングを使用し、特定のデータセットや事前トレーニング済みのモデルに依存しません。生成品質を確保しながら高速化を実現し、柔軟に使用できます。

高解像度の画像を描画できるだけでなく、16×256²の動画を合成することもできます。

一部のネットユーザーは、Appleがついにテキストと画像を関連付けたとコメントした。

では、MDM の「入れ子人形」テクノロジーは具体的にどのように機能するのでしょうか?

全体的かつ漸進的なアプローチを組み合わせる

トレーニングを開始する前に、データを前処理する必要があります。高解像度の画像は、特定のアルゴリズムを使用して再サンプリングされ、さまざまな解像度のバージョンが取得されます。

次に、これらの異なる解像度のデータを使用して、共同 UNet モデリングを実行します。小さな UNet は低解像度を処理し、高解像度を処理する大きな UNet にネストされます。

クロス解像度接続により、さまざまなサイズの UNets が機能とパラメータを共有できます。

MDM トレーニングは段階的なプロセスです。

モデリングは共同で行われますが、トレーニングプロセスは高解像度から始まるのではなく、低解像度から始めて徐々に拡大していきます。

そうすることで、膨大な量の計算を回避できるだけでなく、低解像度の UNet 事前トレーニングによって高解像度のトレーニング プロセスを加速できるようになります。

トレーニング プロセス中に、より高解像度のトレーニング データが徐々にプロセス全体に追加され、モデルが徐々に増加する解像度に適応し、最終的な高解像度のプロセスにスムーズに移行できるようになります。

ただし、全体的に見ると、高解像度のプロセスが徐々に追加された後も、MDM のトレーニングはエンドツーエンドの共同プロセスのままです。

異なる解像度での共同トレーニングでは、複数の解像度での損失関数を一緒に使用してパラメータを更新し、多段階のトレーニングによって発生するエラーの蓄積を回避します。

各解像度には対応するデータ項目の再構築損失があり、異なる解像度の損失は重み付けされて結合されます。生成の品質を確保するために、低解像度の損失重みは大きくなります。

推論フェーズでは、MDM は並列処理と増分性を組み合わせた戦略も採用します。

さらに、MDM は、事前にトレーニングされた画像分類モデル (CFG) を使用して、生成されたサンプルの最適化をより合理的な方向に導き、低解像度のサンプルにノイズを追加して、高解像度のサンプルの分布に近づけます。

では、MDM の効果は何でしょうか?

SOTAに一致するパラメータが少ない

画像に関しては、ImageNet および CC12M データセットでは、MDM の FID (値が低いほど効果が高い) および CLIP のパフォーマンスが、通常の拡散モデルよりも大幅に優れています。

FID は画像自体の品質を評価するために使用され、CLIP は画像とテキスト指示の一致度を記述します。

DALL E や IMAGEN などの SOTA モデルと比較すると、MDM のパフォーマンスも非常に近いですが、MDM のトレーニング パラメーターはこれらのモデルよりもはるかに少なくなっています。

MDM は通常の拡散モデルよりも優れているだけでなく、他のカスケード拡散モデルよりも優れています。

アブレーション実験の結果によると、低解像度のトレーニングのステップ数が多いほど、MDM 効果が顕著になります。一方、ネストされたレベル数が多いほど、同じ CLIP スコアを達成するために必要なトレーニング ステップ数が少なくなります。

CFG パラメータの選択に関しては、複数のテスト後の FID と CLIP のトレードオフの結果です (CLIP スコアが高いほど CFG 強度が増加することを意味します)。