候補手をpacked表現に刷新し、EATの学習速度を3.8倍に高速化した
Entity-Action Transformer（EAT）の教師あり学習では、盤面の状況によって選択可能な合法手（候補手）の数が大きく変動する。平均すると約26.8手であるのに対し、最大局面では195手に達するため、従来の固定長パディングによるバッチ処理では、候補手の少ない局面に対しても最大幅分の冗長な計算リソースを浪費していた。
そこで、候補手データを実在する要素のみで密に詰めた packed 表現へと刷新し、512サンプルからなるバッチをマイクロバッチに細分化することなく一度にGPUへ投入できるパイプラインを構築した。
学習スループットは3.82倍へ向上
| 条件 | スループット |
|---|---|
| 改善前（before） | 4,053 rows/s |
| 改善後（after） | 15,501 rows/s |
| 高速化倍率 | 3.82倍 |
単に packed 表現を導入しただけでは、従来の小さなマイクロバッチ設定のままでは劇的な速度向上は得られなかった。無駄なパディングを排したことでGPUメモリ消費に大きな余白が生まれ、512行を一括処理できるようになったことが飛躍的なスループット向上の主因である。
| ボトルネック比較 | 改善前 | 改善後 |
|---|---|---|
| 候補手のテンソル配置 | バッチ内の最大候補数まで一律パディング | 実際に存在する候補手のみを詰めた packed 表現 |
| 512行の更新処理 | 128行 × 4マイクロバッチに分割 | 512行を一括処理（1マイクロバッチ） |
| モデル構造・損失関数・最適化 | 変更なし | 変更なし |
方策スコアを計算した直後のみ、候補手インデックスを参照して元の [batch, width] テンソルへと整列し直す。学習に使用するデータ行、損失関数の定義、オプティマイザ、ネットワーク構造には一切変更を加えていない。
今回の3.82倍という高速化は、パディングの除去そのものだけでなく、メモリ効率の改善によって大規模なバッチ処理が可能となった複合的な成果である。教師あり学習の実験サイクルを高速に回すための重要な基盤改善となった。
この記事は、実装・実験記録をもとに、本文の大部分をLLMが執筆し、筆者が内容を確認・編集しています。