大規模言語モデルで将棋AIを作る その６（相対位置エンコーダ）

前回までは、位置エンコーダに学習可能な絶対位置エンコーダを使用していた。
今回は、相対位置エンコーダを試す。

位置エンコーダ

Transformerは、トークンが入力の何番目にあるかによらず等価に扱う。
そのため、入力の位置が意味を持つ場合は、位置を何らかの形でエンコードする必要がある。
自然言語処理では、文章をトークン列にした場合、何番目のトークンかをSinとCosでエンコードすることが行われる。

Llama2では、rotary positional embeddings(RoPE)という、トークンの相対位置をエンコードする手法が使われている。

将棋の位置エンコーダ

将棋は、盤上の絶対的な位置だけでなく、相対的な位置関係が意味を持っている。
また、位置は2次元の座標のため、単にトークン間の距離だけではなく、2次元上の位置関係を扱う必要がある。
さらに、駒によって相対的な位置の意味が変わる場合や、飛車や角のように効きが長い駒の場合、間にある駒によっても位置の意味が変わってくる。

そのため、自然言語処理で使われている手法をそのまま適用することはできない。

Leela Chess Zeroでは、Smolgenという、チェスの駒の相対位置関係をとらえるための独自の位置エンコーダを考案している。
トークン列の埋め込みを32次元のベクトルに圧縮して、それをトークンの相対位置に全結合層で変換するという手法である。

この考えを参考に相対位置エンコーダを実装する。

相対位置エンコーダ

相対位置エンコーダの論文で示されている式のは、トークン間の距離に応じた学習可能パラメータになっているが、この部分をLeela Chess Zeroを参考に入力のトークン列を32次元のベクトルに圧縮して全結合層で変換したものに置き換える。
また、1つ目の式の右辺第2項のを、相対位置用の学習可能なパラメータとする。

実装

class TransformerEncoderLayer(nn.Module):
    def __init__(self, channels, d_model, nhead, dim_feedforward=256, dropout=0.1, activation=nn.GELU()):
        ...
        self.relative_linear1 = nn.Linear(channels * 81, 32, bias=False)
        self.relative_linear2 = nn.Linear(32, self.d_model * 81, bias=False)
        self.relative = nn.Parameter(torch.randn(self.nhead, 81, self.depth))

    def forward(self, x):
        ...
        r = self.relative_linear1(x.flatten(1))
        r = self.relative_linear2(r)
        r = r.view(-1, self.nhead, self.depth, 81)
        r = torch.matmul(self.relative, r)

        scores = (torch.matmul(q, k) + r) / math.sqrt(self.depth)
        attention_weights = F.softmax(scores, dim=-1)
        ...

学習結果

相対位置エンコーダを使わない場合と使う場合で精度を比較した。
また参考としてResNet20ブロック256フィルタのモデルとも比較した。

前回まで、学習データにfloodgateの2018年から2023年のデータを同一局面を平均化しないで学習していたが、序盤に偏ったデータになるため、今回は「--use_average」オプションで同一局面を平均化した。

	nhead	feed forward	活性化関数	val loss	policy acc.	value acc.
ResNet				2.5771	0.4147	0.6685
相対位置エンコーダなし	8	256	gelu	2.5840	0.4163	0.6657
相対位置エンコーダあり	8	256	gelu	2.5660	0.4120	0.6774

相対位置エンコーダがある場合が評価損失が若干良い結果になっているが、ランダムの影響の範囲かもしれない。
データ量を増やして学習してみないとはっきりした違いはわからない。

なお、32次元から64次元に増やしたり、ヘッドでパラメータを共有したり、LayerNormを入れたりも試したが、若干悪くなるくらいで大きな差はなかった。

まとめ

将棋の盤上の相対位置関係をとらえる相対位置エンコーダを実装した。
学習の結果、相対位置エンコーダを使用しない場合と大きな違いはなかった。
訓練データが少ないため、もっと大きなデータで差がでるか別途検証したい。

次は、活性化関数にSwiGLUを使った場合や、絶対位置エンコーダの場所を変えたりネットワーク構成を変えて比較してみたい。
また、Leela Chess Zeroで試されているDeepNet initializationについても試したい。