【arXiv 2020】KataGo:把自对弈学习加速 50 倍|从围棋 AI 训练效率视角

摘要

本文解读 arXiv 2020 论文《Accelerating Self-Play Learning in Go》。该论文提出KataGo 自对弈训练流程 ,通过融合playout 上限随机化策略目标剪枝全局池化辅助预测目标 ,只用 19 天、约 1.4 GPU-年 就从随机棋开始训练出超越 ELF OpenGo 最终模型的围棋 AI,其特别之处在于约 50 倍的计算量缩减大部分来自与围棋无关的通用改进 。实验表明各改进组件独立消融的加速乘积约为 9.1 倍(且被低估),为自对弈强化学习的数据效率工程提供了重要借鉴。

视频讲解点击观看 B 站视频

论文基本信息

项目 内容
标题(英文) Accelerating Self-Play Learning in Go
标题(中文) KataGo:把自对弈学习加速 50 倍
作者 David J. Wu
机构 Jane Street Group
会议 arXiv 2020
arXiv https://arxiv.org/abs/1902.10565
项目网站 https://github.com/lightvector/KataGo

背景与动机

自对弈强化学习在围棋上取得了划时代的成功,但代价极其高昂:AlphaGo Zero (Nature 2017)与 AlphaZero (Science 2018)的主运行消耗约 41 TPU-年ELF OpenGo (ICML 2019)复现消耗约 74 GPU-年 ;分布式社区项目 Leela Zero 则跨越了多年训练。论文指出,这种巨大开销主要来自自对弈流程中一个被忽视的数据-目标张力:价值目标只需要少量 playout 生成更多对局,而策略目标需要大量 playout 才能让搜索偏离策略先验、推动策略改进------固定 playout 数无法同时满足两者。

KataGo 由此提出两方面的改进:四项通用改进 (playout 上限随机化、强制 playout 与策略目标剪枝、全局池化、辅助对手策略目标)可迁移到其他 AlphaZero 类学习任务;两项围棋特化改进 (归属/分数辅助目标、棋感输入特征)则揭示通用流程与最优实践之间仍存在真实效率差距。这些技术并非全部原创:全局上下文思想呼应 Squeeze-and-Excitation (CVPR 2018),子事件预测呼应 Darkforest (ICLR 2016)与 MLVGo 多标签价值网络(IEEE ToG 2018),但本文是首次将其系统化引入自对弈强化学习流程。

图 1:1600 访问 Elo 随自对弈计算量(等效 20b×256c 查询,对数坐标)的变化------KataGo(蓝)以约 1/50 于 ELF(绿)的计算量达到同等强度。

研究主线:从问题到结论

图 9:KataGo 研究主线------从算力昂贵的问题出发,经目标解耦设计、多输出头方法、六臂消融实验到效率可工程化的结论。

基准/方法设计

KataGo 整体架构沿用 AlphaGo Zero / AlphaZero:神经网引导的蒙特卡洛树搜索(MCTS)自对弈生成训练数据,节点选择采用 PUCT 公式 V(c) + c_{\\text{PUCT}} P(c) \\frac{\\sqrt{\\sum_{c'} N(c')}}{1 + N(c)}c_{\\text{PUCT}}=1.1),根节点施加 Dirichlet 噪声(\\alpha = 0.03 \\times 19\^2 / N_{\\text{moves}})与根温 1.03。

核心设计是四项通用改进

  1. Playout 上限随机化:以概率 p=0.25 的回合执行 600--1000 节点的全搜索并记录训练,其余回合用 100--200 节点的快搜索。全搜索回合提供高质量策略样本,快搜索回合让价值训练获得更多对局------直接缓解策略/价值目标的数据张力。
  2. 强制 playout 与策略目标剪枝 :根节点每个收到过 playout 的子节点强制至少 n_{\\text{forced}}(c) = (k P(c) \\sum_{c'} N(c'))\^{1/2}k=2)次探索;随后在策略目标中减去 这些多余 playout(以不超过最优子节点的 PUCT 为限)。这实现了策略目标与 MCTS 探索动力学及噪声的解耦------目标分布不再被探索噪声污染。
  3. 全局池化:对 c 个通道计算均值、随棋盘尺寸线性缩放的均值与最大值(共 3c 个值),经全连接后以逐通道偏置形式调制另一组通道,让卷积层条件化于全局上下文。
  4. 辅助对手策略目标 :策略头额外输出对手下一手预测 \\hat{\\pi}*{\\text{opp}},损失权重 w*{\\text{opp}}=0.15,作为近零成本的训练正则化。

分类全景

