ﾐ゜

概要

本記事では、CVPR 2020の "Controllable Orthogonalization in Training DNNs" (Lei Huang, Li Liu, Fan Zhu, Diwen Wan, Zehuan Yuan, Bo Li, Ling Shao) という論文を紹介します。実際に実装して論文に載っている結果を検証していきます。

おおまかには、

勾配消失/爆発の対策として、ネットワークの重みを直交化したい。
既存手法には、高計算コスト、不安定（主に0除算による）などの問題がある。
これらに対処するため、Newton’s iterationを用いて直交化する手法を提案。
これにより直交性を課す強さと表現能力のバランス調整も可能に。
分類タスクでの性能で既存手法を上回ったほか、GANの学習を安定にする効果も確認。

という内容です。

背景

深層ニューラルネットワークを訓練するとき、勾配爆発・勾配消失が大きな問題となります。これらを防ぐために、重みの初期化をうまくやることで各層のactivationの分散を一定にしたり、その発展で直交に初期化する方法が提案されており、実験的にも収束を早める効果が確認されています。しかし、初期値を直交にするだけでは訓練中に直交性が崩壊しないことを保証できません。ロスに項を追加してソフトに何らかの制約を課す手法も提案されていますが、やはり崩壊しない保証はありません。

この問題に対して、学習中にも直交性を明示的に生み出す仕組みも提案されていますが、今度は固有値分解の計算コストや、誤差逆伝播で0除算が起こりうる不安定さが問題になっていました。

手法: Orthogonalization by Newton’s Iteration（ONI）

そこで本論文は、 ニュートン法を用いたiterativeな計算で重みを直交化することを提案しました。直交な行列に収束するように操作を繰り返すことで、固有値分解を避けて安定に直交化することができます。さらにその繰り返し回数を調節することで直交化の具合と表現能力のバランス調整も可能になりました。

全結合層

$\mathbf{y = Wx+b}$

を考えます。入出力が $\mathbf{x,y}$ で、重みが $\mathbf{W}$ 、バイアスが $\mathbf{b}$ です。入出力の次元は $d, n$ で、 $\mathbf{W}\in \mathbb{R}^{n\times d}$ です。まず、 $n \leq d$ の場合を考えます。

本手法の目的である重みの直交化とはつまり、

$\mathbf{WW^T=I}$

としたいということです。

この論文では、まず既存手法の問題点を上げたうえで、その解決策を提案しています。

既存手法では、 $\mathbf{V}$ というパラメータと $\mathbf{V}$ から直交な行列 $\mathbf{W}$ を作り出す関数 $\mathbf{W=\phi (V)}$ を用意していました。ここで、 $\mathbf{V-W}$ のフロベニウスノルム（ベクトルとみなしたときのL2ノルム）が最小になるように $\phi$ を定めると、

$\mathbf{W=\phi (V) = D} \Lambda^{-\frac{1}{2}} \mathbf{ D^T V}$

となるそうです。ここで、 ${\Lambda}$ は共分散行列 $\mathbf{ S=VV^T}$ の固有値からなるベクトルで、 $\mathbf{D}$ は対応する固有ベクトルからなる行列です。

この既存手法には問題点が2つありました。それは、

固有値分解が必要で、特にGPUで計算コストが高いこと
固有値分解のところで誤差逆伝播時に0除算が生じうるので、不安定であること

です。

そこでこの論文は、

$\mathbf{W=\phi (V) = D} \Lambda^{-\frac{1}{2}} \mathbf{ D^T V=S^{-\frac{1}{2}}V}$

となることに着目し、ニュートン法で $\mathbf{S}^{-\frac{1}{2}}$ を計算することで固有値分解を避ける直交化、Orthogonalization by Newton’s Iteration（ONI）を提案しています。

具体的には次のように計算します。

$\begin{cases}\mathbf{B_0=I}\\ \mathbf{ B_t=\frac{1}{2} (3B_{t-1}-B_{t-1}^3S) } \end{cases}$

$\mathbf{\left\|I-S\right\|_2} \lt 1$ のもとで、 $\mathbf{B_t}$ は $\mathbf{S}^{-\frac{1}{2}}$ に収束します。

では、 $\mathbf{V}$ が $\mathbf{\left\|I-S\right\|_2} \lt 1$ を満たすようにするにはどうしたらよいでしょうか。一つの方法は $\mathbf{V}$ の特異値 $\mathbf{\sigma(V)}$ が全て $0\leq \mathbf{\sigma(V)} \leq 1$ を満たすことですが、そのような制約を $\mathbf{\sigma(V)}$ そのものに課すと学習が阻害されてしまいます。そこで、あらたなパラメーター $\mathbf{Z}$ を用意し、

$\mathbf{V=\frac{Z}{\|Z\|_F}}$

とします。

以上で、固有値分解によらない直交化が可能になりました。次に、収束を早くすることを考えます。 $\mathbf{V}$ の特異値が1に近いほど収束が早くなります。そこで、 $\mathbf{Z}$ の中心化を行います。

$\mathbf{Z_c=-\frac{1}{d}Z11^T}$

$d$ は入力 $\mathbf{x}$ のサイズです。さらに、さきほど $\mathbf{V}$ を作り出す式では $\mathbf{{\|Z\|_F}}$ で割っていましたが、もう少し大きい値で割ることもでき、

$\mathbf{V=\frac{Z}{\sqrt{\|Z_cZ_c^T\|_F}}}$

とします。このほうが特異値がより1に近い値になるそうです。

ここまでを図にするとこのような感じです。

ここまでは $n \leq d$ の場合を考えていましたが、 $n \gt d$ の場合はどうなるでしょうか。今度は全ての行が直交する $\mathbf{WW^T=I}$ ようにすることは不可能です。しかし、実はそのまま同じアルゴリズムを適用することはでき、その結果列が直交するように

