Stockfish NNUE 的网络简介

As of August 6, the efficiently updatable neural network (NNUE) evaluation has landed in the Stockfish repo!

Both the NNUE and the classical evaluations are available, and can be used to assign a value to a position that is later used in alpha-beta (PVS) search to find the best move. The classical evaluation computes this value as a function of various chess concepts, handcrafted by experts, tested and tuned using fishtest. The NNUE evaluation computes this value with a neural network based on basic inputs. The network is optimized and trained on the evaluations of millions of positions at moderate search depth.

The NNUE evaluation was first introduced in shogi, and ported to Stockfish afterward. It can be evaluated efficiently on CPUs, and exploits the fact that only parts of the neural network need to be updated after a typical chess move.

Introducing NNUE Evaluation 写得非常简明, 回答了很多重要问题.


Q: NNUE 是什么? 用在哪里?

A: Efficiently updatable neural network. 作为 alpha-beta search 的估值函数.


图片来自 Stockfish NNUE - Chessprogramming wiki

Q: NNUE 结构是什么?

A: 很简单, 就是 embedding + sum pooling + 几个小的全连接网络.

  • 输入是 0 1 值的特征向量 (one-hot encoding), 输出是标量 (估值). 得到输入的 embedding 后全加起来再过几个全连接网络得到标量.
  • 输入维度是 2 * 41024 (= 2 * 64 * 641). 棋盘用 (己方的王, 其他非王棋子位置) 的组合表示. 其中己方的王有 64 个可能的位置. 641 = 64 * 5 * 2 + 1, 这里的 64 表示格子总数, 5 表示非王棋子种类数, 2 表示己方还是对方; 前面用乘法可以理解为 64 个格子每个格子上有个 10 维的 one-hot encoding 表示这个格子有什么棋; 最后 + 1 是表示将棋打入的遗留物, 国际象棋没用.

KP is combination of own King - piece position [plus] combination of opponent King - piece position. HalfKP is combination of own King - piece position only.

参考原论文 高速に差分計算可能なニューラルネットワーク型将棋評価関数.


Q: 为什么 NNUE 在 CPU 上这么快?

A:

  • 结构上, 取 embedding, 加法, 后面小矩阵乘法, 都可以快速完成.
  • Alpha-beta search 时需要评估很多局面, NNUE 在评估这些局面时, 不需要全部重新计算: 比如不移动王的话, 一步棋对一方来说只改变特征向量中两三个位置的 01 值, NNUE 计算这两个局面时, 只要复用前一步局面的结果, 再取出变动的这几个 embedding, 做后面的计算再加上去就行. 需要算的东西很少.
  • 网络都用整数存储, 并且对这个特定结构的计算方法用 CPU 指令做了优化 (SIMD), 类似向量化计算.

把这个东西直接放 GPU 上反而不划算. GPU 算这种小网络瓶颈是 IO, 非常亏算力. 设计之初原作者目的就是在 CPU 上跑.


Q: 训练数据哪里来?

A: “The network is optimized and trained on the evaluations of millions of positions at moderate search depth.” 将之前人工特征得到的 stockfish 引擎的 moderate search depth (alpha-beta search) 的估值作为真实数据, 训练 NNUE. 让 stockfish 自己对局, 就能生成训练数据.

结语

之前好奇为什么除了棋类引擎外, 其他没见到什么领域谈论 NNUE. 现在知道了. 因为它的结构真没什么特别的 (甚至是早就用烂了的), 但这个结构却正好符合棋类估值的应用 (大量相似局面的评估). 简单结构的精彩应用.