Nesterovの加速勾配法（Nesterov's Accelerated Gradient Method）

はじめに

勾配法の加速法にはいろいろ提案されていますが、ここでは前回書いた近接勾配法の加速にも使えるNesterovの加速法について書きます。

Momentum法

Nesterovの加速法の前に、類似した手法であるMomentum法についても簡単に見ておきます。
勾配法の更新式は $\boldsymbol{x}_k - \eta \bigtriangledown f(\boldsymbol{x}_k)$ ですが、Momentum法では次のように更新します。

\begin{align}
\boldsymbol{x}_{k+1} &= \boldsymbol{x}_k - \eta \bigtriangledown f(\boldsymbol{x}_k) + \alpha \Delta \boldsymbol{x}_k
\end{align}

ここで$\alpha$は$0 < \alpha < 1$です。1ステップ前の更新方向$\Delta \boldsymbol{x}_k = \boldsymbol{x}_k - \boldsymbol{x}_{k-1}$にも更新する、すなわち慣性（Momentum）を持たせています。
この式を展開すると、

\begin{align}
\boldsymbol{x}_{k+1}
 &= \boldsymbol{x}_k - \eta \sum_{l=0}^{k} \alpha^l \bigtriangledown f(\boldsymbol{x}_{k-l})
\end{align}

となり、過去の勾配の$\alpha$による指数減衰和方向に更新していることになります。
いい感じに動けば、最近の更新で何度更新しても同様の方向への勾配成分があったときにはその方向への更新が加速され、逆に振動成分は抑制されそうです。が、逆に慣性項が収束に悪さをするときもありそうです。
では、本題のNesterovの加速法に進みます。

Nesterovの加速法（Nesterov's Accelerated Gradient Method）

Nesterovの加速法では次のように更新します。

\begin{align}
\boldsymbol{x}_{k+1}
 &= \bar{\boldsymbol{x}}_k - \eta \bigtriangledown f(\bar{\boldsymbol{x}}_k) \\
\bar{\boldsymbol{x}}_k &= \boldsymbol{x}_k + \gamma_k \Delta \boldsymbol{x}_k \\
\gamma_k &= \frac{\rho_{k-1} - 1}{\rho_k} \\
\rho_k &= \frac{1 + \sqrt{1 + 4\rho_{k-1}^2}}{2}
\end{align}

Momentum法と比較しやすいように書くと

\begin{align}
\boldsymbol{x}_{k+1}
 &= \boldsymbol{x}_k - \eta \bigtriangledown f(\bar{\boldsymbol{x}}_k) + \gamma_k \Delta \boldsymbol{x}_k
\end{align}

となり、違いは勾配を計算する位置にあります。
更新を(1)二つ目の式の「慣性による更新」と(2)一つ目の式の「勾配による更新」の2ステップに分けて考えると、(2)の勾配による更新において、Momentum法では(1)の更新前の点$\boldsymbol{x}_k$における勾配で更新してしまっているのに対して、Nesterovでは勾配法による更新の起点が(1)の更新後の点$\bar{\boldsymbol{x}}_k = \boldsymbol{x}_k + \gamma_k \Delta \boldsymbol{x}_k$における勾配で更新しています。
すなわち、Nesterovの加速法は（広い意味でのMomentum法の一種ですが、上述のシンプルな）Momentum法の改良版と言えます。

Nesterov法は勾配計算の起点を修正しているだけでなく、$\gamma_k$と$\rho_k$の更新式を上記の式のように構成しているところがミソで、これにより$\beta$-平滑な凸関数$f$に対して、収束レートが

\begin{align}
f(\boldsymbol{x}_k) - f(\boldsymbol{x}^*) \leq \frac{2\beta \parallel \boldsymbol{x}_1 - \boldsymbol{x}^* \parallel^2}{k^2}
\end{align}

となることが示されます（[1] p.277-282, [2]）（加速しない場合の収束レートは分母が$k$）。
また、実用上は、Momentum法では$\alpha$がハイパーパラメータになっていたのに対して、Nesterov法では$\gamma_k$は更新式で与えられる利点もあります。$\rho_0=1$からスタートすることで$\gamma_1=0$となり、最初の1ステップはただの勾配法による更新となります。

「加速」の直感的理解

Momentum法と同様に展開すると

\begin{align}
\boldsymbol{x}_{k+1}
 &= \boldsymbol{x}_k - \eta \sum_{l=0}^{k} \left(\prod_{m=0}^{l-1} \gamma_{k-m} \right) \bigtriangledown f(\bar{\boldsymbol{x}}_{k-l})
