今年のNeurIPS 2025でBest Paperとなった「Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free」を、dlshogiのモデルでも効果があるか試してみた。
Gated Attention
AttentionのSoftmaxの出力直後に、入力に応じたゲートを追加するだけのシンプルな変更で、
- 注目すべき情報がない時にゲートを閉じることで、「とりあえず先頭トークンに注目する」という悪い癖（アテンションシンク）を防ぐ
- 学習中の不安定な挙動（Lossスパイク）を抑え、従来より大きな学習率でのトレーニングが可能
- ゲートが加わることでモデルの表現力（非線形性）が増し、必要な情報だけを通過させる選択能力（スパース性）が強化される
という効果がある。
Attentionは、どのトークン間に注目すべきかを、Softmaxで確率として出力するが、注目すべき情報がない場合でも合計が1になるようにどれかに値を割り当てる必要がある。
これが、悪影響を及ぼし、モデルの性能を抑えている。
dlshogiモデルへの応用
「とりあえず先頭トークンに注目する」というアテンションシンクは、言語モデル固有の課題だが、将棋AIにおいても局面に注目すべき情報がないという状況は発生するはずである。
dlshogiのTransformerモデルに、Gated Attentionを適用することで精度が向上することが期待できる。
モデル構成
論文では複数のゲートの実装パターンが試されているが、Softmax出力の要素ごとに、Attentionの入力に応じて学習可能なゲートを追加する構成が、一番精度が高い結果が示されているので採用することにする。
dlshogiのTransformerモデルは、ResNetの最終層をTransformerに置き換えている(ResNet2ブロックをTransformer1ブロックに置き換える)。
そのTransformerブロックのAttentionに対してGated Attentionを適用する。
入玉特徴量ありのモデルを使用する。
実験条件
ゲートなしの場合と、単純に非線形(ReLU)のみを適用した場合と比較する。
訓練データに、floodgateの2019年以降のR3800以上の棋譜と、dlshogiの自己対局の棋譜、NNUE系ソフトとの対局棋譜、NNUE系ソフト間の対局の棋譜を混ぜた、3.88億局面を使用する。
評価データには、将棋AIモデルの評価用データセットを使用する。
モデルサイズは、20ブロック256フィルタとする。
バッチサイズ4096で、学習率0.04から開始し、1エポックごとに学習率を1/2にし、8エポック学習する。
オプティマイザは、Momentum SGDを使用する(dlshogiのtrain.pyのデフォルト値)。
2回測定し、平均をとる。
実験結果
| 条件 | 方策損失 | 価値損失 | 損失合計 | 方策正解率 | 価値正解率 |
|---|---|---|---|---|---|
| ゲートなし | 1.4263293 | 0.4641902 | 1.8905195 | 0.5264361 | 0.76091445 |
| ReLUのみ | 1.42579545 | 0.4649385 | 1.89073395 | 0.52654415 | 0.7604402 |
| ゲートあり | 1.4225403 | 0.4633177 | 1.885858 | 0.527123 | 0.7614168 |
損失の差(ゲートなし基準)
正解率の差(ゲートなし基準)
ゲートありの場合、方策、価値ともに評価精度が向上している。
ただし、差は、方策正解率で0.068%、価値正解率で0.05%程度である。
なお、学習可能なゲートによりモデルのパラメータ数は増えているため、パラメータ数が増加して精度が向上するのは当然であり、Gated Attentionがどれくらい効果的かはパラメータ数の条件をそろえた上での比較が必要であるが、ゲートなしモデルとパラメータ数をそろえるのはどの箇所でパラメータ数を増やすかにも依存するので条件をそろえるのは難しい。
ReLUのみは、ゲートなしとほとんど変わらなかった。
Transformerを2ブロックにした場合
Transformerを2ブロックに増やして実験した結果は以下の通り。
| 条件 | 方策損失 | 価値損失 | 損失合計 | 方策正解率 | 価値正解率 |
|---|---|---|---|---|---|
| ゲートなし | 1.4267187 | 0.46398875 | 1.89070745 | 0.5264139 | 0.7607165 |
| ReLUのみ | 1.42490745 | 0.4648181 | 1.88972555 | 0.52646415 | 0.76050565 |
| ゲートあり | 1.42423775 | 0.46279415 | 1.8870319 | 0.5274401 | 0.7618291 |
ゲートなしとゲートありの差は、方策正解率で0.1%、価値正解率で0.11%となり、Transformer1ブロックの場合より大きくなっている。損失の方は良くなっていない。
ResNet2ブロック分をTransformer1ブロックに置き換えるだけでは精度はほぼ同じな一方、Gated Attentionは精度が向上する。
ただし、上記でも書いた通り、パラメータ数は増えているため、NPSの低下は避けられず、対局して強さの比較をしないと効果は判断できない。
まとめ
NeurIPS 2025 Best PaperのGated AttentionをdlshogiのTransformerに適用し、将棋AIでも有効かを検証した。
実験の結果、方策・価値ともに精度が向上した。
パラメータ増加の影響を考慮した強さ評価は今後の課題とする。
今回、学習可能なゲートの線形層を独立した層で実装したが、QKVの線形層とまとめて計算が可能なため、この最適化を行った後に、強さの比較を行う予定である。
