Zenn
Closed17

BatchNormalizationの復習

#
深層学習
ピン留めされたアイテム
takoroytakoroy

Batch Normalizationの復習

以下の論文を読む。

  • Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift
  • Understanding Batch Normalization
  • How does Batch Normalization Help Optimization?
takoroytakoroy

そもそも何か？

BatchNormalizationは、ミニバッチ\mathcal{B}に含まれるm個のサンプルをもとに以下のように計算するレイヤー。\gamma, \betaは学習対象のパラメータ。

\begin{array}{c} \mu_{\mathcal{B}}=\frac{1}{m} \sum_{i=1}^{m} x_{i} \\ \sigma_{\mathcal{B}}^{2}=\frac{1}{m} \sum_{i=1}^{m}\left(x_{i}-\mu_{\mathcal{B}}\right)^{2} \\ \hat{x}_{i}=\frac{x_{i}-\mu_{\mathcal{B}}}{\sqrt{\sigma_{\mathcal{B}}^{2}+\epsilon}} \\ y_{i}=\gamma \hat{x}_{i}+\beta=\mathrm{B} \mathrm{N}_{\gamma, \beta}\left(x_{i}\right) \end{array}

https://paperswithcode.com/method/batch-normalization

takoroytakoroy

何を実現できるか？

内部共変量シフト（internal covariate shift）を減らすことができる。内部共変量シフトとは、トレーニングの過程で各レイヤーのパラメータが変化することで生じる、特定のレイヤーに入力されるデータ分布の変化を指す。

  • 後続する入力データの平均と分散を一定にできる。
    • scaleとshiftを表す\gamma, \betaによって、いかなる入力に対しても恒等変換が可能にする。例えば、0以上の値しかない特徴マップを正規化すると必ず負の値が生じるが、適切なscaleとshiftによって恒等変換が可能になる。
    • BNの直前のレイヤーでは、biasは不要。BNのパラメータである\betaがその役割を担う。
  • パラメータ（およびその初期値）のスケールへの勾配の依存性を減らす。
    • パラメータがa倍されると、勾配は\frac{1}{a}の大きさになる。
    • レイヤーのJacobianが1に近づくことにより、訓練が安定する。

https://arxiv.org/abs/1502.03167

takoroytakoroy

再現実験

colab

MNISTでの実験。元論文の1つ目の実験を追試。Linear層のパラメータは、\mathcal{N}(0, \sigma^2)で初期化。\sigma^2 = 0.1だと、BN無しの場合は訓練遅く進む一方で、BN有りだと早く訓練が進む。また、BN有りの方が精度は高くなる。

最後のSigmoid活性化に入力されるある特徴量の85%点、50%点、15%点をプロットすると、BN無しの場合では特に訓練の初期の中央値はブレブレだが、BN有りの場合では中央値が0付近で安定する。


以下は元論文の実験結果。

takoroytakoroy

BNの効果の分解

BNの主要な効果は、学習率を大きくすることで、正則化と高速な収束をもたらすことである。

  • BNによって大きい学習率でも収束が可能になる。
  • 大きい学習率によって、正則化と高速な収束が可能になる。

実験的な観察

BNなしのResnetモデルを訓練するための学習率1e-4を、BNありのResNetモデルの訓練にも適用すると、そこまで良いテスト精度が得られない。BNありの状態で学習率を大きくすることで、テスト精度の改善（＝正則化による汎化性能の向上）と、高速な収束が実現できる。大きい学習率だとBNなしではそもそも収束しないので、2つの効果を両立できるのは、BNのおかげと言える。

正則化の理論的背景

大きい学習率が高速な収束をもたらすのは（学習の過程で発散しない限り）ほとんど自明。一方、大きい学習率が正則化をもたらすという点はそこまで自明ではない。

大きい学習率がもたらす正則化の効果を端的に表すのが下式（論文中の(2)式）である。\nabla \ell(x)はデータセット全体の損失から計算される勾配で、\nabla_{S G D}(x)はミニバッチBに関する損失から計算される勾配を表す。下式は、これら2つの勾配の差異が、右辺で抑えられてしまうということを表している。

\mathbb{E}\left[\left\|\alpha \nabla \ell(x)-\alpha \nabla_{S G D}(x)\right\|^{2}\right] \leq \frac{\alpha^{2}}{|B|} \mathbb{E}\left[\left\|\nabla \ell_{i}(x)-\nabla \ell(x)\right\|^{2}\right]

