Amortized Planning with Large-Scale Transformers: A Case Study on Chess

We thank the reviewer for their insightful and positive feedback.

Is the dataset of poor quality if you use Stockfish with only 50ms evaluation time (as there is a significant Elo disparity between 0.05s and 1.5s for Stockfish)? Also, why does Stockfish with 0.05s evaluation time take 1.5s to play a move on average?

The dataset is of high (super grandmaster) quality, as evidenced by the fact that Stockfish with 0.05s, i.e., the data-generating oracle, achieves an Elo of 2713 against other bots on Lichess. While higher time limits would lead to even stronger annotations, there is a trade-off in terms of computing time spent on collecting the dataset (currently 8680 days of unparallelized compute). Note that we did, however, create smaller-scale datasets with higher annotation time limits in Table A2, showing that none of our networks can currently fully match the performance of the data-generating oracle(s), so the current frontier seems to lie with improving models rather than the playing-strength of the data-generating oracle.

To clarify the 0.05s vs. 1.5s evaluation time, we consider two versions of Stockfish:

• Data-generating oracle: For a given board, the oracle evaluates every legal move separately for 0.05s (i.e., if there are, e.g., 30 legal moves, Stockfish uses a total time of 1.5s to play a move). We primarily compare to this version since this is how we constructed the dataset.
• “Standard-play” Stockfish: For completeness, we also include Stockfish as it would typically be used in standard play, i.e., its evaluation time is restricted per state but not per move, meaning that some clearly suboptimal moves can receive very little time (based on Stockfish’s pruning) and the resulting extra budget is instead spent on more promising moves. Since there are roughly 30 legal moves per board on average, we choose a time limit of 30 * 0.05s = 1.5s per board state for “standard-play” Stockfish. This leads to an improvement in play of >200 Elo compared to the version that spends a fixed 50ms on each legal move per board state, but we only include this baseline as a “reference” (i.e., it is not reflected in the dataset).

We agree that the difference between these two versions is subtle, and we will clarify it in the next revision of our paper.

How long did it take to train each model?

The 9M models trained at roughly 270 steps per second, yielding a total training time of 10M / 270 * 3600) = 10.2 hours. The 136M models trained at approximately 26 steps per second, yielding a total training time of 10M / (26 * 3600 * 24) = 4.45 days. The 270M models trained at roughly 13 steps per second, yielding a total training time of 10M / (13 * 3600 * 24) = 8.9 days. We will add these details to Appendix A.6, which describes our hardware setup.

Did you consider non-uniform binning to improve accuracy in critical positions?

We initially experimented with non-uniform binning but quickly abandoned the approach due to its increased complexity (e.g., state-value expansion is non-trivial with non-uniform bins due to its minimax nature) and failure to produce large performance gains. However, we did ablate the number of bins in Table A2, showing that an increased resolution improves performance, but only up to a certain point. Based on our experience, non-uniform binning can lead to gains when the overall number of bins is very small, but these gains diminish relatively rapidly with medium to large numbers of bins.

We thank the reviewer for their helpful feedback.

Is the transformer or the dataset more important for performance improvement?

Although this cannot be determined with certainty given our results, we investigated this question in Table A2, where we ablated model architecture and the time limit used when annotating a move with Stockfish for the dataset creation. Table A2 shows that using a higher time limit does lead to some performance gains (e.g., 1.4% higher puzzle accuracy) but at an extremely high computational cost (i.e., doubling the total unparalleled annotation time from 8680 days to 17360 days). Therefore, we compromise between computational effort and final model performance and use 0.05s per state-action value annotation.

Using transformers for amortized planning is a contribution, but not significant enough.

We respectfully disagree. Though the approach of amortizing a planning algorithm with a fixed-parametric function approximator may be evident to some researchers, it is typically not the mainstream interpretation. Large sequence models have been repeatedly dismissed as mere “statistical pattern matchers”, “stochastic parrots”, or “curve fitters”; implying that this approach is insufficient to capture complex algorithmic behavior. While the results of large foundation models may speak for themselves w.r.t. this criticism, we believe it is important to advocate and thoroughly test the amortization viewpoint. Just to reiterate: the state space of chess is very large - even when only considering the “natural” distribution of board states from games on lichess.org, all but the most trivial games face our networks with mostly unseen board states, where accurate value estimates and/or actions require considerable generalization beyond the training data. Our results cannot be explained as memorization with a bit of interpolation or simple statistical pattern matching, and not too long ago, such results were thought to be unachievable without explicit planning. Our results also show that transformers (up to the size used in our experiments) cannot always fully match Stockfish. Together, this raises two important future questions: (i) What is missing to match Stockfish’s performance; is it just a matter of scale, or can we identify systematic shortcomings of transformers? (ii) How do transformers implement amortized planning, and how is it related to architectural parameters such as depth? While these questions are beyond the scope of our current work, we believe that our results and dataset lay excellent groundwork to investigate these questions in the future.

