Three-Head Neural Network Architecture for AlphaZero Learning

1 INTRODUCTION . Computational advancements of deep neural networks ( Bengio , 2009 ; LeCun et al. , 2015 ) have led to progresses in a number of scientific areas , ranging from image classification ( Krizhevsky et al. , 2012 ) , speech recognition ( Graves et al. , 2013 ) , natural language processing ( Collobert et al. , 2011 ) , playing Atari games ( Mnih et al. , 2015 ) , to the prominent achievement of AlphaGo ( Silver et al. , 2016 ) , AlphaGo Zero ( Silver et al. , 2017b ) , and AlphaZero ( Silver et al. , 2017a ; 2018 ) learning systems for playing multiple classic two-player alternate-turn perfect-information zero-sum games — games of which kind have been a research object as old as the subject of artificial intelligence itself . Arguably , the accomplishments of AlphaGo and its successors represent a result of continual effort from two lines of research directions . The heuristic search ( or sometimes AI planning ) community studies various intelligent search strategies ( Pearl , 1984 ) ( e.g. , best-first search ) given that a space of states can be generated gradually by description of the problem ( e.g. , rules of Go ) , the goal is to find a desired solution using smaller computational cost by exploring promising regions preferably ; the study of machine learning attempts to summarize a general concept from experience data , therefore enabling probably correct prediction even in unseen scenarios . The combination of both learning and search techniques dates back to Samuel ’ s studies in checkers ( Samuel , 1959 ) , where two learning methods , rote learning and generalization learning , were studied and integrated into a look-ahead minimax search ; the later development leveraging advancements from both learning and search techniques ( e.g. , training on games produced by stronger players with more expressive function approximation for better generalization learning and using better αβ pruning ( McCarthy et al. , 2006 ) in search ) achieved better results ( Samuel , 1967 ) . Noticeably , the learning techniques are only used to construct an evaluation function for checkers positions from historical playing records either generated by human or computer players . However , in certain games ( e.g. , chess ) where a large body of human players exist , manually constructing a heuristically correct evaluation function is also possible . Therefore , rather than attempting to automatic construct an evaluation functions , much early research effort focus on analyzing formal properties of search algorithms given certain assumption on the correctness of an existing evaluation function ( Knuth & Moore , 1975 ; Pearl , 1980a ) , proposing new search routines that may lead to more efficient pruning ( Pearl , 1980b ; Plaat et al. , 1996 ) , or devising procedures when the assumed correctness of an evaluation failed to consign in practice ( Harris , 1975 ) . Indeed , optimizing using minimax upon an approximated evaluation has been regarded as a “ deadly sin ” in the eye of statistics , since such an optimization could only compound the approximation error , therefore resulting pathological behavior of search ( Nau , 1982 ) , i.e. , the deeper the search the worse of the obtained evaluation . However , in practice , pathological behavior was seldom observed in a number of games . Pearl ( 1983 ) presented several observations that may explain the phenomenon , such as the approximation error tends to decrease as search goes deeper , the ubiquitous of perfect evaluation in some branches of the search , the fact that the search space is a directed acyclic graph rather than a tree , and the relatively small branching factor of the game of interest . Game-searching by minimax continued to triumph as computers become faster and more general or game-specific searching techniques are introduced , culminating to the successes of champion playing strength against human players in games such as checkers ( Schaeffer et al. , 1992 ) , chess ( Campbell et al. , 2002 ) , Othello ( Buro , 1998 ) and backgammon ( Tesauro , 1995 ) . It follows that game AI research began to switch attention on several more challenging games where reliable human-directed evaluation function is difficult to construct and the branching factor is significantly larger , such as the game of Go ( Müller , 2002 ) and the game of Hex ( Van Rijswijck , 2002a ) . Coulom ( 2006 ) proposed a Monte Carlo tree search ( MCTS ) for computer Go ; it uses Monte Carlo simulation to sidestep the need of a static evaluation function . From the heuristic search perspective , MCTS can be regarded as an instance of the general best-first search depicted by Pearl ( 1984 ) . Different from αβ pruning , which relies on an accurate move ordering for preferable exploration , best-first search iteratively adjusts its search effort by summarizing the new information backed from leaf evaluations , at the expense of explicitly storing the whole search tree in memory . To select a leaf node to expand and evaluation from the search tree , a UCT search can be used ( Kocsis & Szepesvári , 2006 ) for a balance between exploration and exploitation . Notably , unlike αβ minimax search , MCTS does not apply hard-min or max when backing up approximated leaf evaluation but uses average to summarize all previous evaluations while preserving asymptotic convergence . The superiority of MCTS is also manifested by its flexibility of integrating learned heuristic knowledge into the search ( Gelly & Silver , 2007 ) . That is , by incorporating an existing move-ranking function for better leaf node selection ( i.e. , progressive bias by Chaslot et al . ( 2008 ) or PUCT by Rosin ( 2011 ) ) , the search can focus on only promising moves therefore effectively reducing the high branching factor ; by leveraging more informed pattern-based Monte Carlo playout ( Coulom , 2007 ) , better leaf evaluation can be obtained and backed up therefore further enforcing the best-first behavior , although in practice stronger playout policy for leaf evaluation does not always result in stronger performance given the same finite computation of MCTS ( Silver & Tesauro , 2009 ) . Continual effort were made towards adding more accurate prior knowledge to MCTS . The successes of deep neural networks in other AI domains , such as image classification ( Krizhevsky et al. , 2012 ) , revived the application of neural networks for approximate human ’ s evaluation in Computer Go ( Maddison et al. , 2015 ; Clark & Storkey , 2015 ; Tian & Zhu , 2015 ) , leading to the development of AlphaGo Fan ( Silver et al. , 2016 ) which uses policy network for providing prior knowledge in move selection and value network for leaf node evaluation in MCTS — it became the first program defeating human professional player in the game of Go . Subsequently improved versions AlphaGo Lee defeated a more pronounced world champion player in public match . A even stronger version called AlphaGo Master was strong enough that it never lost any game against any top human players in a set of online match . The commonality of these programs is that the neural networks used in search were all first trained on human professional game records . Later , it was found that even without using human data to initialize the neural network weights , a closed loop learning paradigm could still learn MCTS players stronger than AlphaGo Lee and AlphaGo Master ( after doubling neural net size ) ; the paradigm is named as AlphaGo Zero ( Silver et al. , 2017b ) . By further removing other Go specific encoding in AlphaGo Zero , a more general version — AlphaZero ( Silver et al. , 2017a ; 2018 ) — was applied to chess , Shogi and Go , producing strong players in all these games after separately training on each of them . We summarize the major differences between AlphaGo Master , AlphaGo Zero and AlphaZero in Figure 1 . However , one limitation of the Zero algorithms ( AlphaGo Zero and AlphaZero ) is that a large computation cost could be needed because they use search to produce millions of training games , where each move in a game was produced by calling neural networks hundreds of times for building the search-tree — AlphaZero used thousands of TPUs in its week-long training on the game of Go ( Silver et al. , 2018 ) . A three-head neural network architecture was proposed in ( Gao et al. , 2018b ) which can learn an additional action-value head on search-produced training dataset , akin to the two-head network used in AlphaGo Master , AlphaGo Zero and AlphaZero . In this paper , we conduct an empirical comparative study on AlphaZero style learning using two- and three-head neural networks . We find that , in comparison to the two-head counterpart , three-head network can also lead to improvement in AlphaZero style learning paradigm . 2 ANATOMY OF ALPHAGO ZERO AND ALPHAZERO . We provide an anatomy to AlphaGo Zero and AlphaZero , along with discussion with their predecessors AlphaGo Fan , AlphaGo Lee and AlphaGo Master . See Appendix A for preliminaries on Monte Carlo tree search . 2.1 POLICY VALUE MONTE CARLO TREE SEARCH . MCTS is an essential component in AlphaGo Zero and AlphaZero since it is the algorithm upon which training data for the neural network are produced ; that is , the policy network is trained by mimicking the move probabilities generated by MCTS , while the value network is to learn to predict game result generated by MCTS players . A two-head policy-value network fθ provides two estimates given one input game-state , i.e. , p , v ← fθ ( s ) where s is a game-state , p is a policy over actions at s and v is an estimate of the optimal state-value of s. After one evaluation , policy and value estimates are added to the Monte Carlo tree search ( PV-MCTS ) in the following way . At each iteration of MCTS , after a leaf node s is selected , the game state of that leaf node is then passed to neural net for inference , p , v ← fθ ( s ) . The value estimate v is backpropaged to the search tree . The leaf node is subsequently expanded , for each action a from s , the newly created child node stores : { N ( s , a ) = 0 , W ( s , a ) = 0 , Q ( s , a ) = 0 , P ( s , a ) = pa } , where N ( s , a ) the visit count of a at s , W ( s , a ) is the number of wins , Q ( s , a ) is the accumulated action-value ( equivalent to negated V ( s′ ) is acting a at s results s′ ) , and P ( s , a ) is the prior probability for action a. AlphaGo Fan used an expansion threshold nth > 0 , therefore a leaf node is expanded only when it has been visited more than nth times . AlphaGo Zero and AlphaZero used an expansion threshold of 0 because for each selected leaf node , to back up an estimate , neural network has to be called ; however , the two-head neural network provides two estimates p and v ; naturally , to not waste GPU computation , probabilities in p has to be saved in newly created child nodes by node expansion .

Summary: The paper applies three-head neural network (3HNN) architecture in AlphaZero learning paradigm. This architecture was proposed in [1] and the paper builds upon their work. In AlphaGo and AlphaZero 2HNN is used, which predicts policy and value for a given state. 3HNN also predicts action-value Q function. In [1], the three new terms are added to the loss to train such a network, and the network is trained on a fixed dataset. The paper utilizes the same 3HNN idea with the same loss, and the contribution is that 3HNN is trained synchronously with MCTS iterations on an updating dataset (“AplhaZero training style”). Learning speed of 3HNN is shown to be higher than that of 2HNN. The special attention is drawn to varying the threshold expansion parameter, as the 3HNN architecture allows to set it above zero, while 2HNN does not. The approach is demonstrated on the game of Hex. Results are presented on two test datasets: positions drawn from a strong agent’s games and random positions. Labels in both datasets are perfect, obtained by a special solver.