Implicit Search via Discrete Diffusion: A Study on Chess

We sincerely thank Reviewer KzTF for your review and are grateful for the time you spent on our submission. We’re pleased you find our paper novel and intuitive. Below, we provide a point-by-point rebuttal to clarify your concerns.

W1: Line 83. The link to the source code is invalid. I hope to see the code during the rebuttal.

Thanks for your interest. We have attached the code above.

W2: I would expect more detailed explanation of Figure 1 in both the main texts and the caption of Figure 1.

Thanks for the suggestion. MCTS explicitly performs action selection, state evaluation, and value backup in an iterative manner before determining the final action to take, while discrete diffusion implicitly gathers future information during the process of future imagination. We have optimized Figure 1, provided additional explanations in the caption, and placed the more detailed content about explicit search via MCTS in the Appendix A.

W3: The authors should provide a brief explanation of each input parameter and its role in Algorithm 1. For instance, what is $\lambda_t$? And a curiosity question: can you use other random sampling methods to draw t in line 169? Like Gaussian? Will different sampling methods affect the algorithm's performance?

Thanks for the suggestion. $\lambda_t$ is a time-dependent reweighting term (Line 193) that assigns lower weight for noisier $x_t$, which is derived from the KL term in Equation (1) as proved in Appendix B.2.

Yes. During testing, all values of t from [1, T] will be used to denote the denoising phase, so the model should be taught to handle all t values. Other sampling methods beyond uniform sampling are also possible. For example, loss-weighted sampling dynamically adjusts based on the model's learning situation, e.g., by increasing the sample proportion of the currently learned, poorer-performing t, which can lead to a slight improvement but within 1 % action accuracy. Gaussian sampling primarily focuses on sampling more frequently around the mean value and the performance in learning t in low-density regions may not be well. We found Gaussian underperforms uniform sampling by around 3% action accuracy. We've added sampling details in the updated version.

Q1: In Figures 3(a) and (b), why DiffuSearch only runs around 7 depths? Are there any technical limitations that might prevent running DiffuSearch at greater depths?

Thanks for mentioning this point. We have already observed a significant improvement from step 1 to step 7, so we have not considered increasing the depth further. We expect that continuing to increase the depth will lead to further performance improvements. However, the main reason is due to resource and time constraints instead of technical limitations. For example, collecting data previously required 1 week using 1024 CPUs (with a depth mostly around 8), and data with a depth of 24 or 25 might take 3 weeks by trebling the search time in Stockfish. Therefore, we are unable to provide a specific number for depth 24 or 25 at this time. However, this is a good suggestion to investigate the convergence of depth in DiffuSearch, and we will update our manuscript once we have results.

Q2: I noticed that the demonstration plots like Figures 4 and 5 are limited to the comparison between DiffuSearch and Transformer (S-A). Can the authors explain the possible reasons? What about DiffuSearch against Transformer (S-V), against Transformer (SA-V), and Transformer (100 MCTS simulations)?

Thanks for the question and suggestion. In the latest manuscript, we have also added the results for Transformer S-A, Transformer SA-V, and Transformer with MCTS in Figure 6 to make the paper more complete.

We hope our response could address your questions!

We sincerely thank Reviewer SREQ for the review and are grateful for the time you spent with our submission. We wish to address your confusion and concerns by providing detailed responses to each of your comments.

W1: The resulting Chess program is very weak compared to the current state of the art. Lc0 or Stoofvless use MCTS and deep neural networks and have a Elo greater than 3500 according to the Swedish Rating List. This is far above the 1728 Elo of DiffuSearch.

Our research question is to explore whether there are alternative solutions to enhance the planning abilities of LLMs without relying on explicit search techniques like MCTS for solving complex problems. We introduce DiffuSearch based on discrete diffusion as our solution, using chess as our study task, where explicit search is known to be essential. Therefore, the core objective of our work is not to achieve SOTA Elo performance that surpasses all public engines in chess (e.g., lc0 which often relies on a large amount of training data and domain-specific improvements), but rather to use chess as a case study to investigate the above research question through controlled experiments (e.g., data and model size), which can provide insights into building LLMs with enhanced reasoning and planning capabilities.

Q1: What would you propose to improve the Elo of your program?