2つの勾配の差異は小さいことは、訓練に使用しているデータセットにoverfitしたモデルが訓練されてしまうことを意味する。そのため、右辺はある程度大きいほうが良い。直感的には、データセット全体の勾配とはちょっと逸れる方向にいろいろ探索しながらモデルを更新していったほうが汎化性能が高まるよね、という話。

右辺を大きくする操作は、学習率\alphaを大きく、バッチサイズ|B|を小さくすることである。よって、学習率を大きくすることが正則化につながる。

むろん、あまりにも学習率を大きくすると発散してしまうし、バッチサイズが小さすぎると訓練に時間がかかるので、バランスが大切。

takoroytakoroy

BNによる発散の回避

一口に発散とっても、いくつかの観点がある。
ここでは、損失、勾配、activationという3つの観点から実証的な分析が行われている。

損失

BNによってなぜ発散が回避できるのか。訓練の最初の数回のミニバッチにおける損失を調査することで、発散の状況の把握を試みている。relative lossを損失の前ステップとの大きさの比と定義する。このrelative lossが大きいとき、発散が生じているとみなすことができる。

様々なミニバッチのサイズ（1,4,7）について、学習率（step-size）を変化させたときにrelative lossがどのように変化するかをまとめたのが以下のグラフ。BNなしだと学習率は1e-3程度が限界で、それ以上大きくすると発散していくことがわかる。BNありだとrelative lossはほぼ1のままなので、ある程度大きい学習率でも発散しなさそうだといえる。

勾配

BNの有無ごとに、学習の開始時における55番目のレイヤーの勾配のヒストグラムを以下に示す。BNなしだと、勾配のすそ野が重くなり、十分に学習率が小さくないと発散してしまうことがわかる。

activation

BNなしだと深い層に行けば行くほどactivationの平均や分散のスケールが大きくなることも観察できる。BNはこのような平均と分散の発散を、強制的に抑えることができる。

takoroytakoroy

BNによる勾配の偏りの抑制

BNなしだと、初期化時点ですでにいくつかの偏りを生じさせてしまう。この状態でモデルを訓練しようとすると、勾配に偏りが生じる。

  • 特定のクラスに関する勾配が偏る
  • 勾配の符号の偏り
  • 特定のチャネルが支配的になる

特定のクラスに関する勾配が偏る

初期化時点で、BNなしの場合は出力が特定のクラスに偏ってしまう。以下の図は、最初のミニバッチにおける最終出力に対する勾配の大きさを表したヒートマップだが、BNなしの時は一番右のクラスに関する出力の勾配のみが大きい。

勾配の符号の偏り

3 \times 3のConv層による入出力の関係は、以下のように書ける。ここで、S=\{-1,0,1\} \times\{-1,0,1\}は、カーネル内の位置を表す。

O_{b, c, x, y}=\sum_{c^{\prime}} \sum_{x^{\prime}, y^{\prime} \in S} I_{b, c^{\prime}, x+x^{\prime}, y+y^{\prime}} K_{c, c^{\prime}, x^{\prime}, y^{\prime}}

また、勾配は以下のように書くことができる。

\frac{\partial L}{\partial K_{o, i, x^{\prime}, y^{\prime}}}=\sum_{b, x, y} d_{o, i, x^{\prime}, y^{\prime}}^{b x y}, \quad \text { where } \quad d_{o, i, x^{\prime}, y^{\prime}}^{b x y}=\frac{\partial L}{\partial O_{b, o, x, y}} I_{b, i, x+x^{\prime}, y+y^{\prime}}

以下の表の一番右の列は単純な勾配の和の絶対値、その他の列は、絶対値の和になっている。和の絶対値と絶対値の和が近い場合、符号に偏りが存在することを意味する。bはバッチ内のサンプルに対するインデックスを、x,yは特徴マップの空間的な領域を表している。

