Momentum Residual Neural Networks

[e4exp]

• https://arxiv.org/abs/2102.07870v2
• 2021

バックプロパゲーションを用いた深い残差ニューラルネットワーク（ResNets）の学習には、ネットワークの深さに応じて線形に増加するメモリコストがかかります。 この問題を回避する方法として、可逆的なアーキテクチャを使用することが挙げられる。 本論文では、モーメンタム項を追加することで、ResNetの順方向ルールを変更することを提案する。 結果として得られるネットワーク、モメンタム残差ニューラルネットワーク（Momentum ResNets）は反転可能である。 これまでの反転可能なアーキテクチャとは異なり、既存のResNetブロックのドロップイン置き換えとして使用することができます。 我々は、Momentum ResNetsが無限小ステップサイズの領域において、2階常微分方程式（ODE）として解釈できることを示し、Momentum ResNetsの表現能力がどのようにして漸進的に向上するかを正確に特徴付ける。 我々の分析によると、運動量ResNetsは任意の線形マッピングを乗算係数まで学習できるが、ResNetsは学習できない。 固定点への収束が求められる最適化学習において、既存の反転可能なアーキテクチャが失敗する中で、我々の手法が成功することを理論的にも経験的にも示した。 CIFARとImageNetにおいて、モメンタムリセットはリセットと同等の精度を持ちながら、メモリ使用量が非常に少ないことを示し、事前に学習されたモメンタムリセットはモデルの微調整に有望であることを示す。

[e4exp]

1. はじめに

問題設定。 深層学習の特定のインスタンス（LeCunら、2015; Goodfellowら、2016）として、残差ニューラルネットワーク（Heら、2016、ResNets）は、非常に深い表現に起因する大きな経験的成功を達成しており、その拡張は、実データセット上の芸術の状態を上回り続けている（Kolesnikovら、2019; Touvronら、2019）。 深層学習タスクのほとんどは、いくつかの状況でメモリが実用的なボトルネックとなるグラフィックスプロセッシングユニット（GPU）を使用します（Wangら、2018; Pengら、2017; Zhuら、2017）。 実際、ディープアーキテクチャーの最適化に用いられるバックプロパゲーションでは、ネットワークの評価（フォワードパス）中に各層で値（活性化）を保存する必要があります。 したがって、ディープアーキテクチャーの深さは、利用可能なメモリの量によって制約されます。 この論文の主な目的は、新しいモデルであるMomentum ResNetsの特性を探ることである。 このモデルは、反転可能であることによって、これらのメモリの問題を回避している。 つまり、層nでの活性化は、層n + 1での活性化から正確に復元される。 このネットワークは、ResNetのフォワード・ルールの修正に依存しており、実際には正確に反転可能である。 ResNet（残差構築ブロック）のフィードフォワード関係を考える代わりに

を繰り返して、その運動量に対応するものを定義します。

ここで、fはパラメータ化された関数、vは速度項、γ∈[0, 1]は運動量項である。 これにより、次の図に示すように、ネットワークのダイナミクスが根本的に変化します。 既存の可逆モデルとは対照的に、Momentum ResNetsは、残差ブロックをビルディングブロックとして使用するあらゆるディープアーキテクチャにシームレスに統合することができます（参照：セクション3)

貢献しています。 モーメンタム残差ニューラルネットワーク(Momentum ResNets)は、ResNetのフォワードルールを簡単に変更したもので、アーキテクチャに制約がなく、完全に反転可能な新しい深層モデルです。 我々は、モーメンタム項γを変更することにより、モーメンタムResNetsの必要メモリ量が任意に減少することを示し（セクション3.2）、従来のResNetsのドロップイン置き換えとして使用できることを示す。 また、理論面では、他の可逆モデルでは収束しないような最適化学習の場面でも、Momentum ResNetsが容易に使えることを示す（セクション3.3）。 また、モメンタムリセットの近似能力を、連続的な極限では2次のODEとして見ることができることを調べる（セクション4）。 まず、命題3で、Momentum ResNetsは、1次のニューラルODEよりも厳密に大きなクラスの関数を表現できることを示す。 次に、線形のケースを研究することで、より詳細な洞察を与える。 ここでは、線形残差関数を持つモメンタムリセットが普遍的な近似能力を持つことを定理1で正式に証明し、このようなモデルの表現可能なマッピングのセットが、モメンタム項γの増加に伴ってどのように成長するかを正確に定量化する。 この理論的な結果は、Momentum ResNetsの表現能力を理論的に分析するための第一歩となる。

