Qhapaq Lab

NN キャッシュ

NN キャッシュ（推論キャッシュ）は、方策・価値ネットワークの推論結果を局面キーで保存し、同じ局面を再び評価するときにニューラルネットワークの再実行を避けるためのキャッシュである。

AlphaZero 型のモンテカルロ木探索(MCTS)では、葉ノードを展開するときに policy と value をニューラルネットワークで求める。GPU 推論は強力だが、探索中に何度も呼ぶと重い。そのため、同一局面に対する policy/value を保存しておくと、トランスポジションや再訪問時の無駄な推論を減らせる。

何を保存するか

典型的には、次のような情報を保存する。

  • 局面キー
  • value network の出力
  • 合法手ごとの policy prior
  • 必要に応じて、入力特徴や探索オプションに依存する補助情報

注意点として、NN キャッシュは置換表と似ているが目的が違う。置換表は探索結果、深さ、境界値、最善手などを保存して探索を再利用する表である。一方、NN キャッシュはニューラルネットワークの推論結果を保存する。探索結果そのものではなく、葉評価の入力を節約するための仕組みである。

dlshogiでの例

dlshogi の selfplay/self_play.cpp には、自己対局生成用の明示的な NN キャッシュがある。

struct CachedNNRequest {
  float value_win;
  std::vector<float> nnrate;
};
typedef LruCache<uint64_t, CachedNNRequest> NNCache;
typedef LruCacheLock<uint64_t, CachedNNRequest> NNCacheLock;

nnrate は合法手ごとの policy prior、value_win は value 出力に相当する。既定の nn_cache_size は 8388608 で、コマンドラインオプションから変更できる。

探索中に pos->getKey() が NN キャッシュに存在すれば、子ノードを展開したあと、キャッシュ済みの nnrate と value_win を使う。存在しなければ QueuingNode で局面をバッチ推論キューへ積み、EvalNode でニューラルネットワークを実行したあと、得られた policy/value を CachedNNRequest として nn_cache.Insert(...) する。

selfplay/LruCache.h では LRU 方式のキャッシュと LruCacheLock が用意されている。LruCacheLock は参照中の要素を pin し、並列探索中に参照しているキャッシュ要素が途中で破棄されにくいようにするための仕組みである。

バッチ推論との関係

NN キャッシュはバッチ推論と相性が良い。

  • キャッシュヒットした局面は GPU バッチへ送らずに済む。
  • キャッシュミスした局面だけをまとめて推論できる。
  • 同じ局面を複数スレッドがほぼ同時に要求する場合、キャッシュとキューの設計が探索速度に効く。

ただし、キャッシュヒットが多すぎる設計が常に良いとは限らない。巨大なキャッシュはメモリを消費し、探索スレッド間の同期やキャッシュロックのコストも持つ。実装では、キャッシュサイズ、バッチサイズ、探索スレッド数、GPU の推論速度をまとめて調整する必要がある。

キー設計の注意

NN キャッシュのキーは、ニューラルネットワークの入力が依存する情報を正しく表していなければならない。

将棋では、盤上の駒、持ち駒、手番はもちろん重要である。さらに、実装によっては次のような情報も結果に影響しうる。

  • 千日手や連続王手の千日手の扱い
  • 入玉判定や終局規則
  • 直前手や履歴を入力特徴に含める設計
  • 定跡 policy など、network 出力へ後処理で混ぜる情報
  • モデル、入力特徴形式、softmax 温度などの推論オプション

例えば Lc0 には CacheHistoryLength という設定があり、ニューラルネットワークが履歴を入力に使う場合、キャッシュキーにも必要な履歴を含めないと、同一盤面だが異なる履歴の評価を誤って再利用しうる、という注意点が示されている。将棋でも、局面キーだけで十分かどうかは、そのエンジンの入力特徴とルール処理に依存する。

実装上の注意

NN キャッシュは高速化のための実用的な仕組みだが、探索の意味を壊さない範囲で使う必要がある。

  • ネットワークを差し替えたら古いキャッシュを再利用しない。
  • 入力特徴、定跡 policy、温度などを変えたらキャッシュを分けるかクリアする。
  • 経路依存の詰み探索や千日手判定を value と混同しない。
  • 並列探索では、キャッシュ要素の寿命、pin、ロック粒度を設計する。
  • キャッシュサイズを大きくしすぎると、メモリ帯域や同期で逆に遅くなることがある。

関連項目

参考にしたホームページ