图 10:KataGo 改进体系------四项通用改进(playout 上限随机化、强制 playout 与目标剪枝、全局池化、对手策略目标)与两项围棋特化改进(归属/分数目标、棋感特征)。

方法细节

神经网络的输入由两个张量构成:空间特征 b \\times b \\times 18(棋盘宽度 b \\in \[9,19\])与 10 维全局特征:

特征类别 通道数 内容
盘面基础 1+2+3 是否在盘内、己方/对方棋子、1/2/3 口气
规则相关 1+2 劫/超劫禁手、pass-alive 区域(己方/对方)
历史与战术 5+3+1 最近 5 手 one-hot、0/1/2 手前可征子、落子可征吃
全局 10 前 5 手是否 pass、贴目/15、劫规则、是否允许自杀、贴目与棋盘奇偶

网络为预激活残差网,trunk 从 b=6, c=96 渐进增大至 b=20, c=256(与 AlphaZero/ELF 相同规模);策略头双通道输出 \\hat{\\pi}\\hat{\\pi}_{\\text{opp}};价值头包含胜负、逐点归属、最终分数分布三个子头。

图 2:全局池化偏置结构------聚合均值/尺寸缩放均值/最大值共 3c 个值,经全连接后作为逐通道偏置调制另一组通道。

图 3:10 块网络策略对数图(带强制 playout 与策略目标剪枝)。

图 4:对照版本------策略质量明显扩散到棋盘边缘的大量坏棋上,剪枝显著改善了策略聚焦。

图 5:归属预测可视化------每个交叉点颜色表示最终归属,比二值胜负细粒度得多,为价值学习提供直接梯度反馈。

训练细节(附录 D):主运行 19 天生成 420 万局 / 2.41 亿样本 ,批量 256、逐样本学习率 6\\times10\^{-5}(前 5M 样本降 3 倍、末段降 10 倍),均匀采样自 25 万增长到 2200 万的滑动窗口;采用随机权重平均 + 快照指数滑动平均(decay 0.75),候选网须在门控测试(200 局胜 100 局)中击败当前自对弈网络才可接替。游戏随机化(附录 E)让单模型泛化全部变体:棋盘 9--19 路三角分布、超劫/自杀规则随机、贴目 N(7,1) 采样、5% 让子局、2.5% 分支探索局。

损失函数(附录 B)以胜负价值(c_{\\text{g}}=1.5)与策略交叉熵为主,辅助项包括归属损失(权重 1.5/b\^2)、分数分布 pdf/cdf 损失(各 0.02)、分数均值/方差自预测 Huber 损失(\\delta=10,权重 0.004)与 L2 惩罚(3\\times10\^{-5});作者坦承这些辅助系数"mostly guesses"、未精细调优,却带来显著且一致的提升。

实验设计与结果

评测协议:19×19 棋盘、Tromp-Taylor 规则、固定 7.5 贴目、1600 访问、约 21,000 局确定 Elo;对手为 ELF OpenGo 原生引擎Leela Zero 0.17 网络谱系。消融协议:6 条约 2 天的短训练臂(FixedN / NoForcedTP / NoGPool / NoPAux / NoVAux / NoGoFeat),约 147,000 局定 Elo。

与 ELF 的 400 局对决(三种设置全胜):

对局设置 胜场 / 400 Elo 差
1600 playouts/mv 无 batching 239 69 ± 36
9.0 s/mv(ELF batch 16) 246 81 ± 36
7.5 s/mv(ELF batch 32) 254 96 ± 37

图 6:主运行 vs 固定上限(N=100~600)------playout 上限随机化一致占优,印证策略/价值张力假设。

图 7:NoGPool / NoForcedTP / NoPAux 三条曲线在训练后期与主运行差距持续扩大。

图 8:移除归属/分数目标(NoVAux)与棋感特征(NoGoFeat)均明显拖慢学习。

六臂消融总表(附录 G,2.5G 等效查询 ≈ 2 天处):

移除组件 Elo 加速因子
(主运行基准) 1329 1.00×
Playout 上限随机化 1242 1.37×
强制 playout + 目标剪枝 1276 1.25×
全局池化 1153 1.60×
辅助策略目标 1255 1.30×
归属 / 分数目标 1139 1.65×
棋感特征与优化 1168 1.55×

结果对比总结

图 11:结果对比------ELF 74 GPU-年与 Leela Zero 多年分布式训练 vs KataGo 1.4 GPU-年,效率分别提升约 50 倍与 10 倍。

