Implementing Fast AlphaZero

A few teammates and I implemented AlphaZero at Cogito NTNU. I expected a few weekends of work after reading the paper, but the project took months. Much of that time went into making self-play and network evaluation fast enough to train at a useful pace.

How AlphaZero works

AlphaZero uses a neural network to guide Monte Carlo tree search. The network predicts a policy over possible moves and a value for the current position, while the search uses those predictions to explore promising moves. In our implementation, training followed this loop:

1. Self-play to generate training data.
2. Train the neural network on that data.
3. Test whether the new network beats the previous one.

The network has two heads sharing a convolutional trunk. The policy head predicts move probabilities, and the value head estimates the outcome from the current player’s perspective as a number in \([-1, 1]\):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

import torch
import torch.nn as nn

class PolicyValueNet(nn.Module):
    def __init__(self, board_size: tuple[int, int], action_size: int, num_channels: int = 256):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, num_channels, 3, padding=1),
            nn.BatchNorm2d(num_channels),
            nn.ReLU(),
            nn.Conv2d(num_channels, num_channels, 3, padding=1),
            nn.BatchNorm2d(num_channels),
            nn.ReLU()
        )
        
        board_size_flat = board_size[0] * board_size[1]
        self.policy_head = nn.Sequential(
            nn.Linear(num_channels * board_size_flat, action_size),
            nn.LogSoftmax(dim=1)
        )
        self.value_head = nn.Sequential(
            nn.Linear(num_channels * board_size_flat, 1),
            nn.Tanh()
        )

A snapshot of our MCTS implementation on a Tic-Tac-Toe board, with the search tree on the right showing visit counts per node, starting from 500 at the root.

The loss function

The network learns by minimizing a single combined loss:

\[L = (z - v)^2 - \pi^T \log p + c\|\theta\|^2\]

Here \(z\) is the game outcome, \(v\) is the predicted value, \(\pi\) is the move distribution from MCTS, and \(p\) is the network’s policy. The final term applies L2 regularization to the weights.

Making it fast

Once the algorithm was working, training speed became the main limitation. The way we scheduled searches and passed positions to the network made a large difference.

We got the largest improvement by running MCTS simulations in parallel and batching their network evaluations:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

import torch.multiprocessing as mp

class ParallelMCTS:
    def __init__(self, num_processes: int = 8, batch_size: int = 32, num_simulations: int = 800):
        self.pool = mp.Pool(num_processes)
        self.batch_size = batch_size
        self.num_simulations = num_simulations
        
    def run_batch_evaluation(self, positions):
        with torch.no_grad():
            policies, values = [], []
            for i in range(0, len(positions), self.batch_size):
                batch = positions[i:i + self.batch_size]
                batch_policies, batch_values = self.network(torch.stack(batch))
                policies.extend(batch_policies)
                values.extend(batch_values)
            return list(zip(policies, values))

Evaluating one position at a time left much of the GPU unused. Batching gave it more work per call and reduced the time spent waiting on individual evaluations.

We also cached evaluations because self-play often revisited the same positions, especially early in a game:

1
2
3
4
5
6
7
8
9

import numpy as np

class KVCache:
    def __init__(self, max_size: int = 100000):
        self.cache: dict[str, tuple[np.ndarray, float]] = {}
        self.max_size = max_size
        
    def get_cached_value(self, board_hash: str):
        return self.cache.get(board_hash)

Results

With parallel MCTS, batched evaluations, and the position cache, we measured these changes relative to our earlier implementation:

• Self-play got 3.2× faster.
• Training was 2.8× faster.
• Memory usage dropped by 45%.

We also experimented with a path consistency loss, which encourages value predictions along a search path to agree:

\[L_{PC} = \|f_v - \bar{f}_v\|^2\]

Training improved with this extra term, although I don’t have a clear explanation for the size of the improvement we saw.

From the paper to working code

The algorithm in the paper gave us a starting point, but it left plenty of implementation decisions to work through. Batching, caching, and handling unusual self-play states took a large share of our time.

Those details accounted for much of the gap between our first working version and the version we could train efficiently. It was useful to measure them separately, since a slow training run didn’t necessarily mean the learning algorithm was wrong.

The implementation is on GitHub.