Arena評価のPUCT探索を早期終了させ、1世代の所要時間を23%短縮した

モデルの更新採否を判定するarena評価フェーズは、当時のCPU学習ループにおいて全体の実行時間の約77%を占めていた。arenaで必要なのは最終的に選択される着手(action)のみである。そこで、残りのシミュレーション回数をすべて他の候補手に割り振っても1位の手が覆らないと確定した時点で、PUCT探索を早期終了(early stop)できるようにした。

なお、自己対局(self-play)では訪問回数分布をpolicyの学習ターゲットとして用いるため、この早期終了は適用せず最後まで探索を行う。

selected actionが確定したら止める

PUCT探索において、残りシミュレーションをすべて他の候補手へ与えても現在の最多訪問手を追い越せない状態になれば、最終的に選ばれる手は変わらない。そこで実装上はarena評価時のみを「選択手のみを求める探索モード」として扱い、自己対局やベンチマーク局面の評価では従来通り完全な訪問回数分布を要求する仕様とした。

学習ターゲットや入力特徴量、ONNXの推論インターフェース、乱数シード、生成される自己対局データなどの設定は一切変更していない。

1 generationで23.3%短縮

Apple M2、8 workersの環境で、PUCT128自己対局256局、512 updates、clean PUCT128 arena 800 pairsによる1世代の所要時間を比較した。

stagebaselineearly stopchange
self-play236.5 s234.6 s-0.8%
fit207.2 s205.3 s-0.9%
arena1,505.8 s1,056.0 s-29.9%
generation total1,955.6 s1,500.5 s-23.3%

arenaにおけるevaluator行数は11,189,248行から約7,748,000行へと30.8%減少し、プロセス全体のCPU時間も25.4%削減された。4回の試行におけるarenaの勝敗集計はすべて850勝10分740敗で一致し、モデル採否の判定結果も同一だった。

早期終了した探索から完全な訪問回数分布を取得しようとした場合はエラーを発生させ、中途半端な探索分布が誤って学習データに組み込まれるのを防いでいる。

他の高速化は採用しなかった

試したもの結果判断
ORT memory arena無効化130.3→127.7 s、RSS +38%不採用
workerごとにORT session260.6→262.7 s、RSS +43%不採用
sole-candidate shortcut89,378 decisions中3回効果が小さい
partition負荷分散254.5〜261.8 s長いtailなし

プロファイリングの結果、ワーカー処理時間の96%をONNX推論が占めていた。そのため、探索アルゴリズムの内部処理を微小に最適化するよりも、結果に影響しない末端(leaf)の推論呼び出しを間引くアプローチの方が大幅な高速化に寄与した。


この記事は、実装・実験記録をもとに、本文の大部分をLLMが執筆し、筆者が内容を確認・編集しています。