\end{align}

となります。更新式に従って$\gamma_k$と$\rho_k$を計算していくと、$k \rightarrow \infty$で$\gamma_k \rightarrow 1$です。そこで理解のために$\gamma=1$としてしまえば

\begin{align}
\boldsymbol{x}_{k+1}
 &= \boldsymbol{x}_k - \eta \sum_{l=0}^{k} \bigtriangledown f(\bar{\boldsymbol{x}}_{k-l})
\end{align}

となり、勾配の累積和方向に更新していることになります。たとえばステップ幅$\eta$が小さすぎるなどで、更新しても更新しても求まる勾配ベクトルが同じだったとすると、$k$回目の更新での更新幅は$k\eta$、累積更新幅は$k(k+1) \eta / 2$となります。加速をしなかった場合の更新幅は$\eta$、累積更新幅は$k \eta$なので、加速が適切な状況では収束レートのオーダーが$1/k$から$1/k^2$に改善するのも頷けます。実際、加速が進んでいるときには目的関数の減少幅は線形に増えていき、あるところから減少幅が小さくなっていくような挙動が見られます。
ちゃんとした証明は参考文献（[1] p.277-282, [2]）をご参照頂きたいのですが、証明からは「加速」が直感的に理解できなかったので、このような大雑把な解釈をしてみました。

リスタート

Nesterovの加速法は$k$が大きくなるにつれて、このようにアグレッシブになっていくので、更新し続けると目的関数が振動したり、かえって収束性が悪化することがあります。
そこで、停止条件を満たしたらリスタート（$\rho_k$に$1$をセット）する方法が提案されています（[1] p.278, [3]）。文献[3]の実験のグラフだけでも見てみる価値があり、リスタートによって劇的に収束が早くなっています。
$f(\boldsymbol{x})$を最小化するときの停止条件としては

• $f(\boldsymbol{x}_k) > f(\boldsymbol{x}_{k-1})$
• $\bigtriangledown f(\bar{\boldsymbol{x}}_{k-1})^T (\boldsymbol{x}_k - \boldsymbol{x}_{k-1}) > 0$

などが提案されています。一つ目の停止条件は「目的関数が増加したら」というもので、二つ目の停止条件は(1)慣性による更新と(2)勾配による更新の(2)では勾配降下方向に更新するので、慣性項がなければ更新ベクトルとの内積は負になるはず、というところから来ています。これらはどちらかだけ用いるのでも良いようです。近接勾配法で$F(\boldsymbol{x})=f(\boldsymbol{x})+g(\boldsymbol{x})$を最小化する場合は、$\bar{\boldsymbol{x}}_{k-1}$から$\boldsymbol{x}_k = {\rm prox}_{\eta g}( \bar{\boldsymbol{x}}_{k-1} - \eta \bigtriangledown f(\bar{\boldsymbol{x}}_{k-1}) )$への更新方向$\boldsymbol{x}_k-\bar{\boldsymbol{x}}_{k-1}$（の符号を反転したもの）を$\bigtriangledown F(\boldsymbol{x}_{k-1})$の近似として用いることができます（[1] p.277-282）。
前回書いた記事のバックトラッキング法によるステップ幅の探索アルゴリズムに、Nesterovの加速法とリスタートを組み込むと次のようになります。前回に引き続き近接勾配法で書いています。$\hat{f}_\eta(\boldsymbol{x};\boldsymbol{y})$の定義については前回の記事をご参照ください。

