優先度付き経験再生の実装・実験

要約

　優先度付き経験再生はAlphaZero方式の学習でも効果がありそう。

背景

　以前の記事でも軽く触れたが、優先度付き経験再生という手法がある。

　大雑把に言うとリプレイバッファからのサンプリング確率を一様ランダムではなく優先度で重み付けするものである。 $i$ 番目のサンプルの優先度を $p_i$ としたとき、このサンプルを確率

$$ P(i) = \frac{p_i^\alpha}{\sum_j p_j^\alpha} $$

でサンプリングする。優先度をTD誤差に応じて適切に設定すればより重要なデータを多く学習することができ、学習速度や性能が向上するとのことである。Actor-Learnerを並列で動かす場合にはかなり性能が良いという報告もある。

　つまりAlphaZero方式の学習でも効果がありそうだと感じたので実装してみた。

　元論文では優先度の付け方として

TD誤差の絶対値の大きさによって順位付けし順位の逆数を優先度とする
TD誤差の絶対値をそのまま優先度とする

という2つの手法が提案されていたが、ここでは(2)を採用することとする。若干の変更点としてTD誤差の絶対値ではなく損失関数の値を利用する。

実装

　現状はMiacisのperブランチで開発を行っている。十分な安定性が確認できたらいずれmasterブランチにマージする。replay_buffer.cppやsegment_tree.cppが重要な部分となる。

更新・サンプリングの実装

　(2)の手法はSum-Tree(セグメント木)を用いて実装することで、優先度の更新、サンプリングが $O(\log{N})$ でできる。学習データを保持する配列に加えて、各データの優先度を保持するセグメント木を別に持つことになる。

　優先度の更新は通常のセグメント木の更新そのものとなる。

void SegmentTree::update(uint64_t x, float v) {
    sum_[x + n_ - 1] = v;
    for (uint64_t i = (x + n_ - 2) / 2; ; i = (i - 1) / 2) {
        sum_[i] = sum_[2 * i + 1] + sum_[2 * i + 2];
        if (i == 0) {
            break;
        }
    }
}

　サンプリングは0から損失の総和までの一様乱数をサンプリングし、値に応じて最下段まで木を降りていくことで実現する。

uint64_t SegmentTree::getIndex(float value, uint64_t k) {
    if (k >= n_ - 1) {
        //最下段まで来ていたらindexを返す
        return k - (n_ - 1);
    }
    return (value <= sum_[2 * k + 1] ? getIndex(value, 2 * k + 1) : getIndex(value - sum_[2 * k + 1], 2 * k + 2));
}

　データは通常の経験再生と同じように古い順に削除する。前回データを入れた位置を保持しておき、バッファをリング状に巡回する。

uint64_t SegmentTree::getIndexToPush(uint64_t k) {
    return (++next_push_index_) %= n_;
}

新データの挿入

　元論文では新しく生成されたデータは格納されている中で最大の優先度と同じ優先度を与えるとしているが、今回は置換表に保存していた情報から仮の損失値を計算し、それをもとにリプレイバッファへ格納する。

//このデータを入れる位置を取得
int64_t change_index = segment_tree_.getIndexToPush();

//そこのデータを入れ替える
data_[change_index] = std::make_tuple(pos.toSFEN(), teacher);

//priorityを計算
float priority = -2 * std::log(nn_output_policy[move.toLabel()] + 1e-10f) 
+ std::pow(nn_output_value - teacher.value, 2.0f);

//segment_treeのpriorityを更新
segment_tree_.update(change_index, std::pow(priority, alpha_));

　仮の損失値を計算する上での注意点として、置換表に保存されているPolicyやValueは局面に対するニューラルネットワークの生の出力とは異なる。たとえばPolicyでは置換表には合法手でマスクしたのちにSoftmax関数をかけたものを保存しているため、全指し手についてSoftmax関数をかける通常の出力より精度が良くなる。よってこの仮損失を2倍した値を利用している。この2倍という値に根拠はなく、ただの勘により定めた。

