重みのスケールに依存しないSGD: Path Normalized Optimization in Deep Neural Network
表題の論文を読んだのでまとめます！
Path-SGD を考えたモチベーション
ニューラルネットワークがこの論文の主題です。
Rescaling
今、あるニューラルネットワークの $i, i+1$ 番目の隠れ層の重み $W_i, W_{i+1}$ を取り出して $i$ 番目の重みを $x$ 倍して $i+1$ 番目の重みを $1/x$ 倍する操作を考えてみます(バイアスの大きさは 0 とします)。$i$ 層の出力を $y_i$ 入力を $z_i$ と表すことにします。隠れ層での計算は入力に対し重みを行列積することで表現できますから上記の操作後の隠れ層での出力を $\hat{y}_i$ とすると
$$ \begin{aligned} y_{i+1} &= W_{i+1} \cdot y_{i} = W_{i+1} \cdot W_i \cdot z_{i} \\ \hat{y}_{i+1} &= \frac{1}{x} W_{i+1} \cdot \hat{y}_i = \frac{1}{x} W_{i+1} \cdot (Wx) \cdot z_{i} \\ &= \frac{1}{x} \times x W_{i+1} \cdot W_i \cdot z_i \\ &= y_{i+1} \end{aligned} $$
となり出力は変化しません。このような重みに対する操作を rescaling と呼ぶことにします。
勾配法の Scaling への弱さ
ニューラルネットワークの学習では、勾配法の近似である確率的勾配法 (Stochastic Gradient Descent) が使われ、勾配計算には Forward Backward が用いられます。 この時重みへの正則化としてL2正則化を考えて少し重みをゼロに近づけるような効果を持たせて SGD を使うことが多いです。 この正則化は WeightDecay と呼ばれることもあります。
さてこの確率的勾配法、ひいては勾配法はネットワークの rescaling に対して不変ではありません。 これは勾配法の重み更新で使う Forward/Backward の勾配が、関連するレイヤーの重みスケールに依存しているからです。
そのため勾配法をネットワークの各層の重みが違うスケールをもっている Unbalanced な場合に使うと、学習は上手く進みません。 以下のグラフはスケールがそろっているときと揃っていないときのロス関数値をプロットしたものです。
以下の図はSGDでの重み更新の具体です。左のネットワークではすべての重みのスケールが 1~10 程度なので調度良く値が更新されています。一方で右の例では10~100 のものと 0.1 のものが混合しています。そのため勾配更新を行うと大きい値の物はほぼ変わらず、小さい値の重みが一気に更新されてしまっています。
「じゃあ重みが歪んでいる時でもいい感じに更新する最適化方法無いの？」という疑問に応えるのが、この論文の趣旨になっています。それが表題にもある Path-SGD です。*1
準備
考えるのは $d$ 層のニューラルネットワークです。 活性化関数は RELU とします。 このネットワークをDAG $G(V,E)$ として表現します。 ここで $V$ はノードを表し $E$ はエッジを表しています。 $D, C \in \mathbb{N}$ を入力, 出力の次元数とします。即ち入力ノードは
$$ v_{in[1]}, \ldots, v_{in[D]} \in V $$
と表現できます。
一般的な正則化
はじめに、ネットワークの各レイヤーの重みごとにグルーピングされた正則化を考えていきます。 パラメータ $p \ge 1, q \le \infty$ を用いて一般的な正則化関数は以下のように表現できます。
$$ \mu_{p,q}(w) = \left( \sum_{v \in V} \left( \sum_{(u \to v) \in E} | w_{u \to v} |^p \right)^{q/p} \right)^{1/q} $$
各ノードごとの足し算の項 $\sum_{v \in V}$ が有るためこれはグルーピングされた正則化関数の一般化であることがわかります。簡単な例であると $p = q = 1$ の時は L1 正則化に $p = q = 2$ の時にL2正則化(Weight Decay)に対応します。
これら以外にも重みの最大値を正則化とする max-normalization というものがあります。 これは上記の一般形式で $q \to \infty$ を考えた時の値として表現できます。
$$ \mu_{p,\infty}(w) = \sup_{v \in V} \left( \sum_{(u \to v) \in E} | w_{u \to v} |^p \right)^{1/p} $$
重み不変な正則化
重み不変な正則化とは、重みに対して rescaling の操作を行った時に値が変わらない正則化関数を指しています。
max-normalization は重みの上限をとっているためこれは明らかに重み不変性を満たしていません。適当なノードの入力を小さくして出力を大きくすれば、いくらでも上限値を大きくすることができるからです。
ここで、もし多数のネットワークが重み不変、すなわち入力が同じなら出力も同じである、という性質を満たしていて、どれでも好きなものを選べると仮定しましょう。 どれでも出力が同じなら、そのなかでもっとも正則化関数の値が小さいものを選ぶことが望ましいはずです。
現実には重み不変を持つネットワーク全てに対して正則化を考えることは難しいですが、ラッキーなことに max-normalization の場合この最小値を計算することが出来ます。それが以下の path-normalization です。
$$ \phi_p(w) = | \pi(w) |_p = \left( \sum_{v_{in}[i] \to \cdots v_{out}[j]} \left| \prod_{k=1}^d w_{e_k} \right|^p \right)^{1/p} $$
ここで $\pi(w) \in \mathbb{R}^N$ はパスベクトルと呼び、次元数 $N$ は入力から出力までのパスの組み合わせ数です。 上記の式中では $v_{in}[i] \to \cdots v_{out}[j]$ と表されている部分が組み合わせ数に対応しています。 この式もやや込み入っていますが、入力 $i$ と出力 $j$ を色々と入れ替えた時にできるすべてのパス、を表しているのでパスの組み合わせに相当していることが解ると思います。
パスベクトルの各次元の値はそのパス上の重みをすべて積算したものになっています。 そしてパスベクトルの $p$ ノルムが $\phi_p$ です。ちょっとややこしいですね。
ややこしいものを持ちだしたのには理由があって、実はこの $\phi_p$ は max-normalization との間に以下のような面白い性質を持ちます。
$$ \phi_p(w) = \min_{\tilde{w} \sim w} \left( \mu_{p, \infty}(\tilde{{w}}) \right)^d $$
ここで $\tilde{w} \sim w$ で表される $\tilde{w}$ は $w$ を rescaling した中で任意の入力に対して出力が $w$ と同じになるという同値類です。
これは即ち path-normalization $\phi_p(w)$ の値は rescaling な重みの中での max-normalization の最小値であることを表しています。 ということは path-normalization を最小化すれば勝手に max-normalization の最小値も下がるということがわかります。しかもその値は rescaling を許したすべての max-normalization の値の中でもっとも小さいのです。これは嬉しい。
また明らかに任意の $w \sim \tilde{w}$ に対して $\phi_p(w) \sim \phi_p(\tilde{w})$ であるので、これより path-normalization は重み不変性を持っていることがわかります。
以上から path-normalization は
- max-normalization の最小値であること
- 重み不変であること
という良い正則化項であることがわかりました。
Path-Normalization を持つ SGD
では path-normalization を正則化項とするような目的関数を考え、これに対して勾配法を考えていきます。 この時すべての重みを厳密に最適化することが難しいので、特定のエッジの重み $w_e$ のみを更新することを考えます。
今第 $t$ ステップの重み $w^{(t)}$ を得ていて $t+1$ での $e$ の重み $\hat{w}_{e}^{t+1}$ を計算する場面を考えます。 ロス関数を $L$ として偏微分すると以下の式を得ます。
$$ \hat{w}_{e}^{t+1} = w_{e}^t - \frac{\eta}{\gamma_p(w^{(t)}, e)} \frac{\partial L}{\partial w}(w^{(t)}) $$
ここで $\gamma $ は以下の式で表される重みの積になります。
$$ \gamma_p (w, e) = \left( \sum_{v_{in}[i] \cdots \overset{e}{\to} \cdots v_{out}[j]} \prod_{e_k \ne e} | w_{e_k} |^p \right)^{2/p} $$
これは入力から出力への経路のうちで $e$ の重みだけを含まない $\pi(w)$ のノルムになっています。 このことから例えば $e$ に対応する重みがとても大きい値を持っていると仮定すると $\gamma$ の値はその値が除外されるため小さくなり、勾配に係る係数は大きくなります。反対に小さい時には $\gamma$ は大きくなり勾配に係る係数は小さくなります。
結果として自分の大きさに応じた勾配更新が行われるようになることがわかります。 これは通常のSGDでは得られなかった性質です。（前半のレイヤーごとに重みが違う場合の更新を考えてみるとわかりやすいと思います）
このことを数学的に行ったのが Theorem4.1 でこのPath-SGDの更新はスケールに対して不変であることが示されています。
じっけん
有名データセットを用いて数値実験を行っています。比較するのは Path-SGD を Unbalanced な条件で学習させたものと SGD/AdaGrad を Unbalanced/Balanced な条件で学習させたものです。
まず SGD を Unbalanced で実行すると MNIST ですら学習できていないことがわかります。 AdaGrad はもうちょっとがんばっていますがやはり苦しそうです。 一方で Path-SGD では Unbalanced にもかかわらず学習が上手く進み、また Balanced の場合の AdaGrad と並ぶかむしろそれよりも学習が早いことがわかります。*2
まとめ
- Unbalanced な重みでは SGD は機能しない。それは勾配法が重み不変性をもたないから
- Path-Normalization は max-normalization の最小値であり、重み不変。
- それを正則化項として持つ objective を考えた path-SGD も重み不変。故に重みが Unbalanced でも学習が進む。
感想
重み不変性という観点が純粋に面白かった。また max-normalization の重み不変空間上での最小値が path-normalization の値につながっているのも綺麗。
一方でそもそも重みのスケールが異なるようなネットワークって存在する時あるんだろうか (転移学習とかならありうる?) とか数値実験で目的関数違うやつを同じグラフにプロットするのはどうなの（せめて精度とかにして欲しかった）とか思ったり。
気になった点としては、やはり重みごとにその勾配に対する係数を保存しないとだめだという点 ($\gamma_e$ の所)。著者は効率的な実装方法があるよーと述べていて forward backward 方式で1回計算すればいいように済ませる実装方法を提案しているがそれだとしても毎回の更新時に重みの係数も更新するのは面倒だし、重みが増えてくるとメモリ的にもしんどそう（とはいえたかだか二倍程度だからしれているのか?)。自分で実装してみたいなと思った時もこの面倒さがボトルネックでまだ出来ていません。Adam とか他の適合的なアルゴリズムが楽すぎるってことなのかな…
