# AlphaZero 之加速演算法實作

###### tags: `圍棋` `論文` `電腦對局`

~~* 註: 本論文報告對應版本 v0.4 ~ v0.5，其中描述的一些方法在之後的版本有被修正~~

* 註：修正最新的版本，編輯中...

## 簡介

DeepMind 提出了 AlphaGo Zero 演算法。此算法雖然能達到超人（super human）等級，但同時需 5000 張 TPU 訓練。隨後，KataGo 與其他相關實作提出了多項改進，雖然在一定程度上降低了訓練所需的資源，然而仍需依賴大規模 GPU 叢集（cluster）。我所開發的圍棋程式 Sayuri，採用了包含大型卷積網路結構和結合 Gumbel 搜尋的演算法，使其在單張 RTX 4080 顯示卡上訓練出與 Facebook 開發的 ELF OpenGo 相當的棋力，計算資源需求約減少 250 倍。該程式已於 GitHub 開源：https://github.com/CGLemon/Sayuri

關鍵詞： AlphaZero、SE、LCB、Gumbel

## 1. 前言

DeepMind 在 2017 年提出 AlphaGo Zero 演算法，此算法在不需要人類專家的情況下，僅通過自我對戰便在圍棋上達到超人類（super human）等級的水準，而後又提出 AlphaZero 並在圍棋、西洋棋和將棋等領域達到一樣的等級。雖然如此，它卻相當消耗計算資源，在 GTX 1080Ti 上，Pascutto (2017) 估計完整重現此算法約要消耗 1700 年。經過幾年的演進，不同的實作提出了各種對於 AlphaZero 加速訓練的演算法或技巧，例如 Wu (2019) 在 KataGo 內通過加入各種改進，相比於 ELF OpenGo 達到 50 倍的加速性能。

Sayuri 是一個基於 AlphaZero 演算法設計的開源圍棋引擎，並且包含若干改進。在本研究中，主要的改進分別為：

1. Mixer Block：包含大卷積的網路結構，提昇感受野（Receptive Field）
2. SH-PUCT：結合 PUCT 和 Gumbel 的 sequential halving，盡可能少搜尋次數的同時並增加重複利用率

之後我使用一台 RTX 4080 的機器，通過自對戰學習 283.5 萬盤棋之後，將最終的權重和 ELF OpenGo 對比，並估計雙方使用的理論計算資源，Sayuri 約略減少 250 倍左右。相比於之前動輒數十至數千的 GPU 叢集，本研究成功在單張 GPU 上達到相同水平。

## 2. AlphaZero 的基本描述

AlphaZero 有三個重要的組成部件，分別為 Replay Buffer、Selfplay Workers 和 Training Worker。Replay Buffer 是一個 FIFO queue，它儲存最新的 $n_{replay}$ 盤自我對戰訓練資料並排出過舊的資料，Selfplay Workers 使用最新的權重自我對戰產生訓練資料，並放入 Replay Buffer 中，而 Training Worker 使用 Replay Buffer 內的資料訓練新的權重。

AlphaZero 在自我對戰時使用 MCTS（Monte-Carlo tree search）決定次一手並產生訓練資料。MCTS 在 AlphaZero 分為三個階段，Select、Expand and Evaluate 和 Backup ，這三個階段完成一次稱為一次模擬（或稱 playout）。Select 階段從根節點開始，會從根節點開始，每次都從自身所有子節點中選擇有最大 $PCUT(a)$ 值的節點下降，直至到葉節點。

$$
\begin{split}
\
\\ PUCT(a) = Q(a) + C_{puct} P(a) \frac{\sqrt{\Sigma_b{N(b)}}}{N(a)+1}
\\
\end{split}
$$

</br>

其中 $Q(a)$ 是節點 $a$ 的預期分數，$N(a)$ 是模擬時通過節點 $a$ 的次數，也稱為訪問次數，$C_{puct}$ 為平衡左右兩項的超參數，$P(a)$ 為節點 $a$ 之策略分數， $\Sigma_b{N(b)}$ 為所有子節點的訪問次數總和。