最後の貢献は、様々な学習タスクにおけるMomentum ResNetsの実験的検証である。 我々はまず、ResNetsが分離できなかった点群をMomentum ResNetsが分離することを示す（セクション5.1）。 また、画像データセット(CIFAR-10, CIFAR-100, ImageNet)において、Momentum ResNetsはResNetsと同様の精度を持ち、より小さなメモリコストで済むことを示す(セクション5.2)。 また、事前に学習されたモデルのパラメータは、わずか数エポックの学習で同等の精度を達成するMomentum ResNetに容易に移植できることを示している。 我々は、事前に学習されたMomentum ResNetを得るこの方法は、メモリストレージがボトルネックとなる新しいデータに対してネットワークを微調整する際に重要であると主張する。 Pytorchパッケージのtransform(model)メソッドは、torchvision ResNetモデルを受け取り、わずかな修正で同様の精度を達成するMomentum対応モデルを返します。 また、Momentum ResNetsがRevNetsよりも優れた性能を発揮することを確認することで、最適化のための学習という設定で、我々の理論的な発見を実験的に検証しています（Gomez et al.、2017）。

我々のコードはhttps: //github.com/michaelsdr/momentumnetで公開されています。

[e4exp]

2. 背景と過去の作品

バックプロパゲーション

バックプロパゲーションは、スカラー値の関数の勾配を計算する際に選択される手法である。 バックプロパゲーションは、計算グラフを逆方向に走査する連鎖法則を用いて動作します(Bauer, 1974)。 逆モードの自動微分法としても知られています（Baydin et al., 2018; Rumelhart et al., 1986; Verma, 2000; Griewank & Walther, 2008）。 計算コストは、関数自体を評価する場合と同様です。 さらなる仮定をせずにニューラルアーキテクチャーを通じてグラデーションをバックプロパゲートする唯一の方法は、フォワードパス中のすべての中間活性化を保存することです。 これは、Pytorch（Paszke et al., 2017）、Tensorflow（Abadi et al., 2016）、JAX（Jacobsen et al., 2018）などの一般的な深層学習ライブラリで使用されている方法です。 このメモリストレージを削減する一般的な方法は、チェックポイントを使用することです： 活性化はいくつかのステップでのみ保存され、他のものはバックワードパスで必要になると、これらのチェックポイントの間で再計算されます（例えば、Martens & Sutskever (2012)）。

リバーシブルなアーキテクチャ。

しかし、最近ではアクティベーションを保存せずにバックプロパゲーションを行うことができるモデルが開発されています。 これらは2種類のアプローチに基づいています。 1つ目は離散的で、活性化nと活性化n+1を結びつけるルールを簡単に反転させる方法を見つけることに依存している（Gomez et al., 2017; Chang et al., 2018; Haber & Ruthotto, 2017; Jacobsen et al., 2018; Behrmann et al., 2019）。 この方法では、バックワードパス中にオンザフライでアクティベーションを再計算することが可能です： アクティベーションは保存する必要はありません。 しかし、これらの方法は、性能の良い非可逆モデルを可逆モデルに移行する簡単な方法がない制限されたアーキテクチャに依存しているか、活性化を後ろ向きに再計算する際に高速な反転スキームを提供していません。

これに対して、我々の提案は、既存のあらゆるResNetに適用でき、容易に反転させることができます。 第2の種類のアプローチは、連続的で常微分方程式（ODE）に依存しており、ResNetは連続的な力学系として解釈されます（Weinan, 2017; Chen et al., 2018; Teh et al., 2019; Sun et al. これにより、ODEからの理論的および数値的な進歩を深層学習に取り込むことができる。 これらのモデルはしばしばニューラルODEと呼ばれ（Chen et al., 2018）、Adjoint Sensitivity法（Pontryagin, 2018）を用いて、ODEを時間的に逆算して解くことで学習することができます。 この戦略は、ODEソルバーの操作を通じて逆モードの自動微分を行うことを回避し、O(1)のメモリフットプリントにつながります。 しかし、既存の残差アーキテクチャのニューラルODE対応を定義することは簡単ではありません。 ODEブロックの最適化は、非自明な時間離散化を必要とする無限次元の問題であり、ニューラルODEの性能は、ODEの数値積分器に依存します（Gusak et al.、2020）。 また、ODEは安定性の問題から、常に数値的に逆走することはできません。 システムを逆走させると、数値誤差が発生し、蓄積されていきます（Gholami et al. したがって、実際には、ニューラルODEは標準的な深層学習の設定ではめったに使用されない。 とはいえ、最近の作品（Zhang et al., 2019; Queiruga et al., 2020）では、ニューラルアーキテクチャにODEブロックを組み込み、CIFARでResNetsに匹敵する精度を達成している。

