Skip to main content

[概要]高解像度画像合成のためのTransformerの制御

Transformerの表現力を備えた畳み込みアプローチの効率.
Created on February 26|Last edited on February 16
このレポートは、akshay uppalによる「[Overview] Taming Transformers for High-Resolution Image Synthesis」の翻訳です。



このレポートは、論文「高解像度画像合成のためのTransformerの制御」の概要を示しています。この論文では、Transformerの表現力を使用して高解像度の画像を合成する際の基本的な課題について説明します。提案されたアプローチは、画像を永続的に豊富な画像構成要素の合成として表し、次に、Transformerを利用して、高解像度画像内でそれらの合成を効率的にモデル化します。
これは、この論文について叙述したSander Dieleman氏による素晴らしいTwitterスレッドです。

 論文 |  プロジェクトページ| GitHub

Colab Notebookをチェックして、次のような素晴らしい画像を生成してください。

Run set
1


概要

畳み込みニューラルネットワーク(CNN)は、視覚タスクの頼りになるモデルです。これは、CNNが強い局所性バイアス(カーネルの使用による)と、共有重みの使用による空間的不変性へのバイアス(カーネルが画像全体をスイープする)を示すためです。
CNNと比較すると、Transformerには、ローカルの相互作用を優先する誘導バイアスが含まれていないため、入力間の複雑な関係を学習し、表現力を高めます。ただし、これは長期間のシーケンスでは計算コストが高くなります。変圧器の表現力の向上は、二次関数的に増加する計算コストを伴います。
Esser氏等(2021)は、CNNの有効性とTransformerの表現力の組み合わせをもって、高解像度の画像合成を示しています。


生成モデルに関する2つの言葉