$\mathbf{W^TW=I}$

収束します。既存手法では行をグループに分けることで各グループ内で直交化していましたが、本手法では $n,d$ の大きさに関係しない統一的なアルゴリズムで直交化することができるようです。

さて、ここまでは直交化する方法を考えていましたが、繰り返しの回数を調節することで直交化の具合を調節することが可能です。

重みが直交行列になっていると、activationのノルムや勾配の分布をレイヤーの前後で不変にすることができ（詳細*1）、勾配消失、爆発を防ぐ効果があります。一方で、制約が強いためにモデルの表現能力を落としてしまう事も知られており、そのバランスを調節できると嬉しいそうです。本論文の提案手法では、ニュートン法の繰り返し回数を調節することでバランスを調節できるというように述べられていました。

最後に、細かい部分の補足説明を列挙します。

ReLUが続く場合は、重みを $\sqrt{2}$ 倍します。（イメージはHeの初期化と同じ感じでしょうか。）
$\mathbf{W^TW=I}$ は直交化と同時に各行の大きさが1になるような制約になっていますが、これを緩和するために別の学習可能なパラメータを用意して掛けることをします。これにより $\mathbf{W^TW=}\Lambda$ となります。
全結合層について述べましたが、畳み込み層についても同様で、異なる出力チャンネルに対応する重み同士が直交するようにします。

手法の説明は以上になります。

実験

実装・実験をしてONIの確認していきます。まず全結合層+ONIを実装し、Fashion-MNISTのデータで実験を行いました。さらに畳み込み層+ONIも実装し、Cifar-10において実験しました。

実装はこちらです。全結合層と畳込み層の実装はoni_module.pyです

github.com

全結合層+ONI

まず、6層MLPでHeの初期化、直交初期化、ONIを試し、ONIの効果を確認します。詳細な実験設定は以下の通りです。

モデル	6層MLP （ReLU・各層256ノード）
データセット	Fashion-MNIST
バッチサイズ	256
学習率	0.05
エポック数	80
ONIの繰り返し数	5

すると、Trainデータに対する精度の変化は下図のようになりました。これは論文のFigure 5 (a)に対応します。（論文ではHeの初期化ではなくplainとありますが、ここではHeと比較しました。）確かにONIが明らかに早く収束しています。

f:id:y60:20200805191322p:plain — 6層MLP・Fashion-MNISTでの比較

このとき、200 iteration時点の特異値の絶対値は下図のようになりました（論文のFigure 5 (b)に対応します。）。activationの次元数に等しい256個の特異値があり、それをソートして左から並べています。直交性を考えない場合は特異値がばらついて、直交性を課す場合は1に近づく点は論文と一致しました。しかし、よく見るとONIの場合に後ろのほうで急落し、論文とは少し異なる結果になりました。（なぜこうなったのかは分かりませんでした。）

f:id:y60:20200805171009p:plain — 6層MLP・Fashion-MNIST・200 iteration後の特異値

次に、20層MLPでScalingの効果を確認します。Scalingとは、ReLUが続く場合に重みを $\sqrt{2}$ 倍することを指します。特に層が多い場合に勾配消失を防ぎ、学習を安定させる効果があります。実際の結果が下図です。これは論文のFigure 6 (b)に対応します。確かにScalingがあるほうが収束が早くなりました。

f:id:y60:20200805175336p:plain — 20層MLP・Fashion-MNISTでの比較

畳み込み層+ONI

畳み込み層でもONIの効果を確認します。実験設定は次の通りです。

モデル	VGG-Style
データセット	CIFAR-10
バッチサイズ	128
学習率	0.05
エポック数	160
ONIの繰り返し数	5

VGG-Styleモデルというのは、k,gを可変のパラメータとして、32kチャンネルのConv1層→32kチャンネルのConv×g層→64kチャンネルのConv×g層→128kチャンネルのConv×g層からなるネットワークです。論文ではg={2, 3, 4}、k={1, 2, 3}の9通りを総当たりで実験していましたが、今回は(g, k)=(2, 2), (4, 2)の場合を試しました。

結果は以下のようになりました。これは論文のTable 1.に対応します。確かにONIを用いた場合明らかによい結果となりました。

f:id:y60:20200805184435p:plain — VGG-Style・CIFAR-10

f:id:y60:20200805184908p:plain — VGG-Style・CIFAR-10

GAN

論文ではGANの安定性に関する実験も行っていました。学習したGANモデルをFIDで比較すると、Spectral Normalizationを上回る効果が確認できたという内容です。ONIの繰り返し数は1か2が良いそうです。分類の場合とは違い、繰り返し数を大きくしすぎると逆に悪化するようです。時間の都合でこれに関しては実装はしませんでした。

まとめ

勾配消失/爆発の対策としてネットワークの重みを直交化し、かつ学習中も常に直交性を保ちたいというモチベーションがありました。既存手法には、高計算コスト、不安定（主に0除算による）などの問題がありましたが、この論文はNewton’s iterationを用いて直交化する手法を提案し、これを解決しました。また、直交性を課す強さと表現能力の調節が可能になりました。分類タスクで既存手法を上回ったほか、GANの学習を安定にする効果も確認することができました。

*1:

詳しくは以下の通りです。

$\mathbf{y=Wx, WW^T=I}$ （または $\mathbf{W^TW=I}$ ）とします。

出力の勾配 $\frac{\partial L}{\partial \mathbf{y}}$ と入力 $\mathbf{x}$ の平均が0で共分散が $cov(\frac{\partial L}{\partial \mathbf{y}})=\sigma_2^2 \mathbf{I}, cov(\mathbf{x})=\sigma_1^2 \mathbf{I}$ と仮定すると、