こんにちは、@tsubota-kougaです。はじめてのQiitaです。
今日は、Spectral Normalizationの論文の紹介をしたいと思います。
Spectral Normalizationって、あれです、GANに使うと安定して学習できるようになるっていうあれです。
今まで無免許でSpectral Normalizationを使ってきたので、さすがにまずいと思いまして、今回、Spectral Normalizationの論文を読んだのでまとめたいと思います。
「Lipschitz連続とSpectral Normalization」ということで、Lipschitz連続とは何か? Lipshitz連続がどうして必要なのか? というところを見ていこうと思います。
間違っている箇所や誤字脱字などがありましたら、気軽にコメントやプルリクお願いします。あと良かったら「いいねボタン」押してもらえると嬉しいです(おじさんなのでLGTMとかワカラナイ)。
原著論文: Spectral Normalization
Spectral Normalizationの嬉しいところ
- GANの学習を安定にすることができる
- Spectral Normalizationを使っても、使わないときとほとんど変わらない計算速度を実現
- Spectral Normalizationに必要なハイパーパラメータは1つだけ、それもほとんどの場合は弄らなくていい
- 今までの実装に、ほんのひと足し
Lipchitz連続とは
数学的に書くと、関数$f: \mathbb{R}^n \rightarrow \mathbb{R}$とするとき、$\forall x, y \in \mathbb{R}^n$に対して、$\kappa > 0$が存在して、
\|f(x) - f(y)\| \le \kappa |x - y|
を満たすとき、$f$はLipschitz連続であるという。
もっと一般の距離空間に関して定義できるので、気になる方はWikipediaを見ましょう。
Spectral Normalizationとは
DiscriminatorにLipschitz連続を課すために作られた、正規化層です。
詳しい定義は、後ほどー。
ここからは、少し寄り道をして、どうしてLispchitz連続をDiscriminatorに課すといいのか、それを見ていきます。
なぜLipschitz連続が必要になるのか、の前に
Generatorの目的関数
\begin{align}
\min_{f \in \mathcal{F}} J_D &= \mathbb{E}_{z \sim \mathbb{P}_z}[\phi(f(g(z)))] + \mathbb{E}_{x \sim \mathbb{P}_r}[\varphi(f(x))] \\
\min_{g \in \mathcal{G}} J_G &= \mathbb{E}_{z \sim \mathbb{P}_z}[\psi(f(g(z)))]
\end{align}
ここで、$f$はDiscriminatorの関数、$g$はGeneratorの関数。
$\mathbb{P}_z$はノイズの分布、$\mathbb{P}_r$は実画像の分布。$\phi,\varphi,\psi: \mathbb{R} \rightarrow \mathbb{R}$となる損失関数。
$\mathcal{F},\mathcal{G}$はそれぞれ、Discriminator、Generatorの関数の空間です。Discriminatorの関数でパラメータが少しずつ違った関数が入った集まりを$\mathcal{F}$と言っていると思っていいと思います。$\mathcal{G}$も同じくです。
GeneratorがDiscriminatorから受け取る勾配は、先ほどの$\min_{g \in \mathcal{G}} J_G$の式を微分してあげればいいので
\nabla_x J_G(x):=\nabla_x \psi(f(x))=\nabla_{f(x)}\psi(f(x)) \cdot \nabla_x f(x)
となります。式の最後は合成関数の微分ですね。
右辺のそれぞれの項は、以下のように説明できます。
- $\nabla_{f(x)}\psi(f(x))$は勾配の大きさを表すScalar
- $\nabla_x f(x)$は勾配方向を表すVector
GANの学習が不安定になる原因として、主に以下の2つがあります。
- Gradient Vanishing (勾配消失)
- Gradient Uninformativeness
それぞれ、どんな問題なのでしょうか。
Gradient Vanishing
勾配の大きさを表す$\nabla_{f(x)}\psi(f(x))$が小さくなり、学習が進まなくなること、これをGradient Vanishing(勾配消失)といいます。
この問題は、損失関数の定義の仕方によって解決できます。例えば、LSGANはこの問題を解決するために提案されました。ただし、次に説明する、Gradient Uninformativenessは考慮に入れられていません。
Gradient Uninformativeness
勾配方向を表す$\nabla_x f(x)$が$\mathbb{P}_r$に関する有益な勾配情報を持たなくなること、これをGradient Uninformativenessといいます。これによって、生成画像の分布が実画像の分布$\mathbb{P}_r$に収束しなくなります。
Gradient Uninformativenessの方は、Lipschitz連続をDiscriminatorに課すことで解決することができます。
なぜLipschitz連続が必要になるのか
それは、こんな命題が成り立つからなのです(参考論文の定理4)。
$f$が微分可能で、そのLipschitz定数が$\kappa$であるとする。
$\forall x,y(x \ne y)$かつ$f(y)-f(x)=\kappa ||y-x||$のとき、$x_t=tx+(1-t)y,0 \le t \le 1$とすると、
$\nabla_{x_t}f(x_t)=\kappa \cdot \frac{||y-x||}{y-x}$
これのイメージは以下のような感じ。
生成画像と実画像の間の点で、ちゃんとどの方向に進むべきかが定まるということですね。
この定理の仮定にある$f(y)-f(x)=\kappa ||y-x||$は何でしょう?
$f$がLipschitz連続なら、$||f(x) - f(y)|| \le \kappa ||x - y||$なはずです。
ここでは、結果だけ使います(疑問投げかけておいてごめんなさい)。詳しくは、参考論文のLipschitz Adversarial Netsの定理1~3を見てください。
損失関数に以下の条件を課します。これを課すことによって、生成画像$x$と実画像$y$に一対一の関係が生まれ、$f(y)-f(x)=\kappa |y-x|$が保証されます。
\left\{
\begin{array}{}
\phi'(x) > 0, \varphi'(x) < 0 \\
\phi''(x) \ge 0, \varphi''(x) \ge 0 \\
\exists a, \phi'(a) + \varphi'(a) = 0
\end{array}
\right.
例えば、OriginalのGANの式であれば、
\begin{align}
\min_{f \in \mathcal(F)}& \mathbb{E}_{z \in \mathcal{P}_z}[-\log(1-f(g(z)))] + \mathbb{E}_{x \in \mathcal{P}_r}[-\log(f(x))] \\
\min_{g \in \mathcal{G}}& \mathbb{E}_{z \in \mathcal{P}_z}[\log(1-f(g(z)))]
\end{align}
なので、$\varphi(x)=-\log(1-x), \phi(x)=-\log(x)$となります。この場合、$\varphi'(x)=\frac{1}{1-x},\phi'(x)=-\frac{1}{x}$となって、先ほどの3つの式のうちの1つ目、$\varphi'(x) > 0, \phi'(x) < 0$を満たしません。
次に、Non Saturating Lossの場合を考えてみます。
\begin{align}
\min_{f \in \mathcal(F)}& \mathbb{E}_{z \in \mathcal{P}_z}[-\log(sigmoid(-f(g(z))))] + \mathbb{E}_{x \in \mathcal{P}_r}[-\log(sigmiod(f(x)))] \\
\min_{g \in \mathcal{G}}& \mathbb{E}_{z \in \mathcal{P}_z}[\log(sigmoid(-f(g(z))))]
\end{align}
ただし、$sigmoid(x)=\frac{1}{1-e^{-x}}$でsigmoid関数とします。
このとき、$\varphi(x)=-\log(sigmoid(x)),\phi(x)=-\log(sigmoid(-x))$
これは計算多いので省略しますが、上の条件を満たします。
上の条件を満たす損失関数の作り方ですが、$\phi$を真に増加する関数、導関数が非減少関数として、$\phi(x)=\varphi(-x)$とおいてあげれば、簡単にできます。
上の条件を満たせば、あとはLipschitz連続をDiscriminatorに課すだけで、Gradient Uninformativenessを解決できます。
これで、Spectral Normalizationを導入する動機ができましたね!
Spectral NormとLipschitz Norm
ここからは、Spectral Normalizationの論文に戻りましょう。
まずは、行列$A$のSpectral Norm (=作用素ノルム) $\sigma(A)$の定義です。
\sigma(A) :=\max_{h:h \ne 0} \frac{{\|Ah\|}_2}{{\|h\|}_2}
=\max_{{\|h\|}_2 \le 1}{\|Ah\|}_2
=\max_{{\|h\|}_2 = 1}{\|Ah\|}_2
このSpectral Normは、行列$A$の最大特異値と一致します。
特異値分解をざっと説明
行列$P P^T$の固有値を大きい順に並べたものを、$\lambda_1,\lambda_2,\ldots,\lambda_n$とし、対応する$n$次元単位固有ベクトルの正規直交基底を$u_1,u_2,\ldots,u_n$とする。
この固有ベクトルを並べた$n \times n$行列を$U = (u_1, u_2, \ldots, u_n)$とする。
同様に、行列$P^T P$でも行い、$\lambda_1,\lambda_2,\ldots,\lambda_N$として、対応する固有ベクトルを$v_1,v_2,\ldots,v_N$とする。
この固有ベクトルを並べた$N \times N$行列を$V = (v_1, v_2, \ldots, v_N)$とする。
このとき、
P= U \left(
\begin{array}{cccc}
\sqrt{\lambda_1} & & &\\
& \ddots & &\\
& & \sqrt{\lambda_r} &\\
& & &
\end{array}
\right) V^T
と書ける。ただし、$r:=rank P$。
$U,V$の行列のベクトルを$r$本にしたものを$U_r,V_r$とすると、$P=U_r \Sigma V_r$。ただし、$\Sigma$は先程の行列の左上の対角行列。これを特異値分解という。
次に、Lipschitz Normの定義です。
\frac{{\|f(x) - f(x')\|}_2}{{\|x-x'\|}_2} \le M
これを満たす$M$の中で最小のものをLipschitz Norm ${|f|}_{Lip}$と定義します。
もっと簡潔に、
{\|f\|}_{Lip} := \sup \frac{{\|f(x) - f(x')\|}_2}{{\|x-x'\|}_2}
とも書けます。
さて、線形層(Pytorchでいうところのnn.Linearですね)のLipschitz Normを考えてみます。ある線形写像を$g: h_{in} \mapsto h_{out}$で、$g(h)=Wh$とします。
$x \ne x'$を仮定すると、
{\|g\|}_{Lip} = \sup_{x,x'} \frac{{\|g(x) - g(x')\|}_2}{{\|x - x'\|}_2}
= \sup_{x,x'} \frac{{\|g(x - x')\|}_2}{{\|x - x'\|}_2} \\
= \sup_{h \ne 0} \frac{{\|g(h)\|}_2}{{\|h\|}_2}
= \sup_{h \ne 0} \frac{{\|Wh\|}_2}{{\|h\|}_2}
= \sigma(W)
と書くことができます。つまり、線形層の場合は、Lipschitz NormとSpectral Normが一致して、最大特異値になります。
Spectral Normalizationの定義
Spectral Normalizationは、各層の重みをSpectral Normで正規化します。
つまり、
\bar{W}_{SN}(W) := \frac{W}{\sigma(W)}
のように各層の重みを正規化します。
ちなみに、このSpectral Normalizationは、nn.BatchNormのような、1つで完結するものではなく、nn.utils.weight_normのような感じで層と一体化する形の正規化です。
まあ、定義見てもらえば、層の重みに依存しているのが分かるので、それはそうって感じですが。
多層ニューラルネットのDiscriminatorで考えてみる
Discriminatorの関数を、
f(x,\theta)=W^{L+1}a_L(W^L(a_{L-1}(W^{L-1}(\cdots a_1(W_1 x)))))
とします。ただし、$\theta:={W^1,\cdots,W^L,W^{L+1}}$を学習パラメータで、$W^l \in \mathbb{R}^{d_l \times d_{l-1}}, W^{L+1} \in \mathbb{R}^{1 \times d_L}$とします。
つまり、$W^l$は$d_{l-1}$次元から$d_l$次元への線形写像であるような行列、ということですね。また、途中に挟まれている$a_l$は活性化関数です。
また、${||a_l||}_{Lip}=1$と仮定します。
この仮定は、妥当なもので、ReLUやLeaky ReLUなどの活性化関数が、${||a_l||}_{Lip}=1$です。
ノルムには劣乗法性${||g_1 \circ g_2||}_{Lip} \le {||g_1||}_{Lip}{||g_2||}_{Lip}$がありますから、
{\|f\|}_{Lip} \le {\|(h_L \mapsto W^{L+1}h_L)\|}_{Lip} \cdot {\|a_L\|}_{Lip} \cdot {\|(h_{L-1} \mapsto W^L h_{L-1})\|}_{Lip} \cdots {\|a_1\|}_{Lip} \cdot {\|(h_0 \mapsto W^1 h_0)\|}_{Lip} \\
= \prod_{l=1}^{L+1} \|(h_{l-1} \mapsto W^l h_{l-1}\|_{Lip})
= \prod_{l=1}^{L+1} \sigma(W^l)
と書くことができます。最後では、線形写像の場合Lipschitz NormとSpectral Normが等しい、を使っています。
このことから何が言えるか、ですが。
Spectral Normalizationは、各層の重み$W^l$を先ほど定義した$\bar{W}_{SN}(W) := \frac{W}{\sigma(W)}$で正規化します。
よって、Lipschitz Normalizationを課した後のDiscriminatorでは、
\|f\| \le \prod_{l=1}^{L+1} \sigma(\bar{W}_{SN}(W^l))
= \prod_{l=1}^{L+1} 1 = 1
となり、Discriminatorに1-Lipschitz制約を課すことができました。
これで、Discriminatorも立派なLipschitz連続関数です。やったね!
Spectral Normalizationをしたあとの層の勾配を見てみる
\frac{\partial \bar{W}_{SN}}{\partial W_{ij}}=\frac{1}{\sigma(W)}E_{ij}-\frac{1}{{\sigma(W)}^2}\frac{\partial \sigma(W)}{\partial W_{ij}}W \\
=\frac{1}{\sigma(W)}E_{ij}-\frac{{[u_1 {v_1}^T]}_{ij}}{{\sigma(W)}^2}W \\
=\frac{1}{\sigma(W)}(E_{ij}-{[u_1 {v_1}^T]}_{ij}\bar{W}_{SN})
ただし、$E_{ij}$は$(i,j)$成分だけが$1$で他は$0$みたいな行列、$u_1,v_1$はそれぞれ左最大特異値、右最大特異値。
また、ある層の重み$W$での勾配は、$\bar{W}_{SN}$への入力を$h$とすれば、
\frac{\partial V(G, D)}{\partial W}
=\frac{1}{\sigma(W)}(\hat{E}[\delta h^T]-(\hat{E}[\delta^T \bar{W}_{SN}h])u_1 {v_1}^T) \\
=\frac{1}{\sigma(W)}(\hat{E}[\delta h^T]-\lambda u_1 {v_1}^T)
ただし、
\begin{align}
\delta := (\frac{\partial V(G, D)}{\partial (\bar{W}_{SN}h)})^T \\
\lambda:=\hat{E}[\delta^T (\bar{W}_{SN}h)]
\end{align}
とおいてあります。また、期待値が$\hat{E}$となっているのは、minibatchでの平均を表しているからです。
ここで注目してほしいのは、$\frac{\partial V(G, D)}{\partial W}$の式変形の最後です。第1項目の$\frac{1}{\sigma(W)}\hat{E}[\delta h^T]$はSpectral Normalizationなしのときの勾配と同じです。
このように考えると、Spectral NormalizationなしのDiscriminatorの勾配に正則化項がついたように見えますよね。ええ、見えないとは言わせません。
で、どのようなときにペナルティを与えているかというと、左最大特異値ベクトルと右最大特異値ベクトルが同じ方向を向いた時です。ユークリッドの意味での内積を表しているので、2つが同じ方向を向いたときに、$u_1 {v_1}^T$は大きくなります。
$\lambda$は変化するので、Spectral NormalizationはAdaptiveに最大特異値にペナルティを与えていると言えます。
Spectral Normを速く求める
今まで、Spectral Normをなんとなく使ってきましたが、この値は簡単に求まるのでしょうか。
Spectral Normは最大特異値なので、特異値分解をしないといけません。しかし、特異値分解を各層において学習イテレーションごとに計算すると時間がとてもかかります。
そこで、冪乗法という近似計算を使います。
冪乗法は最大固有値、最大固有ベクトルを求める手法です。
冪乗法
最大固有値、最大固有ベクトルを求めます。
正方行列$M$について、固有値を大きい順に並べたものを、$\lambda_1,\lambda_2,\ldots,\lambda_n$として、対応する$n$次元単位固有ベクトルの正規直交基底を$u_1,u_2,\ldots,u_n$とします。
このとき、$\forall x \in \mathbb{R}^n$は、固有ベクトルを使って表すと、$x = \sum_{i=0}^{n} {c_i u_i}$と書ける。
よって、$M x=\sum_{i=0}^{n} {c_i \lambda_i u_i}$、$M^k x=\sum_{i=0}^{n} {c_i {\lambda_i}^k x_i}$。
こうすると、大きい固有値の成分だけが残り、最大固有ベクトルが求まります。
これを$U, V$で使えば、最大特異値も分かります。詳しくは原著のAppendix.AにあるAlgorithm 1を参照してください。
この手法では、最大固有値が飛びぬけて大きいときに速く収束し、固有値にばらつきがないときには収束が遅くなります。
学習にかかる時間は上の通りで、WGANgpは結構時間がかかっていますが、Spectral Normalizationはほとんど使わないときと大差ないんじゃないかというくらいまで速いですね。
Spectral Normalizationと他の手法を比べてみる
- Weight Normalization
- Weight Clipping
- Orthonormal Regularization
- Gradient Penalty
比べる手法はこの4つです。
まずは、Weight Normalizationです。
Weight Normalizationは重み行列の行ベクトルをそれぞれの$l_2$ノルムで正規化するものです。
これは、$\sigma_t(A)$を$t$番目の特異値とすれば、
{\sigma_1(\bar{W}_{WN})}^2 + {\sigma_2(\bar{W}_{WN})}^2 + \cdots + {\sigma_T(\bar{W}_{WN})}^2 = d_o, where\;T = \min(d_i, d_o)
となるように正規化しているのと同じことです(右特異行列と左特異行列が正規直交基底を並べたものということが分かれば分かると思います)。
この正規化、実はとても強い正規化です。
$d_i \times d_o$次元のWeight Normalizationされた行列$\bar{W}_{WN}$があったときに、単位ベクトル$h$として、${|\bar{W}_{WN}h|}_2$は$\sigma_1(\bar{W}_{WN})=\sqrt{d_o}$で他の$t$においては$\sigma_t(\bar{W}_{WN})=0$となるときに、${||\bar{W}_{WN}h||}_2=\sqrt{d_o}$で最大になります。
結局、何が言いたいかというと、行列$\bar{W}_{WN}$が学習していくにつれて、段々と行列のランクが1に近づいていってしまっているということです。
ランクが1の行列は、単一の特徴しかつかめないような行列です。そう考えると、Discriminatorに使ってしまったら、ある特定の特徴を持った画像だけ本物、あとは偽物、というように判定するようになってしまい生成した画像は多様性を失うのではないでしょうか。
そう考えると、より多くの特徴を使って真贋を判定したい場合には、${||\bar{W}_{WN}h||}_2$ができるだけ大きいほうがいいでしょう。
2つ目のWeight Clippingですが、これはWGANでLipschitz連続性を課すため提案された手法です。これもWeight Normalizationと同じようにランクが小さくなり、多くの特徴を持つことができないという結果になってしまいます。
では、Spectral Normalizationはどうかと考えてみると、Spectral Normalizationは最大特異値だけペナルティを与えています。ほかの特異値にはペナルティがかかっていないため、重み行列のランクに一切の制限を加えていないので、様々な特徴を表現することができます。やっほぃ!
3つ目はOrthonormal Reqularizationです。これは、正則化項として、${{||W^T W - I||}_F}^2$を追加します。目的としては、Spectral Normalizationと同じですが、手法としては結構違っていて、この手法はすべての特異値を$1$にするような正則化を行っています。
最後にGradient Penaltyです。これは、言わずと知れたWGANgpで提案されたものです。${|\nabla_{\hat{x}}f|}_2=1$(ただし、$\hat{x}$は生成画像と実画像をランダムな割合で混ぜ合わせたもの)となるような正則化項を与えます。
この方法は、ランクが少なくなったりしないので、良きです。しかし、${||\nabla_{\hat{x}}f||}_2$は$\hat{x}$を出すときに生成画像を使っています。そのため、学習するたびに生成画像は分布が変化するので、学習が不安定になりがちです。また、勾配を正則化に使っているので、逆伝播を計算する必要あって、それにとても時間がかかります。
実験
CIFAR-10、STL-10を使って実験して、Spectral Normalizationが本当にいいのかを見ていきます。
学習はAdamを使い、OriginalのGANをベースに実験をします。
ハイパーパラメータの設定は、それぞれ上のようにします。
結果は上の通りです。Inception Score(IS)なので、高いほうが良いです。Spectral Normalization(SN)が安定していいように見えますね。あとOrthonormal Regularizationを良さげです。
次は、FIDで比べたものですので、低いほうがいいです。これは、Spectral Normalizationが一番いいように見えます。続いて、Orthonormal Regularizationです。
これは、ネットワークアーキテクチャで比べたものです。安定してResNetベースのSpectral Normalizationを使ったものがよさそうです。
上の画像は、CIFAR-10で学習されたモデルで生成されたものですが、WGANgpだと設定によってはうまくいかないこともあるようですね。
設定Fの場合は、Spectral NormalizationでもWeight Normalizationでも色が単調に見えます。
設定Eで比べると、Spectral Normalizationの方が色が多様な気がします。
今度は、それぞれの手法での特異値とその大きさ、層の位置の関係を表したものです。
ここでもSpectral Normalizationが他の手法と比べて大きく特異値を保っています。
また、Spectral Normalizationでも、層が深くなるにつれて特異値が小さくなってしまっていることが分かると思います。
これに対して対策をした論文、Sparsity Aware Normalizationがちょっと前(2021年3月)に公開されています。
上の図は、特徴マップの次元数とISの関係を表したものです。Orthonormal Regularizationの方は、次元数を増やすほどに悪くなってしまっています。これは、Orthonormal Regularizationがすべての特異値を$1$に保つように正則化するので、無理やり様々な特徴を使って判定しているので、精度が落ちていると思われます。
もしかしたら、次元低くしてOrthonormal Regularization使ったほうがSpectral Normalizationのよりも良いんじゃないか、とちょっと思ったんですが、これはOrthonormal Regularizationが最もうまくいった設定Cを使っているので、そういうわけじゃないんですね。
最後の図です。ISとイテレーション数の関係です。Spectral Normalizationでは、学習をまわすほど良くなる傾向にありますが、Orthonormal Regularizationはまわしすぎると逆に悪くなる傾向にあるようです。恐らく、学習をまわしすぎると正則化が効きすぎるのでしょう。
最後に
今まで無免許で使ってきたSpectral Normalizationのことが論文を読んだことでよく分かり、すっきりしました。仮免許ぐらいは取得できたんじゃないかと思います。
はじめてのQiitaでしたが、思っていたよりも書くのに時間がかかりました。。。
原著論文で端折ったところもあるので、もし興味があったら原著の方を読んでみてください。
参考文献
特異値分解: 「これなら分かる応用数学教室―最小二乗法からウェーブレットまで」
冪乗法: Wikipedia