ここで、 pdata(x)p_{data}(x) iはデータセットXを表す真の分布であり、データセットはこの分布からの有限サンプルで構成されます。
X=xxpdata(x)X = {x|x \approx p_{data}(x)}
生成タスクは、 pmodel(x;θ)pdata(x)p_{model}(x; θ) \approx p_{data}(x)のようなモデルを見つけることです。ここで、θはモデルパラメータです。
図1 : 生成モデルの分類法(出典
生成モデルには主に2つのタイプがあります。
  • 尤度ベース(明示的)モデル:ここでのこれらのモデルは、データ分布の「明示的」パラメトリック仕様を提供し、扱いやすい尤度関数を備えています。たとえば、変分オートエンコーダー(VAE)や自己回帰モデルなどです。
  • 暗黙的モデル:これらのモデルは、データ自体の分布を指定しませんが、代わりに、トレーニング後に、基になるデータ分布からサンプルを抽出することを目的とした確率過程を定義します。たとえば、生成的敵対的モデル(GAN)。
このホワイト論文をよりよく理解するための重要なアイデアは次のとおりです。
  • 暗黙的なGANは評価が難しく、通���はデータのすべての「モード」をカバーできないため、モデルが崩壊しやすくなります。
  • 大規模なGANは、高品質で高解像度の画像を生成できるようになりました。ただし、これらのモデルのサンプルが真の分布の多様性を完全に捉えていないことはよく知られています。(出典
  • 尤度ベースの方法は、トレーニングデータの負の対数尤度(NLL)を最適化し、モデルの比較を容易にし、見えないデータへの一般化を改善します。ただし、ピクセル空間で尤度を最大化することは困難であり、計算コストがかかります。
  • 自己回帰モデルでは、例 xXx \in X をシーケンスxiとして表すことができると仮定します。分布は、次のように与えられる確率の連鎖律を使用して、条件の積に因数分解されます:: pdata(x)=q(x1)q(x2x1)..q(xdxd1,...,x1).p_{data}(x) = q(x_1)q(x_2|x_1)..q(x_d|x_{d-1},..., x_1).。モデルは、前の x<ix_{<i}​に基づいて次のxi​を予測します。画像生成は、自己回帰シーケンス生成または変換問題として正常にキャストされています。(出典
  • 条件付きGAN(cGAN)は、通常のGANに対する単純で効果的な変更であり、クラスラベルやその他の画像(部分画像、セグメンテーションマップなど)などの条件に基づいて画像を生成できます。この優れたブログ投稿からcGANの詳細を学び、このColabを使用して簡単なcGANを試すことができます。

Run set
1




提案手法の概要

画像生成にTransformerを適用した以前の作品は、64x64ピクセルのサイズまでの画像で有望な結果を示しましたが、シーケンスの長さとともにコストが2倍に増加するため、より高い解像度にスケーリングできませんでした。したがって、Transformerを使用して高解像度の画像を合成するには、画像のセマンティクスを巧みに表現する必要があります。画像の解像度が2倍になるとピクセル数が二次関数的に増加するため、ピクセル表現の使用は機能しません。

ベクトル量子化変分オートエンコーダー(VQ-VAE)

この作業は、2つの重要な点で従来のVAEとは異なるベクトル化量子化変分オートエンコーダーに触発されています。
  • エンコーダネットワークは、連続コード(画像の潜在的表現)ではなく、離散コードを出力します。
  • 静的な多変量正規分布の代わりに、事前分布が学習されます。
VAEの詳細については、私のレポート「W&Bを使用したディープジェネレーティブモデリングに向けて」を参照してください。しかし、この作業はVQ-VAE2のより自然な進化であり、自己回帰生成モデリングのコンテキストで、表現学習がいかに強力であるかを示しています。
図2: VQ-VAEアーキテクチャの概要。(出典
VQ-VAEは、観測値(画像)を離散潜在変数のシーケンスにマッピングするエンコーダー(E(.)E(.)tthe)と、これらの離散変数から観測値を再構築するデコーダー(G(.)G(.))で構成されます。彼らは共有コードブックを使用します。コードブックは eRK×De \in R^{K \times D}で与えられます。ここで、Kはコードブック内の離散コードベクトルのサイズであり、Dは各コードの次元数 eie_i, i1,2,..,Ki \in 1, 2,..,K. VQ-VAEです。
図2に示すように、画像xは EE を通過してE(x)E(x)を生成します。次に、これはコードベクトルeiまでの距離に基づいて量子化され、各ベクトル E(x)E(x) がコードブック内の最も近いコードベクトルのインデックスに置き換えられます。同じことがデコーダーによって再構築のために使用されます。量子化は、によって与えられます。
Quantization(E(x))=ekQuantization(E(x)) = e_k where k=argminjE(x)ejk = \underset{j}{\operatorname{argmin}} || E(x) - e_j ||
最後の2つの簡単なポインタ、
  • 量子化は微分不可能なステップであるため、エンドツーエンドのトレーニングを可能にするには、再構成エラーの勾配がデコーダーを介して逆伝播され、ストレートスルー勾配推定器を使用してエンコーダーに伝播されます(デコーダーから勾配をコピーしてエンコーダー)。
  • 再構築の損失に加えて、コードブックの損失とコミットメントの損失が使用されます。コードブックが失われると、選択したコードeがエンコーダーの出力に近づきます。コミットメントロスにより、エンコーダーの出力が選択したeに近いままになります。

離散潜在表現が機能するのはなぜですか?

Transformerの制御にある程度似ているVQ-VAE-2を使用して多様な高忠実度画像を生成する作業では、VQ-VAE2(VQ-VAEへの変更)を使用しますが、離散表現が機能するのはなぜですか?
これは、画像のJPEG非可逆圧縮に触発されています。JPEGエンコーディングは、知覚される画質を著しく変化させることなく、データの80%以上を削除します。第二に、ノイズの少ない生成モデルをトレーニングする方がうまくいく傾向があります。

Transformerの制御

図3: Transformerの制御のアーキテクチャ設計。
シーケンスの長さを小さく保ち、Transformerの表現力を活用するために、Transformerを使いこなす著者は、画像 xRH×W×3x \in R^{H \times W \times 3} を空間で表すことができるように、学習した表現の離散コードブック(VQ-VAEから着想を得た)を使用します。コードブックエントリのコレクション zqRh×w×nzz_q \in R^{h \times w \times n_z} ここで、 nzn_z はコードの次元です。
図3に示すように、著者はVQ-GANを使用しました。これは、元のVQ-VAEの変形であるため、識別器と永続的な損失を使用して、圧縮率を上げても良好な知覚品質を維持します。
これは、2段階のトレーニングアーキテクチャ設計です。
  • VQ-GANのトレーニングと、量子化されたコードブックの学習。
  • 量子化されたコードブックをTransformerへの順次入力として使用して、自己回帰Transformerをトレーニングします。

VQ-GANのトレーニング

VQ-GANは、パッチベースの弁別器Dを使用した既知の敵対的トレーニング手順を使用してトレーニングされます。最適な圧縮モデル Q={E,G,Z}Q^* = \{E^*, G^*, Z^*\} を見つけるための完全な目的は、次の式で与えられます。
Q=argminmaxE,G,Z Exp(x)[LVQ(E,G,Z)+λLGAN({E,G,Z},D)]Q^* = \underset{E, G, Z}{\operatorname{argminmax}} \space E_{x \approx p(x)} [L_{VQ}(E, G, Z) + \lambda L_{GAN}(\{E, G, Z\}, D)] 
ここで、 λ\lambda は適応重みです。
このトレーニング手順により、潜在コードを展開する際のシーケンスの長さが大幅に短縮され、強力なTransformerモデルの適用が可能になります。
独自のVQ-GANをトレーニングするために、この論文の公式GitHubリポジトリのクローンを作成し、依存関係をインストールできます。このリポジトリにはすでに重みとバイアスが組み込まれているため、W&Bダッシュボードで必要なすべてのメトリックを自動的に取得できます。データ準備ステップは、著者によって明確に言及されています。VQ-GANをトレーニングするには:
python main.py --base configs/faceshq_vqgan.yaml -t True --gpus 0,
選択したデータセットに応じて、異なる.yamlファイルを選択できます。

トレーニングTransformer

トレーニングされたエンコーダーとデコーダーを使用して、コードブックの観点から画像を表すことができます-それらのエンコーディングのインデックス。画像xでのこの量子化されたエンコーディングは、 zq=q(E(x))Rh×w×nzz_q = q(E(x)) \in R^{h \times w \times n_z}.で与えられます。このエンコーディングをコードブックからのインデックスのシーケンスe s{0,....,Z1}h×ws \in \{0,...., |Z|-1\}^{h \times w} として表現することにより、(zq​)ij​=zkのようなコードブックZ、sij​=k内の各コードをそのインデックスで置き換えることによって取得されます。これで、シーケンスsを対応するコードブックエントリにマッピングし直して、デコードされた画像 x~=G(zq)\tilde{x} = G(z_q)を取得できます。
このシーケンスを使用すると、画像生成を自己回帰次のインデックス予測として定式化できます。インデックスs<iが与えられると、Transformerは可能な次のインデックスの分布を予測することを学習します。つまり、p(sis<i)p(s_i | s_{<i}) は、完全な表現の尤度を p(s)=ip(sis<i)p(s) = \prod_i p(s_i | s_{<i}).。したがって、目的関数は、データ表現の対数尤度を最大化することです。これは可能性についての本当に素晴らしい記事です。


ユーザーが生成プロセスを制御できる場合は、任意の画像生成システムが役立ちます。画像の生成は、クラスラベル(上記のように)や部分的な画像などの追加情報を条件とすることができます。次に、この情報cが与えられた場合のシーケンスの尤度を学習することがタスクです:
p(sc)=ip(sis<i,c)p(s|c) = \prod_i p(s_i | s_{<i}, c)

トレーニング済みのVQ-GANを使用してmodel.params.first_stage_config.params.ckpt_path in configs/faceshq_transformer.yaml, の構成キー
python main.py --base configs/faceshq_transformer.yaml -t True --gpus 0,
チェックポイントパスを調整してから、次のコマンドを実行します。



高解像度画像の生成

シーケンス長の増加に伴って変圧器をトレーニングするための二次関数的に増加する計算コストは、シーケンス長に制限を課します。したがって、メガピクセル領域で画像を生成するには、作成者はパッチごとに作業し、画像をトリミングして、トレーニング中にsの長さを最大限に実行可能なサイズに制限する必要があります。
画像をサンプリングするために、Transformerはスライディングウィンドウ方式で使用されます。
図4: スライディングアテンションウィンドウ
VQ-GANは、データセットの統計がほぼ空間的に不変であるか、空間調整情報が利用可能である限り、利用可能なコンテキストが画像を忠実にモデル化するのに十分であることを保証します。
以下に示すメディアパネルで⚙をクリックし、スライダーを使用して、スライディングウィンドウを使用して高解像度画像がどのように合成されたかを視覚化します。


Run set
1


結果


Run set
1



Run set
1


参考文献及び結語

このレポートの目的は、読者がよりアクセスしやすいように、論文を要約することです。それが情報を伝えるための最良の方法だと思い、論文から必要な情報を少し引用しました。
ここにあなたのお役に立つと思われるいくつかの参考文献があります。
Iterate on AI agents and models faster. Try Weights & Biases today.