自己回帰デコードが64倍高速化、Googleが新たな画像合成モデルMaskGITを提案

生成トランスフォーマーは、高忠実度かつ高解像度の画像を合成するために急速に人気を集めています。しかし、これまでの最高の生成トランスフォーマー モデルは、依然として画像を一連のトークンとして扱い、ラスター スキャン順 (つまり行ごと) で画像をデコードします。しかし、この戦略は最適でも効率的でもありません。

最近、Google Research の研究者は、双方向トランスフォーマー デコーダーを使用した新しい画像合成モデル MaskGIT を提案しました。トレーニング中、MaskGIT はあらゆる方向のトークンに注意を払うことでランダムなマスク トークンを予測することを学習します。推論フェーズでは、モデルは最初に画像のすべてのトークンを同時に生成し、次に前の生成に基づいて画像を反復的に改良します。実験では、MaskGIT が ImageNet データセット上の SOTA トランスフォーマー モデルを大幅に上回り、自己回帰デコードを 64 倍高速化することが示されています。

論文アドレス: https://arxiv.org/abs/2202.04200

さらに、この研究では、MaskGIT が修復、外挿、画像操作などのさまざまな画像編集タスクに簡単に拡張できることも示されています。

関連研究

以前のモデル VQVAE では、潜在空間に 2 段階で画像を生成することを提案していました。

最初の段階はトークン化と呼ばれ、ここでは画像を個別の潜在空間に圧縮しようとします。この段階は主に 3 つの部分で構成されます。

画像 x∈ を潜在的な埋め込み E(x) にトークン化することを学習するエンコーダー E、埋め込みを視覚トークンに量子化するための最近傍検索用のコードブック、および視覚トークン e の予測に基づいて画像を再構築するデコーダー G。

第 2 段階では、まず深層自己回帰モデルを使用して視覚トークンの潜在的な事前確率を予測し、次に第 1 段階のデコーダーを使用してトークン シーケンスを画像ピクセルにマッピングします。

この 2 段階のパラダイムは非常に効果的であるため、DALL-E や VQGAN など、一般的に使用されるいくつかの方法はこのパラダイムに従います。その中で、VQGAN は、画像の忠実度を向上させるために、最初の段階で敵対的損失と知覚的損失を追加します。

マスクGIT

上記の 2 段階の方法では依然として自己回帰モデルが使用されるため、第 2 段階のデコード時間はトークン シーケンスの長さに比例します。この研究の目的は、上記の 2 段階方式に従い、第 2 段階を改良して、並列デコードと双方向生成を活用する新しい画像合成パラダイムを設計することです。最初の段階では、VQGAN モデルと同じ設定を採用し、トークン化ステップの潜在的な改善を将来の作業に残します。第 2 段階では、研究者は、マスクされたビジュアル トークン モデリング (MVTM) を通じて双方向トランスフォーマーを学習することを提案しています。

トレーニング中のMVTM

この研究では、画像を VQ エンコーダーに入力することによって得られる潜在的なトークンを表すためにを使用します。ここで、N は再構成されたトークン マトリックスの長さであり、は対応するバイナリ マスクです。トレーニング中に、トークンのサブセットをサンプリングし、特別な [MASK] トークンに置き換えます。 m_i=1 の場合、トークン y_i を [MASK] に置き換えます。m_i=0 の場合、y_i は保持されます。

サンプリング プロセスはマスク スケジューリング関数によってパラメーター化され、次の手順に従います。

まず0から1までの比率をサンプリングし、次に均一に選択する マスクを配置するトークン。N は長さです。マスクのスケジューリングは、生成される画像の品質に大きな影響を与えます。

反復復号

自己回帰復号化では、以前に生成された出力に基づいてトークンが順次生成されます。このプロセスは並列化できず、画像のトークンの長さは通常言語のトークンの長さよりもはるかに長いため、非常に遅くなります。本研究では、MTVMの双方向自己注意に基づいて、画像内のすべてのトークンが同時に並列に生成される新しいデコード方法を提案しました。

理論上、モデルはすべてのトークンを推測し、1 回のパスで画像全体を生成できるはずですが、トレーニング タスクの不一致がこの研究の課題となりました。推論時に画像を生成するには、すべてのトークンがマスクされた空白のキャンバスから始めます。本研究で提案する反復復号法は、反復ごとに以下のアルゴリズム演算ステップを持つ。