We thank the reviewer for their careful study of our paper and, in particular, for all their clarifying remarks on lc0, which have helped us to improve the description of that baseline.

Lc0’s T82 network is a domain-specific transformer (rather than a ConvNet), trained using supervised learning on an unknown amount of MCTS data (rather than using online RL).

Thank you for these clarifications and all the additional information on lc0! We agree with the reviewer that these details are not well-documented, which makes a scientific comparison beyond evaluating playing strength difficult (e.g., it is unclear which and how much data was used to train T82). We have updated the description of lc0 in Section 2.5 as follows:

“Leela Chess Zero We consider three variants: (i) with 400 MCTS simulations, (ii) only the policy network, and (iii) only value network (where (ii) and (iii) perform no additional search). We use the T82 network, which is a transformer and the largest network available on the official Leela Chess Zero website [5]. T82 uses one token per square of the chess board, has fewer parameters (94M vs. up to 270M for our models), includes some domain-specific architecture changes (e.g., Smolgen [29]), and was trained using supervised learning on MCTS data.”

The grandmaster claim against humans may only hold with Blitz time controls (since human performance improves with more evaluation time, unlike the models trained in this paper), which isn’t always clear in the paper.

We fully agree that our results only imply grandmaster-level performance against humans with Blitz time controls, and we will clarify this distinction throughout the paper. For example, we have changed line 44 to “... playing chess at a high level (grandmaster) against humans with Blitz time controls”.

BC does worse than value predictions because it performs a weak form of “1-ply search” rather than having a poorer training signal (lc0’s/AZ’s policy nets are also worse even though they imitate the full distribution, i.e., not just the best move).

This is an interesting observation. Indeed, multiple factors probably contribute to BC’s relatively weak performance, and we agree that inference-time FLOPS (i.e., performing a less expansive “1-ply search” than the value-based methods) is probably also a major factor. We have expanded the discussion of these results in Appendix B5 accordingly.

AZ and lc0 operate on the past 8 board states rather than the full move history.

Thank you for the clarification! We agree that the reviewer’s proposed distinction of “PGN” for GPT-3.5 and “FEN + past states” for AZ and lc0 is more precise and have updated Table 1 and the description in Section 3.1 accordingly.

Did you try training a single model with three prediction heads for the different predictor targets (AV, BC, SV) to enable transfer between them (e.g., more data for BC)?

Thank you for this interesting suggestion! We have not tried training a single shared model with multiple heads. A priori, it is not clear whether this would increase overall performance, lead to catastrophic interference, or make little difference overall since all three targets are highly related. It is also unclear whether increased model capacity would be needed (or not) and how to combine all three predictions into a single policy (though there are obvious candidates). We believe that investigating these questions is out of scope for our current work, and have added them to a future work section in the discussion of our paper.

Note that we did conduct an ablation where all three predictor targets were trained on the same amount of data in Table A4 to investigate whether our observed differences in performance were simply a matter of different amounts of data.

Why is the state-value predictor not trained on the (much larger) state-action-value dataset, using the state-action values as state-values for the next state?

Indeed, that is an astute observation, which would cut the annotation time in half when creating the state value and action value datasets (assuming that Stockfish’s evaluation is symmetric). For our paper, we only trained the action value predictor on our largest dataset of ~15 billion data points. However, we also conducted the ablation where all three predictor targets are trained on exactly the same amount of data in Table A4, showing that the state value predictor performs more or less on par (slightly better, by 12 Elo) with the action value predictor. A priori, it seems plausible that learning to predict state and action values is of roughly equal complexity (in deterministic environments), but repeating our experiment from Table A4 (i.e., using the same number of training data points for all policies) at even larger scales would be an interesting direction for future work.

We thank the reviewer for their constructive feedback.

The paper claims to aim at approximating Stockfish by distilling its board evaluation into a transformer but evaluates performance w.r.t. chess play rather than Stockfish’s centipawn score, which invalidates this claim. Instead, the paper appears to create a strong chess engine without using explicit search, which, however, underperforms lc0.