关键发现

  • 总效率 :KataGo 以约 1.4 GPU-年 (19 天 × 27 张 V100)超越 ELF 的 74 GPU-年 ,实现约 50 倍 计算量缩减;同网络规模下较 Leela Zero 快约 10 倍
  • 对战全胜 :三场 400 局对决分别以 239、246、254 胜战胜 ELF,Elo 差 +69 / +81 / +96,低搜索与高搜索、固定搜索与固定墙钟设定下均成立。
  • 最强单一因素 :归属/分数辅助目标,移除后 Elo 从 1329 降至 11391.65 倍时间损失)------把二值胜负分解为细粒度子事件监督是最高效的改进。
  • 通用性验证 :playout 上限随机化(1.37×)、全局池化(1.60×)、强制 playout 与目标剪枝(1.25×)、辅助策略目标(1.30×)四项与领域无关的改进合计贡献了绝大部分加速。
  • 领域先验仍有价值 :棋感特征与行棋优化贡献 1.55 倍,说明 AlphaZero 式通用流程远未穷尽领域知识红利。
  • 消融低估 :各因素加速乘积约 9.1 倍,且多数技术在 2 天消融窗口内仍在持续放大收益,真实总加速应更高。

局限性

  • 消融窗口短:约 2 天的消融运行使加速因子偏向近似与低估,且无法复现 ELF/AlphaZero 的千卡级基础设施做精确对照。
  • 部分结论依赖非正式测试:论文多处基于"informal tests"(如 playout 数偏好、剪枝初始观察);辅助损失系数多为经验猜测、未精细调优。
  • 计算量指标是粗略近似bc\^2 成本模型、转置缓存折算与 Leela Zero 缺失数据的插值都会引入误差。
  • 通用方法仍有差距:围棋特化特征贡献 1.55 倍的事实表明,纯通用自对弈流程与领域最优实践之间仍存在数量级层面的改进空间。

常见问题(FAQ)

KataGo 为什么能比 AlphaZero 快 50 倍?

核心在于解耦与增密:playout 上限随机化让价值与策略目标各取所需的数据量,策略目标剪枝把训练目标从 MCTS 探索噪声中解放出来,归属/分数等辅助目标则为有限的胜负信号补充了细粒度监督,四类改进的消融加速乘积约 9.1 倍。

playout 上限随机化解决了什么问题?

价值目标希望每局便宜(少 playout、多对局),策略目标希望搜索充分偏离先验(多 playout)。固定上限无法两全,随机化让 25% 回合执行全搜索供策略训练、其余回合快搜供价值训练,两者都获得足够数据。

策略目标剪枝具体剪掉了什么?

它把根节点因强制探索或噪声而多投入的 playout 从训练目标中减去,只保留 PUCT 自然会选择的比例------除非某个被强制探索的招法被证明是好的。这样策略网络不会被探索噪声带偏,又能受益于强制探索发现的好招。

辅助目标为什么有用?

论文给出一个启发式:当期望目标可以表达为子事件之和、合取或析取时,预测子事件大概率有帮助------因为子事件错误带来的梯度会精确局部化到误判区域,加速信用分配;NoVAux 消融 1.65 倍的损失是最强证据。

这些改进能用到围棋之外吗?

四项通用改进均与游戏规则无关:playout 上限随机化适用于任何 AlphaZero 类流程;全局池化可迁移到其他棋盘/图类任务;辅助对手策略与子事件预测可推广到一般强化学习,甚至单智能体环境预测。

KataGo 现在还在发展吗?

是的。KataGo 开源后演化为社区驱动的分布式自对弈项目,至今仍是顶级开源围棋 AI,被韩国国家队与主流在线围棋平台用于分析;其单卡数日训练到业余强手的能力也让个人研究者可以低成本复现自对弈强化学习。

参考链接


给大家推荐一款自用写文献综述、无虚构文献的 AI:

🌟 复旦大学 FudanNLP 团队自研 切问学术

官网:qiewenpaper.com

覆盖 3.6 亿篇可溯源真实中英文文献,能自动整合文献观点生成规范综述

还能挖掘研究创新点、复现实验,配合视频教学,新手快速上手文献综述写作


🍀后记🍀

博客的关键词集中在编程、算法、机器人、人工智能、数学等等,持续高质量输出中。

🌸讨论QQ群白拾的小屋 (750365700)

⭐B站账号白拾的物理AI组会(活跃于知识区和动画区)

✨GitHub主页YhbCode000(工程文件)

相关推荐
nju_spy1 年前
王树森深度强化学习DRL(三)围棋AlphaGo+蒙特卡洛
强化学习·南京大学·alphago·蒙特卡洛树搜索·策略网络·价值网络·随机梯度算法