1. 予測 2. サンプリング 3. マスクスケジューリング 4. マスク

マスクデザイン

研究者たちは、画像生成の品質はマスクの設計に大きく影響されることを発見しました。この方法では、特定の潜在的なトークンのマスク比率を計算するマスク スケジューリング関数を通じてマスキング プロセスをモデル化します。推論中、関数によって使用される入力はデコードの進行状況を表します。トレーニング中は、[0,1) 内の比率 r をランダムにサンプリングして、さまざまなデコード シナリオをシミュレートします。

実験

この研究では、画像生成におけるMaskGITの品質、効率、柔軟性の観点から実験的に評価しました。

クラス条件付き画像合成

この研究では、ImageNet 256 X 256 および ImageNet 512 X 512 のクラス条件付き画像合成タスクにおける MaskGIT モデルのパフォーマンスを評価しました。主な結果を以下の表 1 に示します。

品質。 ImageNet 256 X 256 では、特別なサンプリング戦略を使用せずに、MaskGIT は FID と IS の両方で VQGAN を大幅に上回ります。

スピード。この研究では、各モデルがサンプルを生成するために必要なステップ数 (フォワードパス) を評価することで、モデルの速度を評価しました。表 1 に示すように、MaskGIT は、すべての非 GAN ベースのモデルの中で、両方の解像度で最も少ないステップ数を必要とします。

MaskGIT と自己回帰モデルの速度の違いをさらに確認するために、本研究では、MaskGIT と VQGAN のデコード プロセスの実行時間の比較を行いました。下の図 4 に示すように、MaskGIT は VQGAN を 30 ～ 64 倍大幅に高速化し、画像解像度 (および入力トークンの長さ) が増加するにつれて高速化がより顕著になります。

多様性。この研究では、サンプルの品質に加えて、分類精度スコア (CAS) と精度/再現率という 2 つの指標を使用して、サンプルの多様性を評価します。 BigGAN のサンプルと比較すると、MaskGIT のサンプルはより多様で、照明、ポーズ、スケール、コンテキストの種類が多くなっています (下の図 5 を参照)。

画像編集アプリケーション

この研究では、クラス条件付き画像編集、画像修復、画像アウトペインティングという 3 つの画像編集タスクに MaskGIT を直接適用する方法を示します。タスクを、MaskGIT が反復デコードで使用する初期バイナリ マスク M に対する制約と見なすと、3 つのタスクのほぼすべてを、MaskGIT が処理できるタスクに簡単に変換できます。

調査では、MaskGIT はアーキテクチャの変更やタスク固有のトレーニングを行わなくても、3 つのアプリケーションすべてで非常に優れた結果を生成できることが示されています。さらに、MaskGIT は、画像の修復と拡張に関して専用モデルと同等のパフォーマンスを実現します。

クラス条件付き画像編集タスクでは、この研究では、MaskGIT の柔軟性を実証するために新しいクラス条件付き画像編集タスクを定義します。このモデルは、コンテキスト、つまりボックスの外側のコンテンツを保持しながら、特定のクラスの境界ボックス内の特定のコンテンツを再現します。予測順序が破られているため、自己回帰法は実行できません。

ただし、MaskGIT の場合、境界ボックス領域を反復復号アルゴリズムの初期マスクの入力と見なすと、この問題は解決されます。下の図 6 にいくつかの結果例を示します。

表 2 は、いくつかの方法の定量的な結果を比較したものです。 MaskGIT は、FID と IS の両方で DeepFill と HiFill を大幅に上回り、SOTA インペインティング手法 CoModGAN に近いスコアを達成しました。

下の図 7 に示すように、MaskGIT は同じ入力と異なるシードを与えても異なる結果を合成できます。

アブレーション実験

新しい設計の有効性を検証するために、本研究では、ImageNet 256×256 のデフォルト設定でアブレーション実験を実施しました。 MaskGIT の重要な設計は、トレーニングと反復デコードに使用されるマスク スケジューリング機能です。実験結果を表 3 と図 8 に示します。

図 8 に示すように、同じ設定では、反復回数を増やすことが必ずしも良い結果をもたらすわけではないことに注意してください。反復回数 T が増加すると、プロセス全体でパフォーマンスが低下する対数関数を除き、他のすべての関数は「スイート スポット」に到達します。つまり、モデルのパフォーマンスはピークに達してから再び低下します。