表現能力。

このようなモデルの表現能力を研究することは、実世界のデータに対する性能に関する洞察を与える意味でも重要です。 つまり、コンパクトセット上の任意の連続関数は、1層のフィードフォワード・フルコネクテッド・ニューラルネットワークで一様に近似できるということです。 しかし、ニューラルODEの表現能力は限られています。 Tehら（2019）は、データのベクトルフィールドを異次元空間のゼロと連結することで高次元の点を持ち上げることを提案し、その結果得られた拡張ニューラルODE（ANODE）が、画像分類と玩具実験において、より低い損失と優れた一般化を達成することを示している。 Liら(2019)は、ODE-Netの出力が終端族の要素で構成されている場合、p＜＋∞のL pノルムでの収束に対して普遍的な近似能力が得られることを示しているが、これでは不十分である(Teshima et al., 2020)。 本研究では、ResNetの前方反復から得られるODEのL∞ノルムでの表現能力を考える。 さらに、Zhangら(2020)は、ODEの次元を2倍にすると普遍的な近似器が得られることを証明したが、この結果は我々の知る限り深層学習には適用されていない。 本研究では、連続極限において、我々のアーキテクチャがNeural ODEよりも優れた表現能力を持つことを示す。 また、線形の場合にその普遍性を証明します。

ディープネットワークにおけるモーメンタム

最近のいくつかの作品（He et al., 2020; Chun et al., 2020; Nguyen et al., 2020; Li et al., 2018）では、深層アーキテクチャにおけるモメンタムを探索している。 しかし、これらの手法は、そのアーキテクチャと目的において我々のものとは異なります。 Chunら（2020）は、反復がResNetに対応していない最適化問題を解くためにモメンタムを導入しています。 Nguyenら（2020）（resp.Heら（2020））は、消失勾配の問題を緩和するために重みが結ばれているRNN（ResNetとは異なる）の場合にモメンタムを追加する（resp.key and query encoder layersをリンクする）。 Liら（2018）は、線形層が結ばれていて、対称な定型行列である特定のケースを検討しています。 特に、言及されたアーキテクチャのどれもが反転可能ではなく、これは私たちの方法の主な資産の1つです

2次モデル

運動量項を加えることが、2次のODEを積分するためのオイラー積分スキームに対応することを示します。 最近提案されたいくつかのアーキテクチャ(Norcliffe et al., 2020; Rusch & Mishra, 2021; Lu et al., 2018; Massaroli et al., 2020)も2階微分方程式が動機となっている。 Norcliffeら（2020）は2次の力学系をモデル化するために2次の力学を導入していますが、我々のモデルは連続的な極限における離散的な方程式のセットに対応しています。 また、我々の手法では、ニューラルネットワークはxにのみ作用するため、モメンタムによって次元が2dに増加するものの、フォワードパスの計算負荷は次元dのResNetと同じです。 Rusch & Mishra (2021)は2次のRNNを提案していますが、我々の手法ではResNetを扱います。 最後に、Luら（2018）のLM-ResNetの定式化は、どちらも2階のODEを導くにもかかわらず、我々のフォワードパス（xn+1 = xn + γvn + (1 - γ)f(xn, θn)）とは異なる。 重要なのは、これらの2階の定式化はどれも反転可能ではないということです。

表記

d∈N ∗ , R d×d , GLd(R), DC d (R)で，Cで対角化可能な実数行列，反転可能な行列，実数行列の集合を表す。

3. Momentum ResNetを紹介します。

これは、任意のResNetを必要メモリの少ないモデルに簡単に変換したもので、連続極限では2次のODEとして見ることができます。

3.1. 勢いのあるResNet ResNetの方程式に勢いのある項を加える。

ここで、(vn)nは、R dのある値v0で初期化された速度であり、γ∈[0, 1]は、いわゆる運動量項である。 このアプローチは、fが最小化すべき関数の勾配である、運動量を伴う勾配降下アルゴリズム（Ruder, 2016）を一般化したものです。

初期速度とモメンタム項