　また、今実装を見直していて気づいたのだが、Policyの仮損失の計算が、実際に指された指し手のみをOnehotベクトルとして計算するものになっている。教師を分布化する前に実装したものをそのままにしていたのだった。どうせ仮の値なのでそこまで影響を与えないとは思うが、いずれ修正する。

　これもまた実装を見直していて気づいたのだが、Valueを漸進的に更新するように実装を変更したのでValue側は探索を終えて精度がかなり高まったものを利用して損失を計算することになっている。上の件とも含めて、仮損失が良い推定値にはなっていないように思える。

　以下の実験はこの実装で計測してしまったのだが、新しいデータのサンプリング確率が不当に低くなっている可能性がある。

重点サンプリングによる補正

　論文で言及されている通り、一様分布ではない分布からサンプリングをするとバイアスが生じてしまうため、本来は重点サンプリングの理論から更新時に補正を行う。しかし今回はそれを入れなかった。なにかの論文でこの補正は行わなくても精度が落ちなかったというような記述を目にした覚えがあるが、どの論文だったか忘れてしまった。

実験

実験設定

　GTX 1080を2枚搭載したマシンでAlphaZero方式の学習を行った。2枚とも自己対局に使用している場合にデータ生成速度は33 pos / secほどとなる。学習1ステップごとに1秒のスリープを行ったため、ステップあたりのデータ生成速度もほぼ33 pos / stepと考えて良い。バッチサイズは128で学習をしたため、同じデータが何度か学習される設定となっている。そのような設定でこそ優先度に応じたサンプリングをする意味があるのではないかと思われる。

　 $\alpha = 0,1,2,3$ の4通りについて実験を行った。元論文では $\alpha = 0.7$ などとしているが、扱う値がTD誤差と損失値と違うのでとりあえず大きい値で実験を行ってみることとする。 $\alpha = 0$ のときは結局優先度が全て1で等しくなるので通常の一様ランダムでサンプリングを行う経験再生に一致する。

実験結果

棋譜から計算する検証損失

　floodgate2016年の棋譜から計算した検証損失は次のようになった。

f:id:tokumini:20190613092231p:plain — 左:Policy損失　右:Value損失

f:id:tokumini:20190613092232p:plain — 左:Policy損失　右:Value損失

　Value損失ではほとんど差がなかったが、Policy損失では $\alpha$ を大きくした方が良い結果となった。優先度とする損失にはPolicy損失とValue損失を1:1で和をとったものを利用しているので、異なるデータ間で値の大小を比較する上ではPolicy損失の影響が大きいのではないかと思われる。Policy損失とValue損失に係数をかけて比を変えられるようにはしているが、このパラメータもチューニングしなければいけないとなると大変である。

技巧2との検証対局

　各 $\alpha$ で学習したパラメータを用いて技巧2(深さ制限4)と1手0.25秒でそれぞれ500局の対局を行った。対局時にはRTX 2080ti1枚を搭載したマシンを利用した。

$\alpha$	勝率	レート差
0	34.6%	-100.6
1	32.6%	-126.2
2	57.5%	+52.5
3	54.5%	+31.4

　今回試した中では $\alpha = 2$ としたときが最も良い結果となった。元論文の結果からするとかなり大きい値に思われるが、TD誤差から優先度をつけている元論文とは異なり、損失で優先度をつけていることが影響しているのかもしれない。たとえばPolicy損失では教師信号が探索回数を正規化した分布になるので、交差エントロピーには定数値として教師信号の自己エントロピーが足されている。Value損失に関しても、報酬の範囲を[0,1]として最終層がsigmoid関数で交差エントロピーを損失とする場合には、教師信号が0.5(引き分け)であれば損失の最小値が0ではない。微分を計算する際には定数値のシフトによる影響はないのだが、今回の手法では問題になり得るのかもしれない。