拡散+超解像モデルの強力な組み合わせ、Googleの画像ジェネレーターImagenの背後にある技術

近年、マルチモーダル学習は、特にテキストと画像の合成や画像とテキストの対照学習の分野で大きな注目を集めています。 OpenAI のテキスト画像モデル DALL-E および DALL-E 2、NVIDIA の GauGAN および GauGAN2 など、一部の AI モデルは、クリエイティブな画像の生成と編集への応用により、広く注目を集めています。

負けじと、Google は 5 月末に独自のテキストから画像への変換モデルである Imagen をリリースしました。これは、キャプション条件付き画像生成の限界をさらに押し広げたようです。

シーンの説明だけを与えれば、そのシーンが現実世界で意味を成すかどうかに関係なく、Imagen は高品質で高解像度の画像を生成できます。以下は、Imagen のテキスト生成画像のいくつかの例です。対応するキャプションが画像の下に表示されます。

これらの印象的な生成された画像を見ると、「Imagen はどのように機能するのだろう?」と疑問に思うでしょう。

最近、開発者インストラクターの Ryan O'Connor が AssemblyAI ブログに「Imagen の実際の動作方法」という長い記事を執筆しました。この記事では、Imagen の仕組みを詳しく説明し、Imagen の概要を示し、その高レベルコンポーネントとそれらの関係を分析して理解します。

Imagen の動作原理の概要

このセクションでは、著者は Imagen の全体的なアーキテクチャを紹介し、その動作原理を高レベルで説明した後、Imagen の各コンポーネントをより詳細に分析します。次のアニメーションは、Imagen のワークフローを示しています。

まず、字幕がテキストエンコーダーに送られます。エンコーダーは、テキストキャプションを、テキスト内の意味情報をカプセル化する数値表現に変換します。 Imagen のテキストエンコーダーは Transformer エンコーダーであり、自己注意アプローチを使用して、テキストエンコーディングがキャプション内の単語が互いにどのように関連しているかを理解できるようにします。

Imagen が単語の関連性ではなく個々の単語にのみ焦点を当てる場合、キャプションの個々の要素を捉えた高品質の画像を取得できますが、キャプションの意味を反映するように画像を適切に記述することはできません。以下の例に示すように、単語間の関連性を考慮しないと、まったく異なる生成効果が発生します。

テキストエンコーダーはImagenのキャプション入力に対して有用な表現を生成しますが、この表現を使用して画像を生成する方法、つまり画像ジェネレーターを設計する必要があります。この目的のために、 Imagen は、複数のタスクでの SOTA パフォーマンスにより近年人気が高まっている生成モデルである拡散モデルを使用します。

拡散モデルは、ノイズを追加してトレーニングデータを破損させることによってトレーニングされ、その後、このノイズプロセスを逆にすることでデータを回復することを学習します。入力画像が与えられると、拡散モデルは一連の時間ステップにわたってガウスノイズで画像を反復的に破損させ、最終的にガウスノイズまたはテレビの雑音を残します。次の図は、拡散モデルの反復ノイズプロセスを示しています。

次に、拡散モデルは逆方向に動作し、各タイムステップでノイズを分離して除去する方法を学習して、発生した破壊的なプロセスに対抗します。トレーニングが完了したら、モデルを 2 つに分割できます。これは、ランダムにサンプリングされたガウスノイズから開始し、拡散モデルを使用して徐々にノイズを除去して、次の図に示すように画像を生成することができます。

要約すると、トレーニングされた拡散モデルはガウスノイズから始まり、トレーニング画像に似た画像を繰り返し生成します。画像の実際の出力を制御できないことは明らかです。単にガウスノイズをモデルに入力するだけで、トレーニングデータセットに属しているように見えるランダムな画像が出力されます。

しかし、 Imagen に入力されたキャプションの意味情報をカプセル化した画像を作成することが目的であるため、キャプションを拡散プロセスに組み込む方法が必要です。これをどうやって行うのでしょうか?

前述のように、テキストエンコーダーは、実際にはベクトルシーケンスである代表的な字幕エンコーディングを生成します。このエンコードされた情報を拡散モデルに注入するために、これらのベクトルが集約され、それに基づいて拡散モデルが調整されます。このベクトルを調整することで、拡散モデルはノイズ除去プロセスを調整する方法を学習し、キャプションによく一致する画像を生成します。プロセスの視覚化図を以下に示します。

画像ジェネレーターまたはベースモデルは小さな 64x64 画像を出力するため、このモデルを最終的な 1024x1024 バージョンにアップサンプリングするために、超解像度モデルを使用して画像をインテリジェントにアップサンプリングします。

超解像モデルの場合、Imagen は再び拡散モデルを使用します。全体的なパイプラインは、基本的にベースモデルと同じですが、字幕のエンコードのみに基づいてスケーリングするのではなく、アップサンプリングされる小さい画像もスケーリングします。プロセス全体の視覚化を以下に示します。

この超解像度モデルの出力は、実際には最終出力ではなく、中サイズの画像です。この画像を最終的な 1024 x 1024 解像度にアップスケールするには、別の超解像度モデルが使用されます。 2 つの超解像アーキテクチャはほぼ同じなので、詳細には説明しません。 2 番目の超解像モデルの出力が Imagen の最終出力になります。

Imagen が DALL-E 2 より優れているのはなぜですか?

Imagen が DALL-E 2 よりも優れている理由を正確に答えることは困難です。しかし、パフォーマンスの差の大部分は字幕とプロンプトの違いから生じています。 DALL-E 2 は、コントラスト目標を使用して、テキストエンコーディングが画像 (基本的には CLIP) とどの程度密接に関連しているかを判断します。テキストエンコーダーと画像エンコーダーは、類似するキャプションと画像のペアのコサイン類似度が最大化され、類似しないキャプションと画像のペアのコサイン類似度が最小化されるようにパラメーターを調整します。

パフォーマンスのギャップの大部分は、Imagen のテキストエンコーダーが DALL-E 2 のものよりもはるかに大きく、より多くのデータでトレーニングされたという事実に起因しています。この仮説の証拠として、テキストエンコーダーで拡張した場合の Imagen のパフォーマンスを調べることができます。以下は、Imagen のパフォーマンスのパレート曲線です。

テキストエンコーダーのスケールアップは驚くほど効果的ですが、U-Net のスケールアップは驚くほど効果がありません。この結果は、強力なエンコーディングが条件である限り、比較的単純な拡散モデルでも高品質の結果を生成できることを示唆しています。

T5 テキストエンコーダーは CLIP テキストエンコーダーよりもはるかに大きく、自然言語トレーニングデータは必然的に画像とキャプションのペアよりも豊富であるという事実と相まって、パフォーマンスのギャップの多くはこの違いに起因すると考えられます。

さらに、著者は、次のような Imagen の重要なポイントをいくつか挙げています。