\begin{array}{|r|c|c|c|} \hline & a=\sum_{b x y}\left|d_{c_{o} c_{i} i j}^{b x y}\right| & b=\left|\sum_{b x y} d_{c_{o} c_{i} i j}^{b x y}\right| & a / b \\ \hline \hline \text { Layer 18, with BN } & 7.5 \mathrm{e}-05 & 3.0 \mathrm{e}-07 & \mathbf{2 5 1 . 8} \\ \hline \text { Layer 54, with BN } & 1.9 \mathrm{e}-05 & 1.7 \mathrm{e}-07 & \mathbf{1 1 2 . 8} \\ \hline \text { Layer 90, with BN } & 6.6 \mathrm{e}-06 & 1.6 \mathrm{e}-07 & \mathbf{4 0 . 7} \\ \hline \hline \text { Layer 18, w/o BN } & 6.3 \mathrm{e}-05 & 3.6 \mathrm{e}-05 & \mathbf{1 . 7} \\ \hline \text { Layer 54, w/o BN } & 2.2 \mathrm{e}-04 & 8.4 \mathrm{e}-05 & \mathbf{2 . 6} \\ \hline \text { Layer 90, w/o BN } & 2.6 \mathrm{e}-04 & 1.2 \mathrm{e}-04 & \mathbf{2 . 1} \\ \hline \end{array}

特定のチャネルが支配的になる

以下の図は、特定のレイヤー(45番目)のパラメータに対する勾配の絶対値を可視化したものである。BNなしの場合は、特定の入出力に関わるチャネルのみが大きくなっている。BNありの場合は、様々なチャネルにばらけている。

takoroytakoroy

重みの初期化と勾配爆発の関係

勾配爆発は重みの初期化がもたらす自然な帰結である。
近年のRandom Matrix Theoryの研究成果によると、各要素をガウス分布で初期化した行列をM個掛けあわせたX = X_1X_2...X_Mを考えると、Xの特異値の分布は、非常にすそ野が広くなる。Mが大きいほど、最小特異値は0に近くなり、最大特異値は大きくなる。

  • Dang-Zheng Liu, Dong Wang, Lun Zhang, et al. Bulk and soft-edge universality for singular values of products of ginibre random matrices. In Annales de l’Institut Henri Poincaré,
    Probabilités et Statistiques, volume 52, pages 1734–1762. Institut Henri Poincaré, 2016.

ここで、以下のような最小化問題を考える。A=A_1A_2 \ldots A_Mとする。

{\min}_{x}\|A x-y\|^2

Mが大きいほどAの最小特異値と最大特異値の比は大きくなる。一般に、最小特異値と最大特異値の比である条件数\kappa(A)が大きいとき、この問題を勾配降下法で解こうとすると、以下のような困難をもたらすことが知られている。

  • 収束が遅くなる
  • より小さい学習率の必要性
  • 異なる部分空間どうしの勾配の比率が増加する

深層学習のネットワークは、深くなるほど上記のような問題が生じやすくなり、Mが大きいほど条件数が大きくなって最適化が困難になるという問題と並行関係にある。

BNは、このような深層ネットワークが持つ困難を回避しているといえる。

takoroytakoroy

内部共変量シフトは無関係？

VGGでCIFAR-10を訓練。そのとき、BNの有無によって汎化性能が異なることは明らか（下図の真ん中)。しかし、あるレイヤーのactivationの分布を訓練ステップごとに可視化してみても、両者に大きな違いは見られない。

  • BatchNormの有効性は、本当に内部共変量シフトに関連しているのか？
  • BatchNormによるレイヤー入力分布の安定化は、内部共変量シフトの削減にも効果的なのか？
takoroytakoroy

レイヤー入力分布の安定化は性能向上と無関係

BNの直後に、毎ステップ平均と分散が変化するノイズを加えるNoisy BatchNormという実験設定を追加する。Noisyな場合でも、BNと同様にStandardと明らかに異なる性能の向上がみられる。

Noisyな場合は、BNによって安定化されたレイヤー入力分布をわざと崩している（強制的に内部共変量シフトを起こしている、Layer #13がわかりやすい）が、それでも性能向上はそのままなので、「レイヤー入力分布の安定化→性能向上」という事実は観察できない。

takoroytakoroy

BNは内部共変量シフトを減らさない

訓練のパフォーマンスに直接関連する内部共変量シフトの悪影響はあるのだろうか。あるとすれば、BNはそれを本当に抑制しているのか。

内部共変量シフトを以下のように定義する。i番目のレイヤーのtステップ目の勾配の差異によって定義されている。

\left\|G_{t, i}-G_{t, i}^{\prime}\right\|_{2}