當到達葉節點時為 Expand and Evaluate 階段，此時擴展當前節點可能的後選手，並將網路推論的次一手機率分配到新增的子節點上，最後 Backup 階段從當前節點延當次模擬被選擇的路徑依序根新節點的預期分數 $Q$ 和訪問次數 $N$。

為了在自對戰階段探索不同的節點，我們在根結點上施加 Dirichlet Noise，讓低優先的節點也有機會被搜索。添加的方式如下，其中 $Dir$ 是 Dirichlet 分布。

$$
\begin{split}
\
\\ P(a) = 0.75P(a)_{raw} + 0.25 Dir(\alpha = 0.03)
\\
\end{split}
$$

</br>

在 $t$ 時刻決定次一手時，MCTS 會模擬 $n_{mcts}$ 次，建立一棵搜尋樹，並從根節點根據各子節點的訪問次數比例，得到搜索的機率 $\pi_t$，例如根節點共 A 、B 和 C 三個子節點，其訪問次數分別為 {A, B, C} = {20, 30, 50}，則相對應的機率為 {0.2, 0.3, 0.5}。$n_{mcts}$ 次模擬後，將得到的搜索機率 $\pi_t$ 和其當前盤面狀態 $s_t$ 並推入 replay buffer 中。當本次棋局完成後，可以得到最終結果 $r$ ，將 $r$ 根據 $t$ 時刻分別轉換成當前盤面狀態的得分 $z_t \in \{-1,0,1\}$ ，最後將 $z_t$ 和 replay buffer 中對應的 $\pi_t$ 和 $s_t$ 放在一起。

每當自我對戰下完 $n_{games}$ 盤棋後，即進入訓練階段，從 Replay Buffer 中隨機取出資料 $s$ 、$\pi$ 和 $z$，用以訓練神經網路，我們將網路定義為 $f_{\theta}$，代表權重為 $\theta$ 的網路，且 $p,\ v = f_{\theta}(s)$，$p$ 代表輸出的優先級分佈，$v$ 代表輸出的預期得分分數。網路的損失函數定義為下。

$$
\begin{split}
\
\\ loss = (z-v)^2\ - \pi^{T}\log{p} + c|\theta|^2
\\
\end{split}
$$

</br>

AlphaZero 使用 Residual neural network（ResNet）作為基礎的結構 [13]，一般而以 $N_{b}$ 和 $N_{c}$ 兩個參數表示網路大小，$N_{b}$ 代表 Residual 的 block 數目，$N_{c}$ 代表 block 使用的 channel 數目。

## 3. 文獻探討

### 3.1 Large Kernel Convolutional Layer

Transformer 不論在自然語言處理和影像辨識上都有突破性的表現。尤其 ViT 相較於傳統卷積神經網路，ViT 具有全局視野，可以更好地捕捉物體整體的特徵以提升準確度。甚至 Liu 等人於 Swin Transformer 中提出滑動窗口以消除 patch 間無法交換資訊的缺陷，並宣稱其性能達到 SOTA。為了彌補傳統卷積神經網路視野範圍太小的問題，Ding 等人提出了 RepLKNet，主幹採用 31 $\times$ 31 大小的深度分離卷積（Depthwise Separable Convolution），有效減少減算量，並使用 re-parameterized 的方法提升準確度以及穩定性。在 ImageNet 的分類任務上，其性能超越了 Swin Transformer。

<div id="replknet-rf" align="center">
</br>
<img src="https://hackmd.io/_uploads/S1sI6swole.png" alt="RepLKNet RF" width="512"/>
</br>
<b>圖一： RepLKNet 的視野範圍 </b>
</div>

</br>