\begin{align}
&\bar{\boldsymbol{x}}_0 = \boldsymbol{x}_0, \; \eta = \eta_0, \; \rho_0 = 1 \\
&{\rm for} \; k=1 \; {\rm to} \; K \; {\rm do} \\
&\quad {\rm while \; not \; broken \; do} \\
&\qquad \boldsymbol{x}_k = {\rm prox}_{\eta g}( \bar{\boldsymbol{x}}_{k-1} - \eta \bigtriangledown f(\bar{\boldsymbol{x}}_{k-1}) ) \\
&\qquad {\rm if} \; f(\boldsymbol{x}_k) \leq \hat{f}_\eta(\boldsymbol{x}_k;\bar{\boldsymbol{x}}_{k-1}) \; {\rm then}\\
&\qquad \quad {\rm break} \\
&\qquad {\rm else} \\
&\qquad \quad \eta \leftarrow \beta \eta \\
&\qquad {\rm end \; if} \\
&\quad {\rm end \; while} \\
&\quad {\rm if}\; F(\boldsymbol{x}_k) > F(\boldsymbol{x}_{k-1}) \;{\rm or}\; (\boldsymbol{x}_k-\bar{\boldsymbol{x}}_{k-1})^T (\boldsymbol{x}_k - \boldsymbol{x}_{k-1}) < 0 \;{\rm then} \\
&\qquad \eta = \eta_0 \\
&\qquad \rho_{k-1} = 1 \\
&\quad {\rm end \; if} \\
&\quad \rho_k = \frac{1 + \sqrt{1 + 4\rho_{k-1}^2}}{2} \\
&\quad \gamma_k = \frac{\rho_{k-1} - 1}{\rho_k} \\
&\quad \bar{\boldsymbol{x}}_k = \boldsymbol{x}_k + \gamma_k \Delta \boldsymbol{x}_k \\
&{\rm end \; for}
\end{align}

実際のところ・・

序盤の数回の更新で$\eta$が調整されると、その後の多くの反復では内部のバックトラッキング法によるステップ幅の探索は1回で抜けることが多いです。上記アルゴリズムでは$f(\boldsymbol{x}_k)$と$f(\bar{\boldsymbol{x}}_k)$両方の関数評価が必要になります。関数評価が重いモデルの学習時にはここがボトルネックになってきます。
私の実験では、単に$F(\boldsymbol{x}_k) < F(\boldsymbol{x}_{k-1})$だけをチェックしてループを抜ける方法が上記アルゴリズムと同程度の反復で収束し、一反復あたりの計算時間が短くなり、動作としても安定していました。$\eta$を$\eta_{\rm MIN}$より小さくしても目的関数を小さくできない（(1)の慣性による更新が大きすぎて、(2)の勾配法による更新をしても(1)と(2)トータルの更新で目的関数が増加してしまう）場合、$\boldsymbol{x}_k=\boldsymbol{x}_{k-1}$（更新しない）としてループを抜けて加速法のリスタートをします。アルゴリズムとしては次のようになります。

\begin{align}
&\bar{\boldsymbol{x}}_0 = \boldsymbol{x}_0, \; \eta = \eta_0, \; \rho_0 = 1 \\
&{\rm for} \; k=1 \; {\rm to} \; K \; {\rm do} \\
&\quad {\rm while \; not \; broken \; do} \\
&\qquad \boldsymbol{x}_k = {\rm prox}_{\eta g}( \bar{\boldsymbol{x}}_{k-1} - \eta \bigtriangledown f(\bar{\boldsymbol{x}}_{k-1}) ) \\
&\qquad {\rm if} \; F(\boldsymbol{x}_k) < F(\boldsymbol{x}_{k-1}) \;{\rm then}\\
&\qquad \quad {\rm break} \\
&\qquad {\rm else \; if} \; \eta < \eta_{\rm MIN} \;{\rm then}\\
&\qquad \quad \boldsymbol{x}_k = \boldsymbol{x}_{k-1} \\
&\qquad \quad {\rm break} \\
&\qquad {\rm else} \\
&\qquad \quad \eta \leftarrow \beta \eta \\
&\qquad {\rm end \; if} \\
&\quad {\rm end \; while} \\
&\quad {\rm if}\; F(\boldsymbol{x}_k) \geq F(\boldsymbol{x}_{k-1}) \;{\rm then} \\
&\qquad \eta = \eta_0 \\
&\qquad \rho_{k-1} = 1 \\
&\quad {\rm end \; if} \\
&\quad \rho_k = \frac{1 + \sqrt{1 + 4\rho_{k-1}^2}}{2} \\
&\quad \gamma_k = \frac{\rho_{k-1} - 1}{\rho_k} \\
&\quad \bar{\boldsymbol{x}}_k = \boldsymbol{x}_k + \gamma_k \Delta \boldsymbol{x}_k \\
&{\rm end \; for}
\end{align}

参考文献

[1] 金森敬文, 鈴木大慈, 竹内一郎, 佐藤一誠. 機械学習のための連続最適化（機械学習プロフェッショナルシリーズ）. 講談社. 2016.
[2] Sebastien Bubeck. ORF523: Nesterov's Accelerated Gradient Descent (I'm a bandit). 2013
[3] Brendan O'donoghue, and Emmanuel Candes. Adaptive restart for accelerated gradient schemes. 2012.