G_{t, i}, G_{t, i}^{\prime}は以下のように定義される。i番目以降のレイヤーは固定して計算していることに注意。

\begin{aligned} G_{t, i} &=\nabla_{W_{i}^{(t)}} \mathcal{L}\left(W_{1}^{(t)}, \ldots, W_{k}^{(t)} ; x^{(t)}, y^{(t)}\right) \\ G_{t, i}^{\prime} &=\nabla_{W_{i}^{(t)}} \mathcal{L}\left(W_{1}^{(t+1)}, \ldots, W_{i-1}^{(t+1)}, W_{i}^{(t)}, W_{i+1}^{(t)}, \ldots, W_{k}^{(t)} ; x^{(t)}, y^{(t)}\right) \end{aligned}

VGGと25層のLinear Network（DLN、SGDではなく、full-batch gradient decent）とで比較する。\ell_2-diffは先ほど定義した内部共変量シフトで理想的には0、CosAngleは2つの勾配間のCosineで理想的には1。

驚くべきことに、BNは内部共変量シフトを減らすどころか、増やす。これはDLNだと顕著で、VGGだとBNなしのときと違いがない。また、一貫してBNありのほうが性能は良い。つまり、内部共変量シフトは性能に関係がないといえる。

takoroytakoroy

BNは損失関数のリプシッツ連続性を高める

以下は、VGGの訓練の過程でパラメータを勾配の方向に動かしたときに、BNによって損失関数の安定性（リプシッツ連続性）が高まることを示した実験結果である。

(a)は、損失そのものの安定性を測ったもの。訓練ステップの過程で損失の勾配を計算し、その方向にパラメータを動かしたときに損失がどのように変化するかを測定している（🤔このとき動かすパラメータは特定のレイヤーに限定している？　影が付いているのはどういう意味？）。BNなしの場合は、特に訓練初期にLossの変動が激しい。これから、損失関数がBNによって滑らかになっているといえる。

(b)は、勾配の信頼性と予測性を測ったもの。勾配がある程度一貫していることは、最適化にとっては重要な指標。勾配と、その方向にパラメータを動かしたときの勾配を比較し、\ell_2距離を測定している。BNなしの場合は、全体的に大きな値になっているので、勾配の方向にパラメータを動かすと、全然違う方向に勾配がかかるということを示している。

(c)は、"effective"な（有効な）勾配の滑らかさを測ったもの。移動した距離に対する、勾配の差異（\ell_2距離）の最大値を指している。（🤔様々なレイヤーに関して、\left|G\left(w_{t}\right)-G\left(w_{t+1}\right)\right| / \left\|w_{t}-w_{t+1}\right\|を計算し、その最大値をプロットしている、ということだろうか）。BNなしの場合は、近隣で勾配が大きく変わることが多く、つまり勾配が滑らかでない、ということがわかる。なお、BNなしの場合は発散してしまうため、\eta=0.4までにしている。

L-Lipschitz

関数fがL-Lipschitzであるとは、以下のような条件をすべてのx_1, x_2で満たしていることをいう。

\left|f\left(x_{1}\right)-f\left(x_{2}\right)\right| \leq L\left\|x_{1}-x_{2}\right\|

\beta-smooth

関数fが\beta-smoothであるとは、勾配が\beta-Lipschitzであることを指す。

takoroytakoroy

その他のNormalizationとの比較

BNの代わりに、\ell_1, \ell_2, \ell_{ \infty }という3種類のNormalizationを試している。いずれの場合もBNと同様の性質（パフォーマンス、収束速度、loss landscapeなど）が得られることを確認している。そのため、本論文の結論としては、BNの発見は一種のセレンディピティであるとしている。

takoroytakoroy

BNに関する理論的な分析

入力をx、出力をy=Wxとする単純なモデルを考える。BNを使った、\hat{y}=\frac{y - \mu}{\sigma} + \betaも考える。
BNなしの時の損失を\mathcal{L}、BNありの時の損失を\hat{\mathcal{L}}とする。

💡数式追うのが大変なのでざっくりと

Activationから見た損失のリプシッツ性

まずは、損失関数のリプシッツ性を捉える\left\|\nabla_{y_{j}} \mathcal{L}\right\|について考える。損失関数のリプシッツ定数は、最適化にとって極めて重要な役割を担うことが知られている。BNを使うことによって、損失関数のリプシッツ定数を小さくできる。

