Pivotal Tuning

一般的な顔編集では、固定した重みを用いたStyleGANで、素晴らしい品質の生成を達成していました。しかし、学習に使用していないデータに関する潜在空間の分布は、編集能力に欠けています。そこで、Pivotal TuningによりStyleGANをチューニングすることで、対象の画像の潜在コードを編集能力の高い潜在コードに遷移させ、高性能な編集性能を達成します。

以下、詳細な手法に関して説明します。

1. Inversion

まず、編集能力の高い潜在コードを探索するInversionを行う必要があります。探索にはLearned Perceptual Image Patch Similarity (LPIPS)による知覚損失を使用します。これは、Perceptual path lengthと言う潜在変数の変化の滑らかさを表す指標でよく使用されます。LPIPS損失($\mathcal{L}_{LPIPS}$)は、VGGなどのモデルにより推定した画像の特徴量のL2距離により計算されます。LPIPS損失を含む以下の式を用いて潜在コード$w$とノイズ$n$を最適化することで、対象とする潜在コード$w_p$を計算できます。

$$w_p, n = \mathrm{argmin}_{w,n} \mathcal{L}_{LPIPS}(x, G(w, n; \theta)) + \lambda_n \mathcal{L}_n(n)$$

$G$は$\theta$を重みとした生成器による生成画像で、$x$は対象となる実際の画像です。$\lambda_n$, $\mathcal{L}_n$はノイズに関する正則化の割合と損失です。ノイズ正則化は、ノイズ情報の影響を防ぎinversionの性能を改善するために使用されています。

2. Pitovatl tuning

pivotal tuningは、Inversionにより求めた潜在コード$w_p$を用いて、生成器$G$の重み$\theta^{\star}$を学習します。損失は、LPIPS損失と実画像と生成画像のL2誤差$\mathcal{L}_{L2}$を用いて、以下の式になります。

$$\mathcal{L}_{pt} = \mathcal{L}_{LPIPS}(x, G(w_p; \theta^{\star})) + \lambda_{L2} \mathcal{L}_{L2}(x, G(w_p; \theta^{\star}))$$

しかし、この損失では、大幅に潜在コードを修正し、逆に視覚的な品質が損なわれることがあります。そこで、局所的な修正に留めるための損失を追加する必要があります。論文では、まず、ガウス分布に則ったランダムなベクトル$z$とStyleGANの射影ネットワーク$f$を用いて計算される潜在コード$w_z = f(z)$を用いて、pivotal tuningの潜在コード$w_p$をわずかにシフトさせた潜在コード$w_r$を計算しています。

$$w_r = w_p + \alpha \frac{w_z - w_p}{||w_z - w_p||_2}$$

そして、この潜在コードを用いてオリジナルの生成器とチューニング中の生成器の画像を比較しています。

$$\mathcal{L}_{R} = \mathcal{L}_{LPIPS}(x_r, x_r^{\star}) + \lambda_{L2} \mathcal{L}_{L2}(x_r, x_r^{\star})$$

ここでは、オリジナル生成器による画像を$x_r = G(w_r; \theta)$、チューニング中の生成器による画像を$x_r^{\star} = G(w_r; \theta^{\star})$としています。 この損失により、オリジナル生成器とチューニング中の生成器において、わずかに遷移させた潜在コードを入力として生成した顔編集画像の誤差を小さくするように学習できます。つまり、顔編集を行った際の潜在コードがオリジナルの生成器から大きくズレないように学習しています。これにより、視覚的な品質の劣化を防いでいます。

上記をまとめて、Pivotal tuningにより学習される重み$\theta^{\star}$は、

$$\theta^{\star} = \mathrm{argmin}_{\theta^{\star}} \mathcal{L}_{pt} + \lambda_R \mathcal{L}_{R}$$

により最適化されます。

まとめ

本記事では、学習済みのStyleGANを対象画像に対してわずかにチューニングすることで、ロバストな編集が可能にするPivotal tuningについて解説しました。学習後に、対象の画像に対してチューニングを行うという処理は、挙動が分かり辛いため実際の機械学習システムに取り入れるのは難しいと思っていました。しかし、チューニングの度合いを制限する処理に関しても提案されており、実際、かなり精度が向上していたこともあり、未知のデータに対する対処法の一つとしてシステムに取り入れても良いかもしれないと思いました。