バッチサイズとステップあたりの学習速度の関係

要約

　バッチサイズとステップあたりの学習速度は比例関係にある(？)ため、強化学習の高速化としてバッチサイズを小さくすることは意味がない可能性がある。

前書き

　前回はLR Range Testによる学習率の決定法について書いた。これをもとに複数のバッチサイズで教師あり学習の実験を行い、そこからAlphaZeroアルゴリズムでのバッチサイズを検討したい。

実験

実験設定

データ
floodgate2015年~2018年の棋譜から両対局ソフトのレートが2800以上、かつ手数が50手以上のもの計10,939,651局面(そのうち9割を学習データ、残りを検証データとして使用)
バッチサイズ
4096, 512, 64, 32の4通り
損失
Policy損失とValue損失の1:1での和
学習率減衰
1エポックごとに検証データ全てを使って損失を計算し、前回までの検証損失の最小値より小さくなっていなければ学習率を0.9倍
学習終了
最小値を更新できなかったことが6回続いたタイミングで学習を打ち切り
学習率の初期値
各バッチサイズについて学習率の初期値をLR Range Testにより決定

バッチサイズ	学習率の初期値
4096	0.03
512	0.025
64	0.015
32	0.01

実験結果

　横軸にエポック数を取ったグラフを次に示す。

f:id:tokumini:20190515112200p:plain

　各バッチサイズについて損失最小値を次に示す。

バッチサイズ	エポック	損失	Policy損失	Value損失
4096	44	2.18	1.76	0.422
512	38	2.14	1.73	0.413
64	40	2.12	1.71	0.410
32	62	2.11	1.71	0.405

　バッチサイズは小さい方が汎化性能が上がるという話¹と一致した(この話には反論もあるようだが²)。バッチサイズ32のときだけ損失が最小となる際のエポック数が大きいが、グラフを見ると40エポックほどでもほとんど同程度の損失値になっており、同程度の性能に至るまでに必要なエポック数は同等と見なすことができると考えられる。

　バッチサイズによらず同じエポック数ではほぼ同等の性能と仮定すると、バッチサイズを $\frac{1}{n}$ 倍したとき1エポックあたりのステップ数は $n$ 倍になるため、ステップ数あたりの学習速度は $\frac{1}{n}$ 倍になるということになる。わかりやすく横軸にステップ数を取ったグラフを次に示す。

f:id:tokumini:20190515112216p:plain

考察:強化学習における学習時間

　強化学習でのステップ数と教師あり学習でのステップ数を同一視して良いものかはわからないが、とりあえず同一視できるものとする。今回の実験からは、バッチサイズを $\frac{1}{n}$ 倍した場合 $n$ 倍のステップ数が必要になる可能性が示唆されている。

　AlphaZeroのようなActorとLearnerを並列に動かす手法では、バッチサイズとデータ生成速度のバランスが重要であると勝手に思っている(もちろんActorの数は可能ならば多ければ多いほど良い³)。AlphaZeroは第1世代TPUを5000基用いてデータ生成を行っているが、たとえばこれの $\frac{1}{100}$ 倍である50基しかマシンを用意できなかった場合に、バッチサイズを4096のまま学習させると同じデータを過剰に何度も学習することになりかねない。おそらくバッチサイズもActorの減る割合に合わせて $\frac{1}{100}$ 倍するのが良いのではないかと思われる(ここが線形かどうかは不明)。

　しかし今回の結果ではバッチサイズを $\frac{1}{100}$ 倍すると同じ性能に至るまで100倍のステップ数がかかることになる。バッチサイズによらず1ステップあたりにかかる時間が一定であるとすると100倍の時間がかかることになる。

　結局、次の三つの仮定

Actorの数を $\frac{1}{n}$ 倍するならば、バッチサイズを $\frac{1}{n}$ 倍しなければならない
バッチサイズを $\frac{1}{n}$ 倍するならば、ステップ数を $n$ 倍しなければならない
バッチサイズによらず1ステップあたりにかかる時間は一定

を認めるならば、Actor-Learner型の手法で一定の性能を満たすのにかかる実時間はバッチサイズによらないことになる。バッチサイズをいじることによる小手先の高速化などは考えない方が良さそうだ。

補足

　3つ目の仮定について補足を加えると、当然ながらバッチサイズが大きい方が1ステップ学習するのにも時間がかかる。実際に今回の実験(RTX 2080ti)では以下のようであった。

バッチサイズ	1ステップあたりにかかる時間(秒)
4096	0.378
512	0.039
64	0.011
32	0.010

　しかしActor-Learner型の強化学習で、Learnerを1ステップごとに数秒スリープさせることで擬似的にActorの数を増やす場合、無視できる差になると思われる。Actorの数が少ない(計算資源が足りない)環境では1ステップあたりにかかる時間の差がそこまで効いてくることはなさそうだ。