傳統觀點認為，多個小型卷積層的堆疊可以模擬單一大型卷積層的感受野。然而，Ding 等人認為小型卷積層堆疊的效益無法達到理論上的性能。圖一展示了 Ding 等人使用不同網路計算的感受野範圍。從圖中可見，ResNet-152 的實際感受野相當狹窄，每個像素最終只影響其局部範圍，而擁有大型卷積核的 RepLKNet 的感受野則明顯較大。使用較大的卷積層固然能提升感受野範圍，但也伴隨著兩個主要問題：(1) 卷積核尺寸的計算複雜度為 $O(N^2)$。若使用 $31 \times 31$ 的卷積核，其理論計算量相較於 $3 \times 3$ 的卷積核將增加 106.78 倍。(2) 單純增加卷積層尺寸並不會線性地提升總體效能，過大的卷積核反而會導致網路效能退化。針對問題 (1)，RepLKNet 採用深度分離卷積取代傳統卷積。深度卷積的理論計算量比原始卷積層減少 $C$ 倍，其中 $C$ 為通道數量，因此能避免整體計算量增加過快的問題。針對問題 (2)，RepLKNet 採用重新參數化的方法來消除訓練不穩定的問題。在訓練階段，RepLKNet 同時並行訓練不同大小的卷積層。小型卷積層有助於提升性能和穩定性，避免單獨使用大型卷積層時的退化問題。而在推論階段，則將所有卷積層合併為一個等效的卷積層，在維持性能提升的同時，避免了推論階段的額外開銷。此合併過程如圖二所示。

<div id="replknet-rf" align="center">
</br>
<img src="https://hackmd.io/_uploads/BkW4knDjle.png" alt="RepLKNet RF" width="512"/>
</br>
<b>圖二： 重新參數化（re-parameterized）過程</b>
</div>

</br>

### 3.2 Gumbel 學習法

雖然實作上 AlphaZero 能訓練出強大的網路，但在數學上它並不保證學習能提昇網路表現。Danihelka 等人 (2022) 提出結合蒙地卡羅樹搜索（MCTS）和 Gumbel 的演算法，並且保證此算法能提昇網路表現。

Gumbel trick 依據機率 $p$ 採集樣本，將原本的機率 $p$ 轉換成對數值，並加入 Gumbel 噪音採樣。

$$
\begin{split}
\
\\ (g \in R^{k}) \ & \sim \ Gumbel(0)
\\ A \ & = \ argmax(g(a) + logit(a))
\\
\end{split}
$$

其中 $logit(a)$ 為 $p$ 中樣本 $a$ 的對數值，$A$ 為被採集的樣本，$Gumbel$ 為 Gumbel 分佈。

</br>

今將此模型使用在樹搜索上並混合節點的預期分數，加入預期分數項可以減少不好的節點被選擇的機率。令 $Actions_{root}$ 為根節點上所有可能合法手，即為根節點上所有的子節點，節點 $a \in Actions_{root}$ ，也就是說此模型僅在根結點才使用 Gumbel trick。

$$
\begin{split}
\
\\ A_t \ & = \ argmax(g(a) + logit(a) + \sigma(Q(a)))
\\ \sigma(Q(a)) & = (c_{visit} + Max(N(b))) \times c_{scaling}Q(a)
\end{split}
$$

</br>

其中 $Max(N(b))$ 為根節點中所有子節點裡訪問數最大的值， $c_{visit}$ 和 $c_{scaling}$ 為平衡輸出的超參數。$A_t$ 為時刻 $t$ 選擇的子節點。

Sequential Halving 會選擇 $n_{seq}$ 個後選手，每一輪會模擬 $n_{seq}$ 次，並淘汰剩剩餘的一半的後選手，直至剩下 1 個後選手或模擬次數達到 $n_{mcts}$ 次為止。而 Sequential Halving with Gumbel 使用 Gumbel trick 作為選擇並淘汰後選手的依據。

當完成 $n_{mcts}$ 模擬後，可以直接使用 $A_{n_{mcts}}$ 作為網路的學習目標 $\pi$ ，但更好的方法為 Learning an improved policy，將原本網路輸出的 logits 值混合子節點的 CompletedQ 值，並再次轉換為新的搜索率 $\pi$，此方法可以有效的提昇訓練表現。

$$
\begin{split}
\
\\ \pi &= Softmax(logits + \sigma(CompletedQ(a)))
\\
\\ CompletedQ(a) &= \begin{cases}
Q(a), \ \ \ if\ \ N(a) > 0
\\ v_{\pi}, \ \ \ otherwise
\\ \end{cases}
\\
\end{split}
$$