下式右辺の真ん中の項と最後の項はそれなりに大きくなること、\sigma_jは経験的に大きくなることが知られているため、\left\|\nabla_{\boldsymbol{y}_{j}} \widehat{\mathcal{L}}\right\|^{2}は、\left\|\nabla_{\boldsymbol{y}_{j}} \mathcal{L}\right\|^{2}よりも小さくなることがわかる。

\left\|\nabla_{\boldsymbol{y}_{j}} \widehat{\mathcal{L}}\right\|^{2} \leq \frac{\gamma^{2}}{\sigma_{j}^{2}}\left(\left\|\nabla_{\boldsymbol{y}_{j}} \mathcal{L}\right\|^{2}-\frac{1}{m}\left\langle\mathbf{1}, \nabla_{\boldsymbol{y}_{j}} \mathcal{L}\right\rangle^{2}-\frac{1}{m}\left\langle\nabla_{\boldsymbol{y}_{j}} \mathcal{L}, \hat{\boldsymbol{y}}_{j}\right\rangle^{2}\right)

損失の2次の滑らかさ

次に、2次のオーダーの性質\hat{\boldsymbol{g}}_{j}^{\top} \boldsymbol{H}_{j j} \hat{\boldsymbol{g}}_{j}を確認する。\hat{\boldsymbol{g}}_{\boldsymbol{j}}=\nabla_{\boldsymbol{y}_{j}} \mathcal{L}、\boldsymbol{H}_{j j}=\frac{\partial \mathcal{L}}{\partial \boldsymbol{y}_{j} \partial \boldsymbol{y}_{j}}とする。こちらも先ほどと同様、BNなしの時よりも小さくなるということが示されている。

\left(\nabla_{\boldsymbol{y}_{j}} \widehat{\mathcal{L}}\right)^{\top} \frac{\partial \widehat{\mathcal{L}}}{\partial \boldsymbol{y}_{j} \partial \boldsymbol{y}_{j}}\left(\nabla_{\boldsymbol{y}_{j}} \widehat{\mathcal{L}}\right) \leq \frac{\gamma^{2}}{\sigma^{2}}\left(\hat{\boldsymbol{g}}_{j}^{\top} \boldsymbol{H}_{j j} \hat{\boldsymbol{g}}_{j}-\frac{1}{m \gamma}\left\langle\hat{\boldsymbol{g}}_{j}, \hat{\boldsymbol{y}}_{j}\right\rangle\left\|\frac{\partial \widehat{\mathcal{L}}}{\partial \boldsymbol{y}_{j}}\right\|^{2}\right)

パラメータから見た損失関数のリプシッツ性

以下のような表記を用いる。

g_{j}=\max _{\|X\| \leq \lambda}\left\|\nabla_{W} \mathcal{L}\right\|^{2}, \quad \hat{g}_{j}=\max _{\|X\| \leq \lambda}\left\|\nabla_{W} \widehat{\mathcal{L}}\right\|^{2}

この時、以下が成り立つ。こちらも、パラメータから見たときのリプシッツ定数が、BNによって小さくできることを示している。

\hat{g}_{j} \leq \frac{\gamma^{2}}{\sigma_{j}^{2}}\left(g_{j}^{2}-m \mu_{g_{j}}^{2}-\lambda^{2}\left\langle\nabla_{\boldsymbol{y}_{j}} \mathcal{L}, \hat{\boldsymbol{y}}_{j}\right\rangle^{2}\right)

パラメータの初期化

\boldsymbol{W}^{*}, \widehat{\boldsymbol{W}}^{*}をBNなしのときとありの時の局所解の集合とする。どんなパラメータの初期値W_0に対しても、以下が成り立つ。ここで、W^{*}, \widehat{W}^{*}は、最も近い最適解とする。

\left\|W_{0}-\widehat{W}^{*}\right\|^{2} \leq\left\|W_{0}-W^{*}\right\|^{2}-\frac{1}{\left\|W^{*}\right\|^{2}}\left(\left\|W^{*}\right\|^{2}-\left\langle W^{*}, W_{0}\right\rangle\right)^{2}

つまり、いかなる初期値に対しても、最適解はBNありのほうが近くに存在することを表す。

takoroytakoroy

理論的な部分は、最適化周りの復習もしないと理解が難しそう。

このスクラップは2021/01/24にクローズされました