本論文では、単純な関係を通じてx0に依存する初期速度v0を考慮する。 最も単純な選択肢は、v0 = 0またはv0 = f(x0, θ0)とすることである。 セクション4では、v0とx0の間のこの依存関係が、Momentum ResNetsが表現できるマッピングのセットに影響を与えることを証明する。 パラメータγは、Momentum ResNetがResNetからどれだけ発散するか、また、メモリの節約量を制御する。 γが0に近いほど、Momentum ResNetsはResNetsに近くなるが、メモリの節約量は少なくなる。 我々の実験では、γ=0.9を使用しているが、これは様々なアプリケーションでうまく機能することがわかっている。

反転させる。

手順(2)が反転するのは

を用いることで、Momentum ResNetのバックワード・パスにおいて、活性化をその場で再構成することができる。 実際には、ダイナミクスを正確に反転させるためには、(2)のγによる有限精度の乗算によって失われた情報を効率的に保存する必要があります。 我々は、この可逆的な乗算を行うために、Maclaurin et al. このアルゴリズムは、情報バッファ、つまり、乗算が可逆的になるように、各反復で失われるビットを格納する整数を維持することで構成されています。 この手順については、付録Cで詳しく説明します。 なお、学習可能なマッピングfの追加により、浮動小数点精度の小さな損失が常に発生しますが、実際にはこれが問題になることはありませんでした。

ドロップイン置換。

我々のアプローチは、既存のResNetを可逆的なものに変えることを可能にします。 言い換えれば、あるResNetは、各層の構造を変えることなく、Momentum対応のものに変換することができます。 例えば、ResNet-152(He et al., 2016)を考えてみましょう。 これは4つの層（深さ3、8、36、3）で構成されており、4つの層の前進式(1)を(2)に変更することで、簡単にMomentum ResNetの対応物に変えることができます。 それ以上の変更は必要なく、Momentum ResNetは全く同じパラメータを入力とし、そのまま置き換えることができます。 これは、他の可逆的なモデルの場合とは異なります。 Neural ODEs (Chen et al., 2018) は、入力として連続的なパラメータを取ります。 iResNets (Behrmann et al., 2019) は、重みのスペクトルノルムが制約付き最適化を必要とするため、プレーンなSGDでは学習できません。 i-RevNets (Jacobsen et al., 2018) および RevNets (Gomez et al., 2017) は、残差ブロックごとに独自のパラメータを持つ2つのネットワークを学習する必要があり、入力を畳み込みチャネルに分割し、ResNetsの半分の深さになります： これらは入力として同じパラメータを取りません。表1は，可逆残差アーキテクチャの特性をまとめたものです． RevNetsとMomentum ResNetsの違いについては、3.3節と5.3節でさらに詳しく説明します。

3.2. メモリコスト

各層に完全なデータを格納する代わりに、γによる各乗算で失われるビットのみを格納する必要があります（「intertibility」を参照）。 深さkのアーキテクチャの場合、これは各サンプルに対してlog2(( 1/γ)^k)の値を保存することに相当します（γが1に近い場合はk(1-γ)/ln(2)）。 ここでは、活性化値の保存がメモリのボトルネックになる2つの状況を考えてみます。 まず、f(x, θ) = WT 2 σ(W1x + b), x∈R^d, θ = (W1, W2, b), ここで、W1, W2∈R^{p×d}, b∈R^p, 深さk∈Nとするトイ・フィードフォワード・アーキテクチャを考える。 学習セットは、n個のベクトルx^1, ..., x^n∈R^dで構成されます。 ResNetsでは、ネットワークの各層に、ネットワークの重みと、トレーニングセットのすべての活性化の値を格納する必要がある。 合計すると、必要なメモリは、反復ごとにO(k × d × nbatch)となります。 Momentum ResNetsの場合、γが1に近い場合、必要なメモリはO((1 - γ) × k × d × nbatch)となります。 これは、深さkのメモリ依存性が、運動量γを変えることで任意に軽減されることを証明しています。 このメモリ節約は、図2に示すように、実際に確認されています。

別の例として、ImageNetの分類に使用できるResNet-152(He et al., 2016)を考えてみましょう(Deng et al., 2009)。 conv4 x」と名付けられた層の深さは36で、40M個のパラメータを持っていますが、活性化を保存するには50倍のパラメータを保存する必要があります。 ここでは活性化の保存が主な障害となるため、γを1に近づけることで、この層の必要メモリ量を任意に減らすことができる。