</br>

Gumbel 不只是在數學上保證能提昇網路表現，它還能大幅提昇訓練效率，相比於傳統的 AlphaZero-like 演算法，在達到相同強度的條件下，可以節省約有 5 到 20 倍的總模擬次數 [8]。

</br>

### 3.3 Playout Cap Randomization

</br>

## 4. Sayuri 之實作

### 4.1 Mixer Block

### 4.2 SH-PUCT

## 5. 實驗結果

## 6. 結論

## 參考文獻

~~[1] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap, Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel and Demis Hassabis, “Mastering the game of go without human knowledge,” Nature, 550:354–359, 2017.~~

~~[2] David Silver, Thomas Hubert, Julian Schrittwieser, Ioannis Antonoglou, Matthew Lai, Arthur Guez, Marc Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy Lillicrap, Karen Simonyan, Demis Hassabis, “A general reinforcement learning algorithm that masters chess, shogi, and go through selfplay,” Science, 362(6419):1140–1144, 2018.~~

~~[3] David J. Wu, “Accelerating Self-Play Learning in Go,” 2019, arXiv preprint, arXiv:1902.10565.~~

~~[4] Yuandong Tian, Jerry Ma, Qucheng Gong, Shubho Sengupta, Zhuoyuan Chen, James Pinkerton, and C. Lawrence Zitnick, “Elf opengo: An analysis and open reimplementation of alphazero,” In Thirty-Sixth International Conference on Machine Learning, 2019.~~

~~[5] Jie Hu, Li Shen, Samuel Albanie, Gang Sun, and Enhua Wu, “Squeeze-and-excitation networks,” In IEEE Conference on Computer Vision and Pattern Recognition, pages 7132–7141, 2018.~~

~~[6] Seth Troisi, “[Experiment] Squeeze and Excitation,” 2019, Minigo project issue, https://github.com/tensorflow/minigo/issues/683 .~~

~~[7] Ti-Rong Wu, I-Chen Wu, Guan-Wun Chen, Ting han Wei, Tung-Yi Lai, Hung-Chun Wu, and Li-Cheng Lan, “Multi-labelled value networks for computer go,” IEEE Transactions on Games, 10(4):378–389, 2018.~~

~~[8] Ivo Danihelka, Arthur Guez, Julian Schrittwieser, David Silver, “Policy improvement by planning with Gumbel,” In International Conference on Learning Representations, ICLR, 2018.~~

~~[9] Gian-Carlo Pascutto, “Leela Zero project main webpage,” https://zero.sjeng.org/, last accessed: 2022-03-14.~~

~~[10] Gian-Carlo Pascutto, “[Computer-go] Zero performance,” http://web.archive.org/web/20190205013627/http://computer-go.org/pipermail/computer-go/2017-October/010307.html, last accessed: 2022-03-14.~~

~~[11] Henrik Forstén, “Add O(sqrt(log(n))) scaling to tree search,” 2018, Leela Zero project issue, https://github.com/leela-zero/leela-zero/pull/2072 .~~

~~[12] Henrik Forstén, “Pick chosen move based on normal distribution LCB,” 2019, Leela Zero project issue, https://github.com/leela-zero/leela-zero/pull/2290 .~~

~~[13] Kaiming He, Zhang, Xiangyu Zhang, Shaoqing Ren, “Jian Sun. Deep residual learning for image recognition,” 2016, In IEEE Conference on Computer Vision and Pattern Recognition, 770–778.~~

## 雜物區
[GCP 1700 years](http://web.archive.org/web/20190205013627/http://computer-go.org/pipermail/computer-go/2017-October/010307.html)
[KataGO 論文](https://arxiv.org/abs/1902.10565)
[AGZ 論文](https://discovery.ucl.ac.uk/id/eprint/10045895/1/agz_unformatted_nature.pdf)
[AZ 論文](https://arxiv.org/abs/1712.01815)

[Resnet Paper](https://arxiv.org/abs/1512.03385)