In this paper, we frame our discussion within a similar setting to that of [1], where we achieve improvements using fewer resources. We propose DiffuSearch, a model that performs implicit search by looking into future states through discrete diffusion modeling, and we have found that effectively modeling the future world within the policy model can enhance the performance (e.g., Elo on chess) without relying on explicit search techniques like MCTS.

The transformer-based one-step policy baselines studied in our work can be also seen as a similar smaller version of that presented in the recent lc0 blog, where the lc0 team has find it outperforms original convolution-based models. Given the experiments in our work, we anticipate that our approach could further enhance state-of-the-art systems as well, although this is not the primary focus of our research question in this paper.

[1] Ruoss, A., Delétang, G., Medapati, S., Grau-Moya, J., Wenliang, L. K., Catt, E., ... & Genewein, T. (2024). Amortized Planning with Large-Scale Transformers: A Case Study on Chess. NeurlPS'24.

We sincerely thank Reviewer LQsg for your review and are grateful for the time you spent on our submission. We are also glad you think our methodology is novel and offers a new perspective on implicit search in AI. Below we would like to give detailed responses to each of your comments.

W1: Declining prediction accuracy: As shown in Figure 2 (left), the accuracy of predicted future states and actions declines significantly for steps further into the future. For a strong world model in chess, the lookahead should ideally be accurate for about 7 steps, similar to top engines like Stockfish.

Thank you for your insightful comments. We believe that one significant factor contributing to the diminishing effectiveness of lookahead as the future steps increase is related to the volume of data. In Figure 2, we evaluated a model trained on 10k games (660k records), and we further discovered that scaling the data to 100k games (6.6M records) in the Table below (Table 8 in the paper) significantly improved performance. In the world model construction (Valid $s_i$ and Match $a_{i-1} - s_i$), the model achieves over 90% accuracy within 3 steps, while the 10k data trained model declined to around 50% in predicting valid future states. Therefore, we are optimistic that further scaling of the data could also enhance the world construction accuracy for the 7 steps.

Additionally, our goal is not to build a highly competitive chess engine, such as the ones that rely on extensive training like Lc0. Instead, we hope to propose a different paradigm for solving complex problems, beyond the one-step policy with explicit search, which may provide insights into building LLMs with enhanced reasoning and planning capabilities.

W2: Training complexity: The paper doesn't provide a clear comparison of the computational requirements (e.g., FLOPs) for training DiffuSearch versus traditional transformer models. This makes it difficult to assess the scalability and efficiency of the diffusion process compared to other approaches.

Thank you for your feedback. In Figure 2 (middle) and Appendix C.2, we have demonstrated that DiffuSearch exhibits scalability as the model and data size increase. Below, we provide a comparison based on FLOPS for reference. Overall, we find DiffuSearch outperforms the Transformer at equivalent FLOPS, and both the Transformer and DiffuSearch show that increasing FLOPS with more data or larger model sizes yields improvements. However, under the same amount of data, Transformers are more prone to saturation, leading to diminishing returns when increasing model size. In contrast, DiffuSearch, due to its more challenging objective, demonstrates more improvements when the model size increases. Additionally, both models show substantial gains when scaling data size.

W3: Limited comparison to state-of-the-art: The paper doesn't compare DiffuSearch to more recent advancements in chess AI [1] that achieved a 2299 Elo rating using transformers. This omission makes it challenging to contextualize DiffuSearch's performance within the current state-of-the-art in chess AI.

Thanks for bringing up this point. On the one hand, [1] annotated 10M games from Stockfish and trained on 128 TPUs, which is unaffordable for us. Due to the resource constraint, we consider a more affordable and controlled comparison between them throughout the experiments, where the Transformer S-A, Transformer S_V, and Transformer SA-V baselines are their models but trained with the same data size with DiffuSearch. We can see that DiffuSearch has a performance advantage compared to these baselines.

On the other hand, similar to W1, our goal is not to build a highly competitive chess engine. Instead, we hope to propose a different paradigm for solving complex problems, beyond the one-step policy [1] and one-step policy with explicit search, which may provide insights into building LLMs with enhanced reasoning and planning capabilities.

We sincerely thank Reviewer XeCm for your detailed review and are grateful for the time you spent on our submission. We are also glad you think our research question is interesting. Below we would like to give detailed responses to each of your comments.

