TransformerでGAN Inversion！【Style Transformer】
3つの要点
✔️ Transformerを用いたGAN Inversion手法
✔️ 再構成品質、編集能力、モデルサイズにおいて既存手法を上回る性能
✔️ 参照画像を用いた編集も可能に
Style Transformer for Image Inversion and Editing
written by Xueqi Hu, Qiusheng Huang, Zhengyi Shi, Siyuan Li, Changxin Gao, Li Sun, Qingli Li
[Submitted on 4 Dec 2021 (v1), last revised 29 Mar 2022 (this version, v3)]
Comments: Accepted by CVPR 2022
Subjects: Computer Vision and Pattern Recognition (cs.CV); Artificial Intelligence (cs.AI)
code：
本記事で使用している画像は論文中のもの、紹介スライドのもの、またはそれを参考に作成したものを使用しております。
はじめに
最近では、StyleGANが高解像度画像を生成できるようになり、実写画像の様々な編集タスクに応用する研究が盛んになっています。実写画像を編集するためにはGAN Inversionと言われる手法によって実写画像からStyleGANの潜在変数を求める必要があります。GAN Inversionで潜在変数を求める際に重要な点が主に2つあります。
- 元の画像を忠実に再構成できること（再構成能力）
- 元のアイデンティティや詳細を保ったまま、編集したい属性のみを操作できること（編集能力）
これらを同時に満たすことは難しい問題です。
埋め込みを行うStyleGANの潜在空間にはZ空間やW空間、W+空間などいくつかの候補があります。既存研究ではこれらの候補の選択が重要であると言われています。Z空間やW空間は単一の512次元のベクトルで表現されますが、W+空間では18個のベクトルで表現され、それぞれが512次元です。このため、W+空間では画像の細部を表現する能力が優れており再構成の能力が高いですが、1つの属性に対して様々な次元が絡み合っていることが多いため、独立した属性編集が難しいです。
再構成能力と編集能力を同時に改善するために、本論文ではTransformerを用いたGAN Inversion手法「StyleTransformer」を提案しました。この手法ではW+空間への埋め込みを行っています。Transformerは自然言語処理を始めとして様々な領域で良好な結果をもたらしています。
上の画像はStyleTransformerを用いて再構成画像と編集画像を出力した結果です。再構成の品質を高く保ちながら、編集もうまくできていることがわかります。
また、この手法では参照画像を用意して、その参照画像のある属性にターゲット画像の属性を転写することが可能となっています。
提案手法
以下の図はStyle Transformerのフレームワークの概要を示しています。
まず、入力画像はエンコーダEによって複数の解像度の画像特徴量F1~F3が生成されます。MLPから出力されるN個の異なるクエリがTransformer Blockを介して、これらの特徴にアクセスし、生成器の潜在変数wへ徐々に更新されていきます。
エンコーダE、MLP、Transformer Block、そして初期値znの全てのパラメータを学習させて、最適な潜在変数wを出力できるようにします。
Transformer Block
上の画像にはTransformer Blockの構造が示されています。
構造は従来のTransformerに似ており、Multi-Head Self-AttentionとCross-Attentionを含む設計となっています。また、残差接続や正規化、FFNモジュールなども従来のTransformerの構造を参考にしています。
ここで、一般的なTransformerデコーダでは入力のクエリトークンをランダムに初期化して、パラメータとして保持することが多いです。しかし、W空間の分布は複雑であり、ガウス分布とはかなり異なることから一般的な方法で学習させてもうまくいきません。
そこで、StyleGANの事前学習済みのMLPを利用して、潜在変数znをwnにマッピングすることでW空間から大きく外れないようにする工夫をしています。また、事前学習済みMLPは学習時に固定はせずファインチューニングします。
Multi-Head Self-Attention
従来のTransformerで使われている計算と同様です。
Self-Attentionでは入力されたクエリの任意のペア間の関係性を見つけ、それらを結びつけるように学習が進みます。これによって、任意の潜在変数wn間の関係を捉えています。
Multi-Head Cross-Attention
Self-Attentionだけでは潜在変数間の関係しか見ていないので、画像の特徴については一切関与していません。
そこでMulti-Head Cross-Attentionを用いて、異なる解像度の画像特徴量F1~F3から情報を得ます。具体的には、キーとバリューを画像特徴量からのものを、クエリをSelf-Attentionの結果を使って計算します。
学習
学習時は、StyleGANの生成器Gは固定し、それ以外の全てのパラメータを調整します。
損失関数についてはpSpというGAN Inversion手法と同様のものを使用しています。詳しくはpSpの論文を参照してください。
Style Transformerを用いた画像編集
最初に述べたように、GAN Inversionでは再構成の性能だけでなく、編集能力が高いことも重要です。
Style Transformerではラベルによる属性編集だけでなく、参照画像を用いて特定の領域を編集することも可能です。
参照画像による編集
参照画像による編集をするために新たに1つのTransformer Blockを学習させます。
まず、W+の潜在変数を入力として受け取り、属性ごとの埋め込み特徴量とラベルを出力する属性分類器Cを学習させます。
次に上の画像のようにして、新たなTransformer Blockを学習させます。まず、参照画像とターゲット画像をそれぞれStyleTransformerによって潜在空間W+に埋め込みます。そして、Transformer Blockのバリューとキーには参照画像の潜在変数を、クエリにはターゲットの画像の潜在変数を入力し、新たな潜在変数weを出力します。
編集したい属性を参照画像の属性に寄せ、それ以外はターゲットの画像に寄せるように損失関数を計算することで、所望の編集がされた潜在変数weが生成できるようになります。
実験結果
定性評価
まず、既存手法と再構成結果および属性編集をおこなった結果が以下の画像に示されています。
２列目が再構成結果となっています。pSpは再構成結果の品質が高いGAN Inversion手法ですが、提案手法の結果はpSpに劣らない結果に見えます。また、e4eは編集能力が高いGAN Inversion手法であり、再構成結果の品質はそれほど高くないように見えます。
編集結果を見ると、提案手法は編集能力が高いとされるe4eよりも、属性ごとにうまく分離できているように見えます。
顔画像だけでなく車の画像でも実験しており、再構成の品質および編集結果の品質のどちらも高いことがわかります。
定量評価
以下の表では既存手法と定量的に比較した結果が示されています。
再構成結果については、ピクセル間の類似度（MSE、LPIPS）や実画像と生成画像の分布間の距離（FID、SWD）で評価していますが、どの指標においても提案手法が最も高い結果となっています。編集結果についてもFIDとSWDで評価していますが、提案手法が最も高いです。
また、結果画像の品質が高いだけでなく、モデルの大きさや推論時間などの観点でも他のモデルより高性能であり、ここに挙げられているすべての指標においてpSpとe4eを上回っています。
参照画像による編集
参照画像による編集をした結果画像が上の画像です。多様性は少し低くなってしまっていますが、参照画像の特定の属性だけを元の画像に反映させることができています。
まとめ
今回はCVPR2022に採択されたStyle Transformerを紹介しました。
GAN Inversionにおいて重要な再構成能力と編集能力のどちら点においても高い性能を持つだけでなく、モデルサイズが小さく推論時間も短い手法となっており、今までの手法をさまざまな観点で上回っています。
GAN Inversionの技術が進歩すれば、高速で高品質な画像編集が行えるようになるため今後のGAN Inversion手法にも注目です。
この記事に関するカテゴリー
