概要
Image2Image手法にドメイン適応を組みあわせた手法DAI2Iの論文を読んだので紹介します。本論文は、CVPR2020採択論文です。
本手法は、訓練時には存在しないドメインの画像を生成するための、ドメイン適応型のImage2Imageフレームワークです。
例えば、ドメイン$A, B$が人の顔、猫の顔、だとします。このとき、モデルの訓練時には真顔の人$A^-$と笑顔の人$A^+$の画像、真顔の猫$B^-$は存在しますが、笑顔の猫$B^+$は存在しないという状況を考えます。$-, +$は、各ドメインに対するサブドメインと呼びます。このとき、任意の真顔の猫$B^-$から笑顔の猫$B^+$を生成するような画像変換モデルを実現したい、というのが本手法の目的です。
論文の流れとしては、まずは$A, B$という2つのドメインがあって、$A$には全サブドメインの教師データがあるが、$B$には一部のサブドメインの教師データしかない、という場合の説明から始まります。その後、$B$という1つのドメインではなく、より多くのドメイン$X$に対しても適用できるようにしています。本記事では、まずは後者の一般化されたモデルを概観し、前者のなかで提案されていた損失関数について説明します。
書誌情報
- Chen, Ying-Cong, Xiaogang Xu, and Jiaya Jia. "Domain Adaptive Image-to-image Translation." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2020.
- http://openaccess.thecvf.com/content_CVPR_2020/papers/Chen_Domain_Adaptive_Image-to-Image_Translation_CVPR_2020_paper.pdf
手法の概観
本手法は、大きく4つのモジュールから成り立っています。以下、図と対応させながらご覧ください。
- (b)入力画像のドメイン$X$のスタイル特徴を抽出するStyle Net
- (c)入力画像からドメイン$A$への適応を行うためのAdapter$F_{X \mapsto A}(\cdot)$
- (d)ドメイン$A$内での学習済みImage2Imageモデル$F_{A^- \mapsto A^{+}}(\cdot)$
- (e)ドメイン$X$のスタイル特徴とImage2Imageの出力をかけ合わせて最終的な出力画像を合成するReconstructor$F_{A \mapsto X}(\cdot)$
Image2Imageのモデルは、上図ではStarGANが用いられていて、すでに写実的な人の顔画像（ドメイン$A$）で訓練済みで、固定されています。StyleNetによって抽出されたスタイル特徴は、ReconstructorのConv層に接続されたAdaINに入力されています。
直感的には、入力画像のコンテンツ情報をAdapterによってドメイン適応させ、そのうえでImage2Imageモデルによってコンテンツ情報を編集し、事前に抽出しておいたスタイル特徴も参照しながら元のドメインに戻す、という流れになっています。
各段階での入出力は以下のようになります。左から入力画像、Adapterを通してドメイン適応した画像、Image2Imageを通してサブドメイン変換を施した画像、最後のReconstructorを経た出力画像、となっています。
損失関数
上記のような枠組みを実現するために、様々な損失関数が使われており、最終的には以下のような損失関数が作られています。
\begin{array}{l}
\mathcal{L}_{F_{X \mapsto A}}=\lambda_{G A N} \mathcal{L}_{G A N}+\mathcal{L}_{A D A}+\mathcal{L}_{r e c}+\mathcal{L}_{P A} \\
\mathcal{L}_{F_{A \mapsto X}}=\mathcal{L}_{r e c}+\mathcal{L}_{P A}
\end{array}
$\mathcal{L}_{F_{X \mapsto A}}, \mathcal{L}_{F_{A \mapsto X}}$はそれぞれAdapter$F_{X \mapsto A}$、Reconstructor（とStyleNet）$F_{A \mapsto X}$に関する損失という意味です。それぞれの中にある細かい損失関数がなぜ必要なのかを見ていきましょう。
Adversarial Loss
まず、本手法はGANの一種なので、当然Discriminatorによる敵対的損失が存在します。LSGANのような損失関数を考えます。ここで$\hat{A}$はAdapter$F_{X \mapsto A}$によって生成された画像です。この損失によって、Adapterの出力結果がドメイン$A$のサブドメインである$A^-$とみなせるように訓練されます。
\begin{array}{l}
\min _{D} \mathcal{L}_{G A N_{D}}=\mathbb{E}\left(\|D(\hat{A})\|_{2}\right)+\mathbb{E}\left(\|D(A)-1\|_{2}\right) \\
\min _{F_{X \mapsto A}} \mathcal{L}_{G A N}=\mathbb{E}\left(\|D(\hat{A})-1\|_{2}\right)
\end{array}
Reconstructorのほうには、この損失は使われません。GANの訓練には、合成画像以外にも真の画像が必要になります。しかし、ドメイン$X$には、サブドメイン$X^-$は存在していても、別のサブドメイン$X^+$は存在しません。そのため、真の画像として$X^-$しか与えることができなくなります。これは、Reconstructorの出力が必ず$X^-$になるように促進してしまうため、$X^+$も出力できるようにしたいという目的を阻害してしまいます。Reconstructorの訓練のためには、もう少し回り道が必要です。
Adaptation Loss
Adapterの訓練がGANに依存すると、必ずしもImage2Imageの入力として最適な出力にならない可能性があります。この問題を克服するために、Image2Imageの出力結果が$A^-$よりは$A^+$に近くなることを保証する損失を加えます。これによってAdapterは、Image2Imageの入力として適切な画像を出力できるように訓練されます。
\min _{F_{X \rightarrow A}} \mathcal{L}_{A D A}=\mathbb{E}\left[-\log \left(\mathcal{C}\left(\hat{A}^{+}\right)\right)\right]
ここで、$\hat{A}^{+}$はAdapterとImage2Imageを通した出力結果です。$\mathcal{C}$は事前に訓練された$A^-, A^+$の分類器です。この分類器とImage2Imageは訓練の対象外なので固定されています。
Reconstruction Loss
次に、サイクル変換を施した時の損失を考えます。AdapterとReconstructorを直接接続すると、入力画像と出力画像が同じ画像になることが期待されます。
\min _{F_{X \rightarrow A}, F_{A \rightarrow X}} \mathcal{L}_{r e c}=\mathbb{E}\left[\left\|F_{A \mapsto X}(F_{X \mapsto A}(X), {Style}(X))-X\right\|_{1}\right]
この損失は、AdapterにもReconstructor（およびStyleNet）の両方にかかわるので、いずれの訓練にも使用されます。
Perceptual Analogy Loss
一見すると$\mathcal{L}_{r e c}$で目的は達成したかのように思えますが、実際には、この損失が保証するのは、$F_{A \rightarrow X}(\hat{A}) = X$という関係だけで、$F_{A \rightarrow X}(\hat{A^+}) = X^+$となることを保証しません。$X^+$はモデルの訓練中に存在しないため、このことを保証するモデルの訓練は非常に難しい課題になります。
ここで、Image Analogyという考え方を導入します。これはWord2Vecなどでよくある、Queen - King = Princess - Princeという関係が画像の特徴量でもある程度成り立つという考え方です。
具体的には以下のような損失となります。
\min _{F_{X \mapsto A}, F_{A \rightarrow X}} \mathcal{L}_{P A}=\mathbb{E}\left[\left\|\mathcal{V}_{X}-\alpha \mathcal{V}_{A}\right\|_{1}\right]
\mathcal{V}_{X}=\Phi\left(\hat{X}^{+}\right)-\Phi(X)
\mathcal{V}_{A}=\Phi\left(\hat{A}^{+}\right)-\Phi(\hat{A})
ここで、$\Psi$は学習済みで固定されたVGG19のReLU3_1, ReLU4_1, ReLU5_1の特徴マップを表します。各画像については、$X$をもとにAdapterによって得られる$\hat{A}$、それをImage2Imageによって変換した$\hat{A}^{+}$、そして最終出力結果$\hat{X}^{+}$という表記になっています。
要は、サブドメイン間での画像特徴の差異が、ドメイン$A$でも$X$でも似た方向になっていてほしいという考え方です。$\alpha$は、ドメイン$A$での差異$\mathcal{V}_{A}$をドメイン$X$でどの程度強調（または抑制）すべきかをコントロールするスカラー値です。
Calibration of Domain Shift
上記の$\mathcal{L}_{P A}$は一見うまくいきそうですが、2つのドメイン間の分布シフトによって、そのままでは比較が困難です。ここでは、domain-specific batch normalizationを用いて、この問題を回避します。
$\Phi(X)$の$i$番目のチャネルを$\Phi(X)_{i}$とし、その移動平均を$\sigma_{i}^{X}$とします。実際に$\mathcal{V}_{X}$を計算するときは、$\Phi(X)_{i}$を$\sigma_{i}^{X}$で正規化した$\Psi(X)_{i}$を使います。同様に、以下のようにすべての特徴量が正規化されます。これらを用いて、$\mathcal{V}_{A}, \mathcal{V}_{X}$を計算します。
\begin{array}{l}
\Psi(X)_{i}=\frac{\Phi(X)_{i}}{\sigma_{i}^{X}}, \Psi\left(\hat{X}^{+}\right)_{i}=\frac{\Phi\left(\hat{X}^{+}\right)_{i}}{\sigma_{i}^{X}} \\
\Psi(\hat{A})_{i}=\frac{\Phi(\hat{A})_{i}}{\sigma_{i}^{A}}, \Psi\left(\hat{A}^{+}\right)_{i}=\frac{\Phi\left(\hat{A}^{+}\right)_{i}}{\sigma_{i}^{A}}
\end{array}
結果の確認
StarGANと本手法の出力結果を比較してみましょう。
StarGANはリアルな人の顔画像で訓練されているので、入力画像の色を無視して、肌色に着色してしまう傾向があります。また、サブドメイン（Happy, Angry, Sad,...）などに対して、良い変換結果になっているとは言えません。
DAI2Iを用いた変換結果を見てみると、入力画像の色あいを保持したままの出力結果が得られています。また、歯を見せて笑ってる猫の画像や口を丸く開けて驚いている猫の画像など、写実的でありながら現実には絶対存在しない出力結果が得られています。
また、Image2Imageモデルの部分を新視点画像合成のモデルに変更することで、以下のようにスケッチや油絵を別の視点から見たときにどうなるのか、という出力結果も得られています。
まとめ
ドメイン適応とImage2Imageモデルを組み合わせたDAI2Iについて紹介しました。多様な損失が無駄なく適用されており、出力結果もあまり見たことがないようなものになっていて面白い研究だと思います。Image2Imageの部分をいろいろと変えることで、いままではできなかった様々な画像生成ができるように思います。