SGDにおける「順番」の問題
この記事は atma Advent Calendar
の 12/1 分の記事です。大分遅くなってしまいましたがこの記事では Stochastic Gradient Descent における順番が与える影響とそれにまつわる論文をいくつか紹介したいと思います。
Stochastic Gradient Descent とはなにか
Stochastic Gradient Descent (SGD) は連続変数に対する最適化手法の一つで、ベースとなっているのは勾配法 (Gradient Descent) です。ここで以下の説明のためにいくつか用語を定義します。まず考える最適化問題は以下のように関数 $f_n$ の和で表される関数 $F$ の最適化です。
$$ {\rm minimize}_{x \in \mathbb{R}^d} F(x) := \sum_{n=1}^N f_n (x) $$
こういう関数は例えば機械学習・統計分析で頻出です。というのも $\sum$ の部分というのは期待値の計算とみなせるので、期待値を取ったなにかを最大化・最小化する問題もこの形式で記述できるためです。
また尤度関数の最大化問題も対数を取ってあげると上記の形式に帰結することができます。この場合 $N$ はデータの数に相当します。典型的なニューラルネットワークや線形モデルによる学習はすべてこれに該当する問題です。
最適化手法・勾配法では $i$ 番目の iteration において、今の最適化変数 $x_i$ の周りで計算した目的関数全体の勾配 $\nabla F$ を用いて, 全体の勾配の意味で最も下がる方向へ更新します。 $$ \begin{aligned} x_{i + 1} &= x_i - \alpha_i \nabla F(x_i) \\ &= x_i - \alpha_i \sum_{n=1}^N f_n(x_i) \end{aligned} $$
ここで $\alpha_i \in \mathbb{R}$ は $i$ 番目の更新で用いるステップサイズです。
SGD では全体の勾配の代わりに適当なインデックス $k \in \{1, 2, ... N\}$ を選んでそれを用いて更新を行います。すなわち以下のとおりです。
$$ x_{i + 1} = x_i - \alpha_i \nabla f_{k}(x_i) $$
更新時にはただひとつの勾配を計算すれば良いだけなので、GD に比べ一回あたりの更新の計算量が小さいのが特徴です。特に $N$ が膨大な値になる場合には GD では都度 $N$ 回の計算が必要になるのに比べて有利です。一方で、都度ただひとつのデータに対しての勾配しか計算しませんので、目的関数全体に対する正確な勾配との乖離が起こり、収束は遅くなります。
SGD における「順番」の問題
前置きが長くなってしまいましたが、ここからが本題です。先ほど SGD では毎回更新に使うデータの index $k$ を選ぶ、という話をしました。
ではこの $k$ はどのように選ぶのが良いのでしょうか。
まずナイーブに思いつくのは $k$ を $1, 2, \cdots, N, 1, 2, \cdots$ の用に順繰りに選んでいく方法です。この方法を以下では Cyclic と呼びます。
他のやり方は都度 ${1, 2, ..., N}$ を取りうる一様な分布から一つサンプルする方法です。この方法を以下では Random と呼びます。
$n$ ごとに重みをつけて偏りがある分布からサンプルする方法 (Importance Sampling) という方法もありますが以下では確率はいじらずに、一様な分布からサンプリングする方法に絞って議論していきます。
さらに上記の中間のような感じで、都度ランダムにサンプルはせず $N$ 回に1回だけ順番を入れ替え(シャッフル) して、そのあとはその順番通り選んでいく方法もあるでしょう。これを RR: Randomly-ReShuffle と呼びます。
実はやり方で収束速度が違う
この3つの方法、どれも各 $n$ が選ばれる確率は $i \to \infty$ では $1/N$ で同じなのですが、面白いことに収束レートは違うことが経験的に知られています。
Curiously Fast Convergence of some Stochastic Gradient Descent Algorithms 2009 では Logistic Regression に対して、上記3つの方法を用いた場合の収束レートを CCAT Dataset で実験した様子が報告されていて、Random < Cyclic < RR の順番で収束が早くなっています。更に Cyclic / Shuffle の場合には理論研究で得られている収束レート $t^{-1}$ よりも早く収束しています。
SGD に対する理論研究では $k$ の選び方は、前後の関係によらず独立にサンプルされたものである、という立場にたって分析されています。一方で Cyclic / Shuffle は前後に依存関係があるという違いがあり、この違いがレートの違いを生んでいるとして、選び方に対してより分析の余地があるのではないか? とコメントされています。
この問題に $f_n$ が二次関数で定義されるもので、一般の凸関数に対する議論ではないという限定的な課題設定ではあるものの、理論的な証明を与えたものが Why Random Reshuffling Beats Stochastic Gradient Descent です。
証明の気持ち
証明のアイディアは Cyclic や RR では $N$ 回ごとに必ず1度すべての関数 $f_n$ を使って計算が行われる、という点です。この性質を利用して各 $N$ 回の更新ごとの場所を $x_j$ のように定義して $x_j$ と $x_{j+1}$ の位置関係に対して成り立つ収束レートを元に、 Polyak-Rupper Averaging と呼ばれる方法で使われているような「過去 k 回の平均値を使う」ことで RR の方がレートとして早くなることを証明しています。
ライブラリでの実装はどうなっている?
順番に関する実装ですが、pytorch や keras では普通に実装するとデータセットを1度全部見るまでは順番が固定で取り出して、一周するごとに順番を並び替える用になっています。これはまさに Randomly Reshuffle ですね。
まとめ
- SGD には順番をどう選ぶか、という問題が存在しています。
- 特定の関数に対してではあるが RR が Random よりも早いという証明を与えた論文を紹介しました。
最後論文をガッツリ紹介しようと思っていたのですが、内容も詳しく追えておらず(最適化特有の数式 & 僕のパワーのなさ…) アイディアの一欠片ぐらいしかご紹介できませんでした。が、こういう問題意識や、その周辺の議論もあるんだなーというのを知っていただいて、あわよくば最適化おもろいな、と興味を持っていただけたのであれば幸いです。