3.3. 勢いの役割

(2)でγを0にすると、ResNetが回復する。 したがって、Momentum ResNetsはResNetsの一般化である。 γ -→ 1のとき、f → 1 1-γ fをスケーリングして(2)のシンプレクティック・スキーム(Hairer et al., 2006)を得て、他の人気のある反転可能なニューラル・ネットワークの特別なケースを回復することができる。 RevNets (Gomez et al., 2017)とHamiltonian Networks (Chang et al., 2018)である。 RevNetは次のように反復します。

ここで、φとψは学習可能な2つの関数である。 このようなアーキテクチャの有用性は、タスクに依存する。 RevNetsは分類と回帰で成功を収めている。 しかし、我々は、RevNetsがいくつかの設定では機能しないことを主張する。 例えば、穏やかな仮定の下で、パラメータが各層で同じである場合、RevNetの反復は魅力的な固定点を持っていません：θn = θ, 我々は、Ψ(v, x) = (v + ϕ(x, θ), x + ψ(v + ϕ(x, θ), θ')) で、(4)を(vn+1, xn+1) = Ψ(vn, xn)と書き直します。

命題1（固定点の不安定性）。

(v ∗ , x∗ ) をRevNetの反復（4）の固定点とする。x ∗ (resp. v ∗ ) においてφ (resp. ψ) が微分可能であり、ヤコビアン行列A (resp. B) ∈ R d×d であるとする。 (v ∗ , x∗ )でのΨのジャコビアンは、

です。AとBが反転可能であれば、｜λ｜≧1かつλ neq 1となるλ∈Sp (J(A, B))が存在する。

このことから、(v ∗ , x∗ )は安定した固定点ではないことがわかります。 つまり、(4)によれば、xn が収束すれば vn も収束し、その極限が固定点となるはずです。 前述の命題は、それが不可能であることを示しています。 この結果は、ネットワークの反復が収束することが期待される問題では、RevNetsの性能が低いことを示唆しています。 例えば、5.3節の実験で示されているように、最適化を学習するためにリバーブル・ダイナミクスを使用すると、このようなことが起こります(Maclaurin et al., 2015)。 これに対し、提案手法では、運動量項γが厳密に1以下であれば、固定点に収束することができる。

備考

命題1は連続的な対応を持つ。 実際、連続極限では、(4)はv˙ = ϕ(x, θ), x˙ = ψ(v, θ0 )と書く。 これに対応する(v ∗ , x∗ )のヤコビアンは

です。 この行列の固有値は、ABの固有値の平方根であり、すべての実部が<0になることはない（連続の場合の安定性の問題と同様

3.4. Momentum ResNets as continuous mode

Neural ODEs: 一次のODEとしてのResNets。 初期条件x0(ResNetの入力)を持つResNetの方程式(1)は、x(0)=x0のODE x˙=f(x, θ)の離散化されたオイラースキームと見なすことができます。 Tを時間地平とすると、ニューラルODEは入力x(0)を出力x(T)にマッピングし、Chenら(2018)と同様に、損失L(x(T), θ)を最小化することで学習されます。

二階のODEとしての運動量ResNets。

ε = 1 / 1-γ とします。これは、微分方程式εx¨ + ˙x = f(x, θ)のステップサイズ1のVerlet積分法(Hairer et al., 2006)に対応しています。 このように、ResNetsが1次のODEの離散化であるのと同じように、Momentum ResNetsは2次のODEの離散化であると考えられる。 図3はこれらのアイデアをまとめたものです。

[e4exp]

結論

本論文では、ResNetに比べて大幅に少ないメモリ・フットプリントで動作する新しい反転可能な残差ニューラル・ネットワークであるMomentum ResNetsを紹介します。 既存の反転可能なアーキテクチャとは対照的に、ResNetのフォワードルールを簡単に変更することで可能になりました。 この単純さは、理論的な利点（より優れた表現能力、線形力学の扱いやすい分析）と実用的な利点（ドロップインでの置き換え、モデルの微調整のための速度とメモリの改善）の両方を提供します。 運動量ResNetsは、ResNets(γ=0)とRevNets(γ=1)の間を補うもので、ニューラルODEの自然な二次拡張です。 そのため、非同型のダイナミクスや収束するイテレーションを捉えることができます。 この論文で示されているように、後者は既存の反転可能な残差ネットワークでは不可能ですが、最適化のための学習という設定では重要です。