Transformerで将棋の学習

　最近Transformer(ViT)モデルで教師あり学習をやっていたところ、損失値についてはCNNと近い程度の低さが出せるようになってきたので記事としてまとめておく。

　(※ 教師あり学習での損失値を比較しているだけなので、CNNより強いという主張をしているわけではない。特に推論時の速度は懸念材料であり、TRTorchではTransformerモデルのTensorRT変換が上手くいっていないため比較できていない。onnx経由で変換するにしてもCNNと同程度の速度が出るのかな？とちょっと怪しんでいるところはある)

前提知識

　Transformerとは？ → 様々な人が解説記事を書いているのでそれらを参照のこと。たとえばこれ

　今回使うのは、Transformerを使う中でもVision Transformerに近いやりかたを採用(谷合さんが以前にやっていたのとは異なる)。

　Vision Transformerについても様々な人が解説記事を書いているのでそういうものを検索すれば良いと思う。たとえばこれ

使用モデル

　Vision Transformerに似たモデルを用いる。

　画像であれば入力サイズは3×H×Wであり、HやWが大きいことから適当サイズでのパッチ化が必要となるが、将棋の盤面は9×9であり、これが十分小さいためそのままの解像度で処理を行う。

　Miacisでは各マスの情報を42チャンネルで表現しているので、42×9×9が入力となる。まずこれを1層MLPの線形変換により各マス独立でチャンネル数を $C$ チャンネルへ変換する。 $C$ はハイパーパラメータである(256とか512とか)。 $C \times 9 \times 9$ を、 $C$ チャンネル81系列の $C \times 81$ とみなし、 $L$ 層のTransformerEncoderに入力する。Transformerはシェイプ不変の変換であるため、最終的に $C \times 81$ の表現が得られる。これを $C \times 9 \times 9$ と解釈し直し、CNNの場合と同じ構造のPolicyヘッド、Valueヘッドに入力する。

　Vision Transformerでは $N \times N$ のパッチ化の後に、クラストークンを追加し $N \times N + 1$ 系列として処理していくが、今回は比較のためにCNNとシェイプを統一したかったため、そのような処理は行わなかった。

　位置エンコーディングは81マスそれぞれについてランダム初期化した学習可能な $C$ チャンネルのパラメータを用意し、最初の1層MLPによる線形変換直後に、足し算によって与える。

ソースコード

学習の工夫

　3ヶ月ほど前にTransformerを試した際は、あまり上手くいかなかった。

ほぼ同じ条件で盤面を81個の系列とみなしたTransformerを学習させたけど、CNNに比べてはっきり悪い。12万ステップ程度では、ということ？ホントかなぁ pic.twitter.com/3eUcwAGoOy
— t (@tokumini_ss) 2021年5月18日

　このときから今回までの変更点として

学習ステップ数の増加(12万→60万)
ベース学習率の低下(0.025→0.01)
学習率のWarm-up(最初の2万ステップで0から0.01まで線形に増加させる)
学習率の線形減衰(最後の60万ステップで0になるように0.01から線形に減衰させる)

などを入れた。それぞれ明確な切除実験を行ったわけではないためどれが重要だったかは不明だが、ベース学習率の低下と、学習率のWarm-upは大きく寄与しているのではないかと思われる。

　学習率の制御などについてはScaling Vision TransformersというVision Transformerについていろいろ調べた論文や、以下のツイート

私のViT体験談。optimizerの設定がシビアですが、きちんと調整すれば、CNNを上回ることがしばしばあります。

CPMPの神pdfが参考になります。（4.4 training transformers参照、lrを小さくするのがコツ）

データ数が50/classでも精度出ます（タスクによりけり）。https://t.co/yugLSPRny7
— shinmura0 (@shinmura0) 2021年8月20日

およびここで紹介されているPDFなどを参考にした。

実験

使用データ

　学習にはAobaZeroの棋譜000011500000〜000014300001を用いた。(データ取得用スクリプト)

　検証にはfloodgate2015年の棋譜から、どちらかの対局者のレートが3200以上となっている棋譜のみについて検証損失を計測した。

　学習データ、検証データ両方についてそれぞれ重複局面を削除した。削除の際には指し手や勝率などは同じ局面について統計を集め、平均化したものを教師情報とした。

CNNとの差異

　CNNは34万ステップ、Transformerは60万ステップ。CNNはGoogle Colabで回したものなので、24時間制限で60万ステップまで回りきらなかった。平等な条件でやり直したいところではあるが、計算資源が足りないので……。

　モデルサイズは学習時間が同じくらいになるよう調整して、TransformerではCNNの半分のブロック数にした。

　学習条件が異なる部分は主に以下の表でまとめられる。

項目	CNN	Transformer
学習ステップ数	34万	60万
学習率減衰	15万ステップごとに1/10	最終ステップで0になるように線形減衰
ブロック数	20	10
チャンネル数	256	256

　Valueは両方とも51分割したカテゴリカル分布でValueを表現するモデルである。

検証損失

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

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

Policyの低下はCNNの方が速い
Valueの安定感はTransformerの方が良い

最終ステップでの損失値

モデル	Policy損失	Value損失
CNN	1.940291	0.516135
Transformer	1.915562	0.516003

　学習量約2倍というハンデはあるが、一応両方の値でCNNを下回っている。

34万ステップ時点での学習時間

モデル	経過時間	GPU
CNN	20:42:20	V100(Google Colab)
Transformer	22:03:56	2080ti(自宅のPC)

　2080tiよりV100の方がわずかに学習が速いので、だいたい同じ程度だと思われる。

学習損失

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

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

　CNNのValue損失がギザつくのは以前からよく見られていた現象ではあった。これは、学習で使用しているAobaZeroの棋譜が1万棋譜ごとにまとめられていて、メモリの関係上、その単位で読み出して使用しているので、その間はほぼ同じ時期に生成されたデータを使うことになっているためだと思われる。

　TransformerだとValue損失でのギザつきが発生しておらず、またPolicy損失についても「学習損失と検証損失のギャップ」が小さい。単純に解釈すると、汎化しやすいような学習ができているのではないかと思われる。

　CNNは1万棋譜ごとに過学習しては別に塊に移って反省するということを繰り返しているように見える。CNNの改善を目指すならば(特にValueの)正則化を強めるのは選択肢として挙がりそう。 Scaling Vision Transformersの論文には、ネットワークのボディ部分とヘッド部分で重み減衰の係数を変えた方が良いという話があり、それはCNNでも有効かもしれない。