AlphaZeroの探索はすごく速いか?

 あんまり調べず適当なことを言ってしまったのでちゃんと調べた。結論からすると、今のMiacisの実装よりは速そうだがそこまで非現実的な(おかしいと思えるほど)速いわけではなさそうだった。

背景

 AlphaZeroの論文の5ページ目あたりに検証対局時の探索速度の情報が載っている。

AlphaZero searches just 80 thousand positions per second in chess and 40 thousand in shogi, compared to 70 million for Stockfish and 35 million for Elmo.

であり、これは4基のTPUを用いた状況での速度とのことである。おおむねTPUの数とNPSが線形の関係にあるとすると、TPU1基では10K NPS出ると考えられる。これと比較して現在のMiacisの探索速度がどの程度であるか検証を行った。

前提

 現状、Miacisは残差ブロック中におけるチャンネル数が128(AlphaZeroは256)、残差ブロックの数が10個(AlphaZeroは20個)という比較的小さいネットワークを使用している。

 また探索においてはGPUが一度の計算で複数の局面を計算した方が並列計算ができて得であるため、バッチ処理を行っている。基本的にはねね将棋、dlshogiの手法

を真似したものであり、CPUでモンテカルロ木探索の選択ステップを b回行い、GPUで評価したい局面をキューに溜めてから、一括で b個の局面についてGPUで計算を行う。

 このCPUスレッドは複数用意することができる。「CPUでの評価要求を溜める時間 << GPUで計算する時間」である場合、CPUスレッドを1スレッドから2スレッドにすると効率的になるが、上の記事でも言及されているように2スレッドから3スレッドにしても効率的にはならない。

実験

 Miacisを用いて初期局面における探索速度を測定した。以下、使用したマシンはRTX 2080tiが1枚搭載されているものである。

学習済みパラメータ(10ブロック、128ch)

 まずは現在手元にある最も性能の高い学習済みパラメータで探索を行った場合にどの程度の探索速度が出るのかを確認した。

f:id:tokumini:20200219111102p:plain

 バッチサイズを大きくすると10K NPSは超える結果となった。しかしバッチサイズ、スレッド数を増やすことによるNPSの増加が性能向上に繋がっているかどうかは自明ではない。これらを増やしたときには探索が幅広くなっているだけの可能性もあり、PV上の評価要求を投げた局面がなかなか計算されなければむしろ性能低下にも繋がりうる。このあたりは実際に対局させてみないとなんとも言い難いところであり、簡単な検証ではバッチサイズ32、スレッド数2が今のところ一番高い性能となっている。

 注意すべき点としては、バッチサイズ16までは2スレッド→3スレッドでNPSが伸びることはないが、それより大きいバッチサイズにするとスレッド数を増やした方がNPSが増えている。これはつまり「CPUでの評価要求を溜める時間 << GPUで計算する時間」が成り立っていないことを示している(実際にGPU使用率を見ていても40%程度でとどまっている)。どうせGPU律速なのでCPU部分の高速化は気にしなくてもよいだろうと思っていたが、高速化で性能が伸びる余地があるのかもしれない。

ランダムパラメータ(10ブロック、128ch)

 AlphaZeroと同様の20ブロック、256chで検証を行いたいが、そのネットワークでの学習済みパラメータはないのでランダムパラメータでの探索速度を検証することになる。学習済みパラメータとランダムパラメータでは探索速度に差が出そうに思えたため、まず今のネットワークの大きさでランダム初期化したパラメータについても探索速度の検証を行った。

f:id:tokumini:20200219141351p:plain

 直接的な比較ではないためわかりづらいかもしれないが、学習済みパラメータと比べて特にバッチサイズを大きくした際にNPSが2倍程度まで大きくなっている。これはランダムパラメータだと方策、価値がほぼ一様になるため選択ステップ選ばれる局面がバラけやすく、並列化が効率よく行われるためだと思われる。

ランダムパラメータ(20ブロック、256ch)

 AlphaZeroと同じネットワークサイズで測定を行った。

f:id:tokumini:20200219141402p:plain

 バッチサイズを256あたりまで上げてもまだNPSが伸びる傾向があり、8K NPSあたりは目指せそうである。

 ネットワークをこれだけのサイズにすると「CPUでの評価要求を溜める時間 << GPUで計算する時間」が成り立ち、スレッド数を3以上にしても効果がないことがわかった。

全体の結果

 スレッド数2について上の3つの結果を一枚のグラフにまとめると

f:id:tokumini:20200219143932p:plain

 バッチサイズを大きくすると学習済みパラメータであるかランダムパラメータであるかの差が大きくなるが、これはMiacisの実装上の問題かもしれない。

TPUとGPUの性能差も考慮して結論

 対局で使われたTPUが第1世代のものなのか第2世代のものなのかはわからないが、第2世代のものだろうと仮定して話を進める。第2世代TPUは一基で180 TFLOPSとのことである。

 (TPUのチップ一つが45 TFLOPSであり、これを4つまとめたものをTPU一基と見なすようだ。論文での「4TPUs」という言及はこのまとまったTPUを4つ使うのだと解釈した。)

 RTX 2080tiはFP16が108 TFLOPSらしい。

 つまりTPUは2080tiの5/3倍の速度が出ると考えられる。

 先の結果では2080tiでAlphaZeroサイズのネットワークを使うとバッチサイズ256, スレッド数2で7192 NPSであり、学習済みだと半分なる、TPUを使うとして5/3倍になることから換算するとだいたいほぼ6000 NPSとなる。

 AlphaZeroの4TPUで40K NPS(TPU一基あたり10K NPS)というのは、今のMiacisの実装よりも速そうだが、そこまで異様なほど速いというわけでもなさそうだ。