Thank you for pointing out a potential source of confusion!

We do evaluate the learning targets (i.e., the loss over win percentage/centipawns) in Appendix B.4 (Fig. A2 and A3). These loss curves show the quantitative comparison against the discretized ground-truth centipawn score.

We respectfully disagree that our goal was not to distill Stockfish’s value estimation behavior. The objective we directly optimize (and report in the appendix) is our model’s prediction error w.r.t. Stockfish’s value estimates, and as a modeling assumption, we chose to phrase this as a discretized prediction problem (similar to distributional RL). This is the main idea behind general algorithm distillation (i.e., minimizing prediction error over an algorithm’s outputs). We mainly focused on two behavioral metrics for comparing models since they imply good quantitative predictions and allow easy comparison across learning targets and models. One could think of these metrics as using a chess-relevant distortion function. If, instead, our goal were to build a strong chess engine, we would have used explicit search and a domain-specific transformer. The fact that lc0’s policy net performs ~160 Elo better is interesting but does not undermine the main goal of our paper (which is to carefully investigate various architectural and training factors, not to outperform).

Why are transformers interesting for this study if prior work showed that convolutional networks achieve strong performance on chess?

We focus on investigating transformers’ amortized planning capabilities (in the context of chess). While convolutional networks have been used successfully in the past (e.g., AlphaZero, or past versions of lc0), the more recent consensus seems to be that transformer-based architectures are at least equally suitable, perhaps even stronger. For example, the strongest modern lc0 networks are transformer-based, as also pointed out by R-LZwX, which is the result of many ablations by the lc0 community. While it is intuitive that 2D convolutions implement good inductive biases for the 2D chess board, it is also plausible that (self-)attention leads to good inductive biases for considering relative relations between sets of relevant pieces regardless of their precise spatial arrangement on the board. To avoid speculation, we also performed an ablation with a convolutional network in our paper (see Table A2).

What does it mean that lc0’s source code is available but that there is no (peer-reviewed) research report explaining and comparing lc0’s methods? Does this mean lc0 may include Stockfish already? Is it still fair to compare with lc0?

The difficulty with lc0 is that it is a collection of models (originally re-implementing AlphaZero, but with significant evolution since then) that is developed by an active chess engine community, which makes heavy use of the lc0 blog and discord server to share results. This means that detailed architectural information is often only available via the source code (~90k lines of code), and, more importantly, full details about the training process and data sets are typically not part of the open-source release and are thus simply not available for many lc0 networks (though we cannot rule out that this information could be found on lc0’s discord server). For example, even R-LZwX, who seems very knowledgeable about lc0, is not fully aware of all the training details. While the lc0 community has made impressive contributions and progress over AlphaZero, the lack of independent reproducibility due to missing information makes the networks currently less suitable for academic research (one of the key requirements of peer-reviewed research articles is providing all information necessary for independent reproduction of results). One positive side-effect that our paper and dataset/model release hopefully have is to encourage and enable others to evaluate their methods in a way that allows for scientific reproducibility.

Consequently, we do not know whether lc0 includes Stockfish evaluations. However, this does not have an impact on our work and main claims. Similarly, whether comparing to lc0 is “fair” cannot be answered without knowing the full training details of lc0, which is a caveat we clearly state in our paper (lines 190 - 192): "We show these comparisons to situate the performance of our models within the wider landscape, but emphasize that some conclusions can only be drawn within our family of models and the corresponding ablations that keep all other factors fixed."

Can you conduct ablation studies on the training objectives?

Yes, we have already conducted the following ablations of the training objectives:

• loss function (Table 2): log-loss (classification) vs. HL-Gauss loss (classification) vs. L2 loss (regression)
• predictor-target (Table 2): action-value vs. state-value vs. behavioral cloning
• number of value bins (Table A2): 16 to 256 bins for the HL-Gauss loss

If the reviewer has any other ablation of the training objectives in mind, we are happy to include it in the next revision of our paper.

Did you try to balance the labels and make sure that they are semantically interconnected?

We did not balance the labels as we did not want to perturb the natural game/value distribution. We are not quite sure what the reviewer means by “semantic interconnectedness”. Could you please clarify? Note that we did perform ablations on the number of bins we used for our discretization (see Table A2) and an ablation on a non-discretized loss (see Table 2).