W1: About the paper claim.

Thank you for your insightful suggestion. The position of this work exactly aligns with the weaker claim you mentioned: given the popularity of enhancing one-step policy (e.g., LLMs) inference with explicit search, we want to explore alternative solutions to enhance their planning abilities. We find implicit search via discrete diffusion indeed has the potential to compete with one-step policy with explicit search through controlled experiments. Sometimes we use "explicit search" to directly refer to the “one-step policy with explicit search” for brevity, which may confuse as you mentioned. In the updated version, we have made proper corrections.

W2: I would have liked to see much more detail on the MCTS baseline. I couldn't tell whether it uses a perfect world model or a learned one, and if the world model is learned, I couldn't find information on how it is trained. These details are crucial for understanding exactly what method diffusion is improving on.

Thanks for your suggestion. To combine MCTS and the one-step policy, we follow the approach of AlphaZero, which uses a perfect world model to perform action-state transitions. We have added more detail about the MCTS-enhanced policy baseline in the Appendix B of the latest version.

How does the paper combine the policy model with MCTS to do search for each different "future paradigm"?

Thank you for your question. The term "different future paradigm" specifically refers to DiffuSearch; we did not adjust the future paradigm for the policy model with MCTS. The employed policy model with MCTS is the standard approach following AlphaZero. We have added more detail about the MCTS baseline in the latest version.

Furthermore, the future paradigms in section 3.1 include s-arsar, which seems to be different from s-avsav, and while these names provide an intuition about what's going on, it's still unclear precisely what either of these mean.

Thanks for noting this typo. s-arsar means the same as s-avsav and we will use s-avsav for consistency. We have fixed the typo and added specific examples of each paradigm in the Appendix C.3 to make it clearer.

W3: The text is often rather imprecise, and in some places even a bit misleading.

Thanks for your detailed comments and pointing out some potential confusion. We will explain each point below and refine the corresponding descriptions in the updated version.

1. AlphaGo pre-dates the transformer, but the abstract makes it seem like people were working on LLMs prior to AlphaGo.

2. Heuristic search existed long before the deep learning revolution, but the intro makes it sound like NNs introduced the idea of more efficient search.

Thank you for pointing out this potential confusion. We have made updates to avoid any ambiguity.

3. The intro (line 044) references three papers, all of which discuss model-based RL and the compounding error problem, but presents them as evidence that such errors occur due to recursive invocation of the value function, rather than the world model.

Thank you for your insightful feedback. We would like to clarify that the one-step policy with MCTS is based on a learned value model to guide search at each step and a perfect world model to perform action-state transitions. This paradigm inherently leads to a recursive invocation of the value function, which can result in a compounding error problem.

4. The problem setting defines the value function with respect to a single policy, despite the fact that the players generally have different policies.

Thank you for your feedback. The modeling of DiffuSearch is independent of whether it involves multiple policies. Since we are focusing on chess, we used it as a research example and adopted a simplified definition. However, extending this to a multiple-policy scenario is certainly possible and represents an interesting and valuable direction.

5. In Sec 4.5 (line 288), S-AVAV does not lead to a significant performance improvement, though S-ASS does.

Thank you for your comments. In Line 288 (i.e., Line 293 in the updated version), we focus on comparing different paradigms under the same DiffuSearch modeling framework. We illustrate the significant improvements achieved by incorporating the future state "S" into the model. For instance, we observe enhancements from S-AA (15.07) to S-ASA (41.31), as well as from S-AVAV (17.63) to S-AVSAV (40.69). If I understand correctly, the reviewer may interpreted from another angle, i.e., the comparison between Transformer and DiffuSearch within the same paradigm, which is also correct. We have added clarifications in the text to prevent any ambiguity.

6. I'm not sure why you call a perfect world model with random actions "random world". That still feels confusing to me. Why not just call them "No world model", "Random world+policy", "Random policy", and "Stockfish"---or something similar?

7. I agree that would be a valuable direction. But if what you have isn't actually a scaling "law", I would suggest you don't call it that.

Q2.

The observed decline in the performance of [...].

I'm not talking about the decline in performance, I'm talking about the initial performance. These curves only start at ~35-40% accuracy. This says even at the very first step, the model isn't able to accurately predict the best action or the resulting state.