Implicit Search via Discrete Diffusion: A Study on Chess
Authors: Jiacheng Ye
Zhenyu Wu
Jiahui Gao
Zhiyong Wu
Xin Jiang
Zhenguo Li
Lingpeng Kong
Venue: ICLR 2025 Poster (2025)
Keywords: discrete diffusion model
search
planning
chess
MCTS
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose a model that does implicit search by looking into the future world via discrete diffusion modeling.

Abstract: In the post-AlphaGo era, there has been a renewed interest in search techniques such as Monte Carlo Tree Search (MCTS), particularly in their application to Large Language Models (LLMs).
This renewed attention is driven by the recognition that current next-token prediction models often lack the ability for long-term planning. Is it possible to instill search-like abilities within the models to enhance their planning abilities without relying on explicit search? We propose DiffuSearch , a model that does \textit{implicit search} by looking into the future world via discrete diffusion modeling. We instantiate DiffuSearch on a classical board game, Chess, where explicit search is known to be essential. Through extensive controlled experiments, we show DiffuSearch outperforms both the searchless and explicit search-enhanced policies. Specifically, DiffuSearch outperforms the one-step policy by 19.2\% and the MCTS-enhanced policy by 14\% on action accuracy. Furthermore, DiffuSearch demonstrates a notable 30\% enhancement in puzzle-solving abilities compared to explicit search-based policies, along with a significant 540 Elo increase in game-playing strength assessment. These results indicate that implicit search via discrete diffusion is a viable alternative to explicit search over a one-step policy. All codes are publicly available at \href{https://github.com/HKUNLP/DiffuSearch}{https://github.com/HKUNLP/DiffuSearch}.

Published as a conference paper at ICLR 2025
IMPLICIT SEARCH VIA DISCRETE DIFFUSION: STUDY ON CHESS
A
Jiacheng Ye', Zhenyu Jiahui Gao?, Zhiyong Wu²
Xin Jiang', Zhenguo Li", Lingpeng Kong'
1 The University of Hong Kong 2 Shanghai Jiaotong University
3 Huawei Noah's Ark Lab SShnggaai Al Laboratory
carsenye@connect. www.annneethh............ hku.hk
ABSTRACT
In the post-AlphaGo era, there has been a renewed interest in search techniques such Monte Carlo Tree Search (MCTS), particularly in their application Large Language Models (LLMs). This renewed attention is driven by the recogni- tion that current next-token prediction models often lack the ability for long-term planning. Is it possible to instill search-like abilities within the models to enhance their planning abilities without relying on explicit search? We propose DIFFUSE- ARCH, a model that does implicit search by looking into the future world via discrete diffusion modeling. We instantiate DIFFUSEARCH on a classical board game, Chess, where explicit search is known to be essential. Through extensive controlled experiments, we show DIFFUSBARCH outperforms both the search- less and explicit search-enhanced policies. Specifically, DIFFUSEARCH outper- forms the one-step policy by 19.2% and the MCTS-enhanced policy by 14% on action accuracy. Furthermore, DIFFUSEARCI demonstrates notable 30% en- hancement in puzzle-solving abilities compared to explicit search-basee based policies, along with a significant 540 Elo increase in game-playing strength assessment. These results indicate that implicit search via discrete diffusion is a viable alter- native to explicit search over one-step policy. All codes are publicly available https://github.comm http:/iiiuuccoHKKUUUUU WWI.//KKIUI/////SSAAeeee
INTRODUCTION
Search is central to problem-solving in AI (Russell & Norvig, 2010). One of the most notable examples IBM's Deep Blue (Campbell et al., 2002), which performs extensive search over large space through a strong search algorithm (alpha-beta pruning: Knuth & Moore 1975), defeated the world chess champion Garry Kasparov 1997. Search has also been utilized neural networks. A noteworthy advancement in this progression is exemplified by AlphaGo (Silver et al., 2016) and its successors (Silver et al., 2017b;a), where the policy is guided by an extra value network through Monte Carlo Tree Search (MCTS; Coulom 2006; Browne et al. 2012). By explicitly searching into the future, the decision to be taken can be iteratively refined (Silver et al., 2016; 2017b;a; Schrittwieser al., 2020).
Recent research on Large Language Models (LLMs) demonstrates the utilization of a similar frame- work. Although scale-up, the autoregressive one-step policy addresses only a portion of the prob- lems and relies on explicit search on complex tasks (Hao et al., 2023; Yao et al., 2024; Zhao al., 2024; Trinh et al., 2024, inter alia), highlighting their inherent limitations in long-term plan- ning (Valmeekam et al., 2022; Bubeck et al., 2023; Bachmann & Nagarajan, 2024). This explicit derrndddm approach, however, is not quite satisfactory, as the repeated invocation of the value model can result in an accumulation of errors if the value model is inaccurate and increased inference costs for long-horizon rollouts ( (Yao al., 2024). Given the essence of explicit search ge. MCTS) over one-step policies lies in iteratively looking into the future and leveraging the future to enhance the next token (or action) prediction, our research question is:
Can the policy model predict and utilize the future by itself to improve the next token (or action) prediction without relying on explicit search during inference?

---

Published as a conference paper at ICLR 2025
This paper explores the potential transition from uti- lizing an explicit search algorithm (e.g.. MCTS) over the one-step policy to implicitly searching over fu- ture representations by teaching the policy to predict and utilize the future. Firstly, to reduce the difficulty of future prediction, we take inspiration from diffu- sion models (Sohl-Dickstein et al., 2015; Ho al., 2020), which perform multi-step generative pro- cess for sample generation. Secondly, iterative refine the current policy prediction based on future information, we directly rely on the internal bidi- rectional self-attention mechanism (Vaswani et al., 2017) and the multi-step diffusion generative pro- cess. Finally, we represent the future to be learned and predicted with the multi-step interaction infor- mation between policy and the world (e.g., states and actions), such that the generation of the future shares similar spirits implicit searching in the fu- ture world. We name our approach as DIFFUSE- ARCH, method that looks into the future world via diffusion modeling without any explicit search dur- ing inference. Alternatively, DIFFUSEARCH can be seen as containing world model that predicts the future. However, rather than having a separate world model simulating the environment's transition dynamics and another policy performing action pre- diction through interaction with the world model us- ing planning algorithms such as value iteration (Put-
2. Evaluation
1. Selection
One-step Policy
Po O udd Play ag
3. Backup
(a) Explicit Search via MCTS
Multi-step Diserete Diffission Policy
Play
Imagination
(b) Implicit Search via Discrefe Diffusion
Comparison between explicit search via MCTS and implicit search via dis- crete diffusion. MCTS explicitly performs action selection, state evaluation, and value backup in iterative manner before deter- mining the next action to take (as detailed in Appendix A), while discrete diffusion im- plicitly gathers future information during fu- ture imagination to improve the next action.
Figure I:
erman, 2014) MCTS (Schrittwieser et al., 2020), DIFFUSEARCH internalizes the world model directly within the policy without intermediate components. We show the comparison between ex- plicit search via MCTS and implicit search via discrete diffusion in Figure 1.
We take a specific focus on the chess-playing task, where explicit search is known to be essen- tial (Campbell et al., 2002: Silver et al., 2017a). The ideas and techniques learned in this con- trolled task may eventually be useful in natural-languaal settings as well. We conduct extensive experiments and take a deep look into various paradigms represent and learn the future. When measured by action accuracy, DIFFUSEARCH outperforms the one-step policy (Ruoss et al., 2024) by 19.2%, and MCTSSenhanced policy by 14%. DIFFUSEARCH demonstrates 30% increase in puzzle-solving capabilities in comparison to the MCTS-enhance higher level of game-playing proficiency, as evidenced by a 540 more Elo rating, which showcases the potential of substituting one-step policy with explicit search with learned discrete diffusion model that looks into the future world by itself.
policy. Furthermore, it attains
Our contributions include: 1) propose DIFFUSEARCH to foresee and utilize future information via diffusion modeling as an alternative to explicit search via designed search algorithms ($3); 2) we instantiate DIFFUSFARCH for chess-playing and demonstrate its superior performance compared to both the one-step policy and the MCTS powerrd policy in a rigorous evaluation, such as solving over 30% more puzzles and 540 Elo playing strength in the tournament ($4.2); 3) we provide de- tailed analysis of the design considerations for future representation and diffusion modeling (84.3), as well as unveiling the working mechanism and appealing advantage compared to MCTS-bass policy regarding effectiveness and efficiency ($4.4). These findings demonstrate the possibility of moving from the one-step policy with explicit search algorithms to the future world-awan multi-step diffusion policy with implicit search ability.
2 PRELIMINARIES
This section introduces key concepts and notations in the chess-playine problem and diffusion mod- eling.
2

---

Published as a conference paper at ICLR 2025
Problem Setting Chess, along with other games of perfect information like checkers, othello, backgammon, and Go, fits the framework of alternating Markov games (Littman, 1994). In chess, there exists state space S, action space A, state transition function f(s, a) that determines the subsequent state after taking action a in state 8, and two reward functions (>))) and r (8) representing the two players' reward in state (rewards being zero except at the final time-step). The outcome of the game = =1 is the terminal reward at the end of the game from the perspective of the current player at time-step Chess is also zero-sum game which indicates r0(s) (r) A policy p(a|s) is probability distribution over actions space A. A value function (s) represents the expected outcome when all actions for both players adhere policy p. denoted v"(s)
B|os| 8, di..., p). The goal is to build a policy that, when actions are taken based on it, results the highest possible final outcome.
Discrete Diffusion Modeling Discrete diffusion models (Sohl-Dickstein et al., 2015; Hooge- boom et al., 2021; Austin et al., 2021) are a class of latent variable models characterized by forward and a backward Markov process. Suppose Zo g(xo) is discrete random variable with K possible categories and represented as a one-hot vector. The forward process q(V1:1/20) = I (<1)/<1-1) corrupts the original data zo into sequence of increasingly noisy latent variables #1.7 := ..,..., TT. The learned backward process (o p(æT) 11" pe (x; gradu- ally denoises the latent variables to the data distribution. In order to optimize the generative model (xo) to fit the data distribution ((xg). we typically optimize a variational upper bound on the negative log- likeliidd due to the intractable marginalization:
Lub = By(ee) DKI Deill((er(ao))leerrr
By By(er1 (wo) DKI. (ik,,g((x,1)) (et, ac (20 [20+
B (m)) [log x1)]
Lo
where T.r is constant when fixed prior p(æT) ) is employed. In discrete diffusion, both the forward and backward distribution are defined as categorical distribution, e-g., g( (2+-1) = Cat (at; p = 0)) 0)) (Hoogeboom et al., 2021), where Qr is pre- defined transition matrix of size K x K. Therefore, the forward process posterior g( | 20) and each KL term can be calculated analytically. We provide more details about discrete diffusion in Appendix B.I.
Q, 2t-1) and pa( [x)) q( fz
METHODOLOGY
In this section, we introduce DIFFUSEARCH, an approach that looks into the future world via discrete diffusion modeling without any explicit search at inference time, with focus on the chess-playing task.
3.1 MODELING
In order to endow the model with the capability to predict and utilize the future, we consider train- ing the model in a supervised way following (Ruoss al., 2024), leaving self-play training from scratch (Silver et al., 2017b) for future work. We provide the current state s; as the history repre- sentation following prior studies (Silver et al., 2016; 2017b; Ruoss et al., 2024). For future world representation, we consider variety of alternative variants, such as purely future actions (denoted as s s-aa), action-states (denoted as s-asa), and aetion-state-values (denoted as s-avsav, etc. We analyze the performance of different future paradigms in Section $4.3. The s-asa approach is ultimately chosen as our modeling paradigm considering the effectiveness and simplicity. The policy distribution at state 8; considering the future is given by:
Po (0 8i+1, 0i+1...., 8ith-1, -141
(2)
where h > 1 is the future horizon.
3

---

Published as a conference paper at ICLR 2025
3.2 TRAINING
In order to train a policy that models (2.(2), we consider a supervised training approach lever- aging Stockfish (Romstad et al., 2008). We utilize Stockfish 16, currently the world's strongest search-based engine, as an oracle to label board states extracted from randomly selected games on lichess.org. We approximate the optimal policy with & and obtain each action by taking a = arg maxa, Q (s), For a given world horizon h, we construct a dataset SF D {(s;, (a (a 8441, 1141..... Sith-1. aFF ((4h-1))}, where the oracle future path means playing some move that has the maximum evaluation for the best opponent's 's reply for both players.
An intuitive way use D is train network to directly predict the entire concatenated next action S.F. a-fh-1)- Nonetheless,
and future sequence a. 2: (= == 8i+1 || a=-1 || ... || 8ith-1 we observe that this approach not only fails to predict the future but also impedes the learning of the next action ai (see Section $4.3). Therefore, we resort to diffusion modeling (Sohl- (Sohl-Dickstein al., 2015) as a powerful sequence modeling approach with strong expressive capabilities. The bidirectional mullti-layer elf-aateenioon and iterative denoising mechanism are expected to enhance the prediction of the next action by considering future information. Specifically, we consider dis- crete diffusion modeling and streamline Luh in (1) into a weighted cross-entropy loss moti- vated by Austin et al. (2021); Zheng et al. (2023); Shi et al. (2024); Sahoo et al. (2024). The KL term D xo) (2 for each individual random variable is simplified as -Ac Do log fax;; 0). where and Lib becomes:
Lib -Feg(w00) Σ de (aa (ea) la mi Fy log (a O),
(3)
I-da E (0, 1] is a time-dependent reweighting term that assigns lower weight for noisier Zt, and Cty E [0, belongs to a predefined noise scheduler that controls the level of noise in H, at timestep We explore multiple variants of X₂ in Section $4.3. To enable conditional training with a given state, we freeze the state tokens and perform denoising on the next action a² and all futures tokens 2/S We employ Monte Carlo sampling with regard to æp. H, and when optimizing T.b- We provide detailed derivations in Appendix B.2. We elaborate the training procedure in Algorithm
where d, = 0tt-1-a
3.3 INFERENCE
During inference, taking arg max,, (a.lss) as in one-step policy in DIFFUSEARCH requires marginalizing over all future with horizon h, i.e., (a(a))))= E Po (a; 2 zi(8;), which is intractable search space when h goes larger, e.g., the game tree contains b²
due to the sxponential-growing nodes and the branching factor is around 31 on average in chess (Barnes, 2019). One simplified approach to comparing actions is measure the best future if one action is taken, which can be reflected by the joint probability Po (a,, 2; s, Therefore, we resort to arg max,,.;; Po (a; 2; (s;), which does not involve marginalization and can be achieved by sampling from the trained model. During diffusion sampling, we adopt an easy-first decoding strategy (Savinov et al., 2021; Chang al., 2022), which achieves better performance compared the random decoding approach em- ployed by Austin et al. (2021). Specifically, at diffusion timestep f, the tokens within the least 100 * t-1% predictive log-likelihood are selected to be reset the noise state. To change search depth, we mainly train separate models on D with different h, and study a single model on D with mixed h in Appendix C.I. We elaborate the inference algorithm in Algorithm 2.
4 EXPERIMENTS
4.1
SETUP
Baselines We compare our model with three Transformer models proposed in Ruoss et al. (2024); State-action model (S-A) which learns to predict next move via behavioral cloning; Staue-value model (S-V) which predicts next move via comparing the value of next states; and Action-value model (SA-V) which predicts next move via comparing the value of each legal actions at the current state. We also integrate the trained S-A and S-V models into MCTS following AlphaZero (Silver al., 2017 70)
4

---

Published as a conference paper at ICLR 2025
Algorithm DIFFUSEARCH Training
Algorithm 2 DIFFUSEARCH Inferenee
Input: dataset D {(s. (d,)))), neural network
Input: board state 8, trained network ff (::0, timesteps T.
(::0) d) timesteps 7'
Output: model parameters 0. Denote state length / = |s|: repeat
Output: next action
Denote state length = |s|:
Initialize XT.14 & and X7,++1:N ~ Quoisei for t = T. 1 do
Draw (s. (a, ()) ~ D and obtain X0.1:N =
= [Ik concat);
for = 1 =1+1,-.. 1,..... do
Draw E Uniform(f1,..... TT):
Draw No.0 Cat (f(xx,a))) 0):
Draw X q(xx,, [ for n E {I+1,...,NN;
Draw xt-1,n ~ g(x(-1n n Xt,ns No.
L(O) -x -x En-141 tx. xo,n log (x 0);
end for end for
Minimize L(0) with respect to until converged
Return X0,4+1.
Table 1: Data statistics.
Data We construct dataset for supervised train- ing by downloading games from lichess recorded February 2023. When analyzing the scaling behay- ior, we use up to 100k games, while reverting to the default 10k games for other experiments due to re- source constraints. We show the data statistics in Ta- ble 1. Following Ruoss al. (2024), convert the centipawns returned by Stockfish to the win percent- age and then discretize it into 128 bins represent value in S-V and SA-V. We encode the state as a fixed-length FEN string with 77 characters by padding with s if needed. Actions are stored in UCI notation with 1968 possible moves total. We provide example training data for each paradigm in Appendix C.3.
Stage
Records
Games
Train SA-V (TÖDK)
193,539
100,000
Train SA-V (10k)
17,448,268
10,000
Train others (100k)
6,564.661
100,000
Train others (10k)
659,576
10,000
Action Test
62,561
1,000
Puzzle Test
36,816
10,000
Implementation Details For all the neural models in this paper, we use the same decoder-onke GPT-2 transformer architecture (Vaswani et al., 2017; Radford et al., 2019) for rigorous compar- ison. For DIFFUSEARCH, we convert casual attention into full attention without introducing addi- tional learned parameters. We train all baseline models until convergence and set a maximum of 200 epochs for diffusion models due to their slow convergence. We use the Adam optimizer (Kingma & Ba, 2015), learning rate of 3e-4, and batch size of 1024 for all models. By default, we set the horizon h to be 4, the number of network layers to be 8 (with a total parameter size of 7M), the diffusion timesteps to be 20, and an absorbing noise type. By default, 100 simulations are utilized in MCTS-enhancee determining the level of exploration in MCTS, on a held-out set and set them to Gpnet = 0.1 and = for its superior performance. All experiments are done on 8 NVIDIA V100 32G GPUs.
policy, and its impact is analyzed in Figure 3. We adjust Gpuct and т, constants
Evaluation Metries We mainly consider three metrics to evaluate the policies following (Ruoss al., 2024): Action Accuracy: the percentage of the test set in which the model selects the same action as the ground truth; 2) Puzzle Accuracy: the percentage of puzzles where the policy's ' action sequence exactly matches the known solution action sequence and we use 10k puzzles with difficulty rated by Elo from 399 to 2867 provided by (Ruoss et al., 2024); 3) Tournament Elo: the Elo ratings calculated using BayesElo (Coulom, 2008) in an internal tournament involving all policies, where each pair of policies played 400 games, resulting in total of 6000 games.
4.2 MAIN RESULTS
We report the prediction and playing strength comparison for our model against basclines in Ta- ble Additionally, we report the performance of Stockfish 16 with time limit of 0.05s per legal move, which stands as the oracle used to generate our dataset. We find DIFFUSEARCH signifi- cantly outperforms the S- A model by 653 Elo and 19% action accuracy, indicating the effectiveness of DIFFUSEARCI in improving next action prediction through future prediction. Remarkably, de- spite utilizing 20 times fewer data records than the SA-V model, our model demonstrates superior performance with approximately 10% higher action accuracy. Our model demonstrates superior per- formance over the MCTS- based agent by achieving a higher Elo difference of 542 and an increased action accuracy of 14%. This highlights the effectiveness of DIFFUSEARCH in modeling multi-step
5

---

Published as a conference paper at ICLR 2025
Table Prediction and playing strength comparison for our model against baselines and the oracle Stockfish 16. The S-V and SA-V models can be seen as depth-one search. The best results are bold.
Accuracy
Agent
Search
Tournament Elo
Puzzles
Actions
Stockfish 16 (0.05s) [oracle]
2689
99.10
100.00
10k training games
Transformer (S-A)
1075
3.95
22.10
Transformer (S-V)
1028
12.20
21.45
Transformer (SA-V)
1294
12.74
31.50
Transformer (100 MCTS simulations)
1186
6.85
27.34
DIFFUEEARCH (Ours)
1728
39.49
41.31
100k training games
Transformer (S-A)
1467
20.83
36.58
Transformer (S-V)
1078
17.42
28.89
Transformer (SA-V) V)
1521
24.25
39.76
Transformer (100 MCTS simulations)
1469
20.71
38.05
DIFFUSEARCH (Ours)
1995
58.46
48.66
simulations when compared with the step-by-step MCTS- policy, which relies on a robust value model and necessitates careful balance between the policy and value models.
4.3 ABLATIONS
Future paradigm matters We compare baselines and different future paradigms during the training of DIFFUSEARCH with horizon h = 4 in Table 3. For each future paradigm, we compare training with au- toregressive Transformer and DIFFUSEARCH We find that directly performing future action prediction (S -A) (Chi et al., 2023) with DIFFUSEARCH hurts performance compared to S-A due to the difficulty of future move prediction in chess. However, af- ter we integrate future states, we observe significant performance improvements when comparing S-AA (15.07) to ASA (41.31), and also when comparing S-AVAV (17.63) to S- S-AVSAV (40.69). No further improvement is observed when integrating the values in DIFFUSEARCH, which may be attributed training on the optimal future trajectory rather than all possible trajectories. For training using autoregressive Transformer, we observe that the overall performance hovers around 26%. This per- formance level superior to that achieved by S-A (22.15 (2.1%)). attributed to the utilization of more (S, A) pairs (e.g., each S GSASA record containing h (S, A) pairs). However, the performance falls short when compared to DIFFUSEARCHI which underscores the importance of modeling bidirectional context to leverage future information for subsequent action prediction.
Table 3: Action accuracy comparison of
baselines and different future paradigms.
Paradigms
Transformer
DIFFUSEARCH
S-A
22.10
S-V
21.45
SA-V
31.50
S-AA
26.62
15.07
S-ASA
27.39
41.31
S-ASS
24.93
41.19
S-AVAV
25.92
17.63
S-AVSAV
25.59
40.69
Ensuring the validity of future world dynamics crucial After we discuss the future paradigm, we now investigate the effect of future quality on performance, as shown Table For better illustration, denote a sequence of future horizon as [s] fs0, on), 41 = g(s1), 82 = f(1 a1), g(22)). where f is world dynamic function and is a pol- icy function. 5g is the current state and ag gested by Stockfish. We first utilize random state-action se- quences for future steps, where both actions and states were randomly selected (i.e., random world f and random policy g). This methodology did not yield performance enhance-
Table Future world quality in su- pervising the model for the S-ASA paradigm.
Future Quality
Acc.
Without future
22.10
move sug-
+ Random world+policy
22.69
+ Random policy
39.47
+ Stockfish policy
41.31
ments. Subsequently, explore selecting random actions and incorporating the resulting state

---

Published as a conference paper at ICLR 2025
100
44
Valid 9
40
Random
Percentage
75
Valid
iu y*
Accuracy
Accuracy
30
50
40
25
20
Transformrr (S-A)
38
DuruStancc (5 ASA)
36
0
2
3
2
16
1
2
4
16
32
64
128
Future Horizon
Self-attention Layers
Diffusion Timesteps
Figure 2: (Left) Prediction quality analysis for DIFFUŠEARCH at different future steps. (Middle) Action accuracy when scaling self-attention layers. (Right) Action accuracy when increasing diffu- sion timesteps.
from executing those actions (i.e., random policy but an oracle world f), which notably outper- forms the initial strategy. This underscores the significance of aligning states with corresponding actions, mirroring the dynamics of the world. Finally, we investigate incorporating high-qualit future actions suggested by Stockfish (i.e., Stockfish and oracle world f) and observe additional performance improvements compared the random action selection approach.
Proper discrete diffusion modeling helps Given the dataset annotated with future states and actions, we in- vestigate alternative ways to train the model, presented in Table 5, We first observe is hard to teach the model to directly output the entire future sequence, leading to lower performance compared to auto-regressive train- ing. Secondly, we employ continuous Gaussian diffu- sion VDM (Kingma et al., 2021) and observe its superior performance compared the Direct and auto-regressive methods, but inferior compared to discrete approaches. The absorbing diffusion with reciprocal X 1/t ob- tained by setting a; - & in (3) is a simplified expression from D3PM (Austin al., 2021), which we find significantly outperforms continuous diffusion. Fi- nally, we discover linear A (Bond-Th Taylor al., 2022; Zheng et al., 2023) further exceeds the constant and reciprocal ones, as well as the multinomial counterpart.
Table 5: Comparison of training meth- ods. Direct: train the model to predict the entire future sequence at once.
Method
Acc.
Direct
20.61
Auto-regressive Gaussian
27.39
31.91
Absorbing, A
39.66
Absorbing. Ae 1/t
39.07
Absorbing. d =
41.31
Multinomial, X T
40.08
4.4 ANALYSIS
Does DIFFUSEARCH predict accurate future information? We analyze the percentage of valid actions and the optimal action recommended by Stockfish for each predicted action. The best dg metric is exactly the action accuracy by definition.. Additionally, we assess whether each predicted state a valid representation and if 8; corresponds to the resulting state when action dj-1 taken at 8:-1- The initial state 8o provided as input is excluded, and the results are presented in the left figure of Figure 2. We observe that the first action, denoted as ap, are almost 100% valid. As we progress through future steps, both the valid rate and the optimal rate decline. However, even at i = 3, where the valid rate stands at 50%, surpasses the random move baseline of approximately 1.6% (calculated as the average number of legal actions per move, 31, divided by the total number of moves, 1968). This indicates that the model retains a certain level of predictive capability for future moves, albeit with reduced performance. A similar pattern appears in the evaluation of states, where the accuracy is perfect for the first predicted state 81 but diminishes in subsequent predictions. In Appendix Table 8, demonstrate that further increasing the training data enhances the effective- ness of the world model within DIFFUSEARCH, achieving over 90% accuracy in predicting valid and matched future states corresponding the preceding action.
How does DIFFUSEARCH leverage future information? We attributes the future- ability of DIFFUSEARCH mainly to self-attention and iterative decoding process, as shown in the middle and right figures of Figure 2, respectively. When employing single self-attention layer, our model
7

---

Published as a conference paper at ICLR 2025
50
40
Latency (ms per anooe
Transfoment MCTS DIFFUSEARCH
10*
400 sim.
Accuracy
35
sim.
Accuracy
40
100 sim.
30
50 sim.
200 sim.
400 sim.
10a
25 sim.
30
25
SEARIES
Transforee-- MCTS
Transformer (S-A)
25 sim.
10-
20
Dar OSaAnnn (5 ASA)
10 15
15
20
5
10
20
40
80
Average Search Depth
Average Search Depth
Data Size (k)
Figure 3: (Left) Action accuracy when increasing average search depth in MCTS through more simulations and DIFFUSBARCH through context length extension. (Middle) Latency measured by ms per second when increasing search depth. (Right) Action accuracy when scaling data size.
exhibits inferior performance compared to the S-A model, yet surpasses with two layers. More- over, its performance steadily enhances as augment the number of layers. This suggests that with additional layers, there is more chance for the subsequent actions and future to interact recipro- cally, akin the enhancement in the action prediction with increased MCTS simulations. We do not observe a similar upward trend in the performance of S-A model when increasing attention layers as in (Ruoss al., 2024), possibly indicating that the available data (10k) does not necessitate the integration of more layers. In the right figure of Figure 2, it is evident that employing an appropriate decoding strategy (such as likelihood (eelihood-bbsed) further enhances next-action prediction as the number of iterations grows. However, the overall improvement is relatively modest compared to increasing the attention layers.
Explicit search vs. Implicit search Based on our previous analysis, we can consider DIFFUSE- ARCH as performing implicit search through the inner self-attention layers and the multi-step dif- fusion process. Now, we aim to evaluate the efficiency and effectiveness of this implicit search in comparison to explicit search using MCTS when conducting deeper searches. In DIFFUSEARCH, deeper search is realized by increasing the context length (80 tokens per search depth), whereas in MCTS, it is achieved through running more simulations. In the left figure of Figure 3, it is evident that DIFFUSEARCH exhibits significant enhancement when increasing search depth, while MCTS becomes stagnant after 50 simulations at a search depth of around 4. This could be attributed to the accumulated errors caused by the value network due to a limited horizon. In the middle figure of Figure 3, measure the latency per move for Transformer with MCTS and DIFFUSEARCH on a single V100 GPU with batch size 1. The performance of Transformer combined with MCTS is notably affected by the necessity of invoking the value network for every simulation. In contrast, DIFFUSEARCH experiences only slight rise in latency as it requires just one call for greater depth.
Scaling In Figure 2, the effectiveness of model scaling in DIFFUSEARCH has been observed. Here explore the impact of increasing the dataset size on the performance. Specifically, we conduct experiments training the DIFFUSEARCI S-ASA model with a horizon of 4 and the Transformer S-A using game sizes ranging from 5k 100k, as shown in the right figure of Figure 3. Both the Transformer and DIFFUSEARCH models exhibit a log-2 scaling behavior, showing that doubling the training data results in a linear increase in accuracy. Scaling also enhances future prediction significantly, leading to a more valid and accurate representation of future actions and states, as well as near- perfectt level of capturing the state-action transition dynamics, as detailed Appendix C.2.
Case study We sample several challenging puzzles from Lichess (with Elo ratings above 1800) to compare the predictions of DIFFUSEARCH and Transformer (S-A). Two instances are shown in Figure 4, with additional cases provided in Appendix C.4. DIFFUSEAR ARCCH demonstrates superior foresight, accurately predicting critical exchanges and piece sacrifices that lead to long-term strategic advantages. In the left puzzle, DIFFUSEARCI strategically sacrifices the rook to set up a long-term checkmate situation against the opponent. This maneuver compels the opponent to defend and creates an opportunity to capture the queen, facilitating valuable piece exchanges. The S-A model, unfortunately, makes critical error by focusing achieving direct checkmate without considering the possibility of the opponent's 's queen launching a counterattack. Similarly, in the right puzzle, DIFFUSFARE anticipates an exchange sacrifice, correctly valuing the long-term positional benefits
8

---

Published as a conference paper at ICLR 2025
Transformer (S-A)
DiffuSearch
Transformer (S-A)
DiffuSearch
SAM
A
Figure 4: Two examples of Transformer (S-A) A) and DIFFUSEARCI solving challenging puzzles. The predicted next move is in blue for both policies. The predicted future actions from DIFFUSEARCI are in light blue and red representing the two players, respectively, along with the numerical counters 1, 2, and 3 indicating future steps.
of opening lines by sacrificing the rook for its queen. Conversely, the S-A model misjudges the value of this exchange, leading to suboptimal moves. These findings highlight the effectiveness of DIFFUSEARCH in long-term planning without relying explicit search.
RELATED WORK
5.1 NEURAL NETWORKS FOR CHESS
The development of chess Al has undergone a significant transformation, shifting from the explicit design of search strategies and heuristics to the more data-driven and learning- based approaches. The early research, exemplified by Turing's investigations (Burt, 1955) and NeuroChess (Thrun, 1994), heavily depended handcrafted search algorithms and heuristics, eventually leading the development of powerful search engines like Deep Blue (Campbell et al., 2002) and Stockfish (Rom- stad et al., 2008). However, the emergence neural network-based approaches, typically Alp- haZero (Silver et al., 2017a), marked a paradigm shift, where deep reinforcement learning equipped with Monte Carlo Tree Search (MCTS) enabled the system learn its own heuristics, i.e., the policy and value networks, without the need for manual design (Klein, 2022; McGrath et al., 2021). The rise large language models (LLMs) has also inspired innovations in chess AI, such as the eval- uation (Toshniwal et al., 2022; Carlini, 2023) and interpretation (Li et al., 2023a; Karvonen, 2024) of LLMs' ability to play chess, the integration chess-related text data into training (Feng et al., 2024), and the exploring of searchless models by scaling the policy networks (Ruoss et al., 2024). Despite this, lookahead search methods like beam search (Feng et al., 2024) and even depth-one search with the value network (Ruoss et al., 2024) remain superior the policy models as action predictors, which is the same as in the AlphaZero era (Silver et al., 20172 Team., 2018). This under- scores the continued significance of lookahead information for move prediction in chess. In contrast to prior research, we explore directly teaching the policy model to look ahead, thereby eliminating the requirement of handcrafted search algorithms or separate value networks.
5.2 DIFFUSION MODELS
Diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Austin et al., 2021). 1), powerful class of generative models, have been applied to various fields such image generation (Dhariwal & Nichol, 2021; Rombach et al., 2022: Croitoru et al., 2023, inter alia). text generation (Li et al., 2022; Gong al., 2022; Zheng et al., 2023; Lou et al., 2023; Li et al., 2023b, inter alia) and reinforcement learning (Janner et al., 2022: Ajay et al., 2022; Chi et al., 2023; Zhu et al., 2023, inter alia). Theo- retically, diffusion models perform multi-step denoising process to progressively convert a random noise into data sample, and the denoising procedure can be seen as parameterizing the gradients of the data distribution (Song & Ermon, 2019), connecting them to score matching (Hyvärinen & Dayan, 2005) and energy- based models (LeCun al., 2006). Particularly, diffusion models have demonstrated their effectiveness in tasks that require global control and future planning, such as math reasoning (Ye et al., 2024b; Gong al., 2024), paragraph generation (Zhang al., 2023b),

---

Published as a conference paper at ICLR 2025
trajectory planning (Janner et al., 2022) and robot manipulation (Chi et al., 2023). The superior planning ability of diffusion compared to autoregression has been studied in Ye et al. (2024a). Compared to Diffusion Policy (Chi et al., 2023), DIFFUSEARCH internalizes world model inside the policy, which we find is crucial in Section $4.3. Furthermore, we focus on exploring diffusion models for implicit search complex tasks that require search.
alternative to the one-step policy with explicit search to deal with
5.3 WORLD MODELS
The primary goal of a world model is to capture the underlying dynamics of the environment and predict the future outcome of certain actions in the context of model-based reinforcement learning (MBRL) (Wang et al., 2019; Moerland et al., 2023). The learned world model can be used for policy optimization of RL agent (Sutton, 1991; Feinberg et al., 2018; Hafner al., 2020a) and allow the agent to explicitly reason about the future consequences of its actions (Hafner et al., 2019; Schrittwieser al., 2020; Ye et al., 2021). Most of the conventional world models (Hafner al., 2020a;b; 2023) rely on single-step prediction, which suffer from compounding errors (Asadi et al., 2019; Xiao et al., 2019; Lambert al., 2022). Recently, there has been growing interest in building multi-step world models utilizing diffusion models (Zhang et al., 2023a; Rigter et al., 2023; Jackson at. 2024; Ding et al., 2024), which, however, separate the world model and policy. Similar to ours, Diffuser (Janner et al., 2022) and Decision Diffuser (DD; Ajay et al. 2022) also unify the world model with the policy. However, the modeling details, training paradigm, and action prediction differ. Specifically, both of them employ continuous diffusion while we use discrete diffusion. In addition, Diffuser trains an unconditioned model and requires a guidance function to obtain desired actions, while we model the best action and future trajectory condition state-only future trajectories and predicts the action through an inverse dynamics model while we model both future states and actions. Finally, the comparison of diffusion world model and explicit search has not been rigorously explored in domains that require precise and sophisticated lookahead such chess, the best of our knowledge.
given state. DD models
6 CONCLUSION AND DISCUSSION
In this study, we present evidence showcasing the potential transition from employing explicit search on one-step policy to implicit search within future-aware policy on the classic board game Chess. The proposed model, DIFFUSEARCH, demonstrates not only superior performance compared to the searchless policy but also the policy empowered by explicit search. We provide extensive experi- ments to demonstrate and analyze DIFFUSEARCI, More broadly, the ideas and techniques discussed in this controlled task may eventually be valuable in natural language settings improve the current next-token prediction LLMs as well.
We now discuss some limitations and workarounds in our study. Firstly, one usage of explicit search such as MCTS is to enhance policy performance through self-play training, such that is able to achieve amazing performance without any human supervision (Silver et al., 2017b). However, our model currently relies on an oracle (Stockfish) to provide future supervision. The integration of DIFFUSEARCH with self-play is an interesting direction to explore. Secondly, our model achieves a deeper search by increasing the context length, with the current training limited to depth of 7, corresponding to a context length of 648. For scenarios requiring more tokens to represent a state or deeper searches, integrating techniques for long-context models may be useful for efficient training or inference (Dao et al., 2022; Gu & Dao, 2023; Xiong et al., 2024; An et al., 2024). Finally, our model's performance is currently constrained by the relatively small training dataset of up to 100k games due to resource restrictions, considerably less than the 10 million games used in the study by Ruoss al. (2024). Continuing to scale the model and data remains valuable direction.
ACKNOWLEDGMENTS
This research was supported in part by the joint research scheme of the National Natural Sci- ence Foundation of China (NSFC) and the Research Grants Council (RGC) under grant number N HKU714/21.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision makinn? In NeurIPS 2022 Founda- fion Models for Decision Making Workshop, 2022.
Chenxin An, Fei Huang, Jun Zhang, Shansan Gong, Xipeng Qiu, Chang Zhou, and Lingpeng Kong- Training- free long-context scaling of large language models. In Forty-first International Confer- ence on Machine Learning, 2024.
Kavosh Asadi, Dipendra Misra, Seungchan Kim, and Michel L Littman. Combating the compounding-error problem with a multi-step model. arXiv preprint arXiv: (tt5:/9905/33220 13320, 2019.
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. In Mare' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural In- formation Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp- 17981-17999, 2021.
Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. arXiv preprint arXiv:2403.06963, 2024.
David Barnes, What is the average number of legal moves per turn? https://chess. stackexchange com/ a/24225, 2019.
Sam Bond-Taylan Peter Hessey, Hiroshi Sasaki, Toby P Breckon, and Chris G Willcocks. Un- leashing transformers: Parallel token prediction with discrete absorbing diffusion for fast high- resolution image generation from vector-quantized codes. European Conference on Computer Vision, pp. 170-188. Springer, 2022.
Cameron B Browne, Edward Powley, Daniel Whitehouse, Simon M Lucas, Peter I Cowling, Philipp Rohlfshagen, Stephen Tavener, Diego Perez, Spyridon Samothrakis, and Simon Colton. A survey of monte carlo tree search methods. IEEE Transactions Computational Intelligence and Al in games, 4(1): 2012.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, al. Sparks of artificial general intelligence: Early experiments with gpt-4. arXiv preprint arXiv:2303.12712 2023.
Cyril Burt. Faster than thought: A symposium digital computing machines. edited by b. bowden. British Journal of Statistical Psychology, 1955.
Murray Campbell, A. Joseph Hoane Jr., and Heinn Hsu. Deep blue. Artif. Intell., 2002.
Nicholas Carlini. Playing chess with large language models. https:/nicholas.ccm
com/ww wr www.iiin///22i wttschhesssseee 11m. html, 2023.
Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman. Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 11315-11325 2022.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shu- ran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
Rémi Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In International conference on computers and games, pp. 72-83, Springer, 2006.
Rémi Coulom. Whole-history rating: A bayesian rating system for players of time-varying strength. In Computers and Games, 2008.
Florinel-Alin Croitoru, Vlad Hondru, Radu Tudor lonescu, and Mubarak Shah. Diffusion models in vision: A survey. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(9); 10850-10869, 2023.
11

---

Published as a conference paper at ICLR 2025
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances in Neural Information Processing Systems, 15.16344-16339, 2022.
Jacob Devlin, Ming-Wor Mii Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Pupers), pp. 4171-41 Minneapolis, Minnesota, June 2019. Association for Com- putational Linguistics. doi: 10. 10.18653/v1/N19-11 1423, URL https: http///caaannttooooeo .org///goor//ooooooooo.oo N19-1423.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-877 18944, 2021.
Zihan Ding, Amy Zhang, Yuandong Tian, and Qinqing Zheng. Diffusion world model. arXiv preprint arXiv:2402.03570, 2024.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael Jordan, Joseph Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free reinforcement learning. arXiv preprint arXiv: 000000 2018.
Xidong Feng, Yicheng Luo, Ziyan Wang. Hongrui Tang, Mengyue Yang, Kun Shao, David Mguni, Yali Du, and Jun Wang. Chessgpt: Bridging policy learning and language modeling. Advances in Neural Information Processing Systems, 36, 2024.
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. In The Eleventh International Conference on Learning Representations, 2022.
Shansan Gong, Shivam Agarwal, Yizhe Zhang, Jiacheng Ye, Lin Zheng, Mukai Li, Chenxin An, Peilin Zhao, Wei Bi, Jiawei Han, al. Scaling diffusion language models via adaptation from autoregressive models. arXiv preprint arXiv:2410. 17891, 2024.
Albert Gu and Tri Dao. Mamba: Linear-tima sequence modeling with selective state spaces. arXiv preprint arXiv:2312.0077 2023.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 2555-2565. PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learn- ing behaviors by latent imagination. In International Conference Learning Representations, 2020a.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. arXiv preprint arXiv:2010.02 093, 2020b.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.0411 2023.
Shibo Hao, Yi Gu, Haodi Ma, Joshua Hong, Zhen Wang, Daisy Wang, and Zhiting Hu. Reasoning with language model is planning with world model. Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, pp. 1154-8173 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Mare Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin (eds.). Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020. NeurIPS 2020, December 6-12, 2020, virtual, 2020.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin. Perey Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Process- ing Systems 2021, NeurlPS 2021, December 6-14, 2021, I, virtual, pp. 12465, 2021.
12

---

Published as a conference paper at ICLR 2025
Aapo Hyvärinen and Peter Dayan. Estimation of non-normalized statistical models by score match- ing. Journal of Machine Learning Research, 6(4). 2005.
Matthew Thomas Jackson, Michael Tryfan Matthews, Cong Lu, Benjamin Ellis, Shimon Whiteson, and Jakob Foerster. Policy-guided diffusion. arXiv preprint arXiv:2404.06335 2024.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.0999)) 2022.
Adam Karvonen. Emergent world models and latent variable estimation in chess-playing language models. arXiv preprint arXiv:2403. 15498, 2024.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Ad- vances in neural information processing systems, 34:21696-21707, 2021.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR (Poster), 2015.
Dominik Klein. Neural networks for chess. arXiv:2209.01506, 2022.
Donald E Knuth and Ronald W Moore. An analysis of alpha-beta pruning. Artificial intelligence, (4):293-326, 1975.
Nathan Lambert, Kristofer Pister, and Roberto Calandra. Investigating compounding prediction errors in learned dynamics models. arXiv preprint arXiv:2203.096:3 2022.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
Kenneth Li, Aspen K. Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. Emergent world representations: Exploring sequence model trained on synthetic task. In ICLR, 2023a.
Xiang Lisa Li, John Thickstun, Ishaan Gulrajani, Percy Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. In Conference on Neural Information Pro- cessing Systems, NeurlPS, 2022.
Yifan Li, Kun Zhou, Wayne Xin Zhao, and Ji-Rong Wen. Diffusion models for non-autoregressive text generation: A survey, arXiv preprint arXiv:2303.06574, 2023b.
Michael L Littman. Markov games as framework for multi-agent reinforcement learning. In Machine learning proceedings 1994, pp. 157-163. Elsevier, 1994.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. ArXiv preprint, abs/23 10.1144..6,,,, 2023.
Thomas McGrath, Andrei Kapishnikov, Nenad Tomasev, Adam Pearce, Demis Hassabis, Been Kim, Ulrich Paquet, and Vladimir Kramnik. Acquisition of chess knowledge in alphazero. arXiv:2111.09259, 2021.
Thomas M Moerland, Joost Broekens, Aske Plaat, Catholijn M Jonker, et al. Model-based rein- forcement learning: survey. Foundations and Trends@ in Machine Learning, 16(1):1-118, 2023.
Martin Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Mare Rigter, Jun Yamada, and Ingmar Posner. World models via policy-guided trajectory diffusion. arXiv preprint arXiv:23/2.088..0 028333,, 2023.
13

---

Published as a conference paper at ICLR 2025
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- on computer vision and pattern recognition, pp. 10684-10695, 2022.
Tord Romstad, Marco Costalba, Joona Kiiski, Gary Linscott, Yu Nasu, Motohiro Isozaki, Hisayori Noda, and et al. Stockfish, 2008. URL https: htt///toccfssschcccc
org.
Anian Ruoss, Grégoire Delétang, Sourabh Medapati, Jordi Grau-Moya, Li Kevin Wenliang, Elliot Catt, John Reid, and Tim Genewein. Grandmaster-level chess without search. arXiv preprint arXiv:2402.0444 04494, 2024.
Stuart J. Russell and Peter Norvig. Artificial intelligence modern approach, third international edition. 2010. https//:: www.nticscoooiii.t......icc..mee wtw.ttppee//a.....mmiiiccco.ooa
Subham Sekhar Sahoo, Marianne Arriola, Yair Schiff, Aaron Gokaslan, Edgar Marroquin, Justin Chiu, Alexander Rush, and Volodymyr Kuleshov. Simple and effective masked diffusion language models. arXiv preprint arXiv:2406.07524 2024.
Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step- unrolled denoising autoencoders for text generation. In International Conference on Learning Representations, 2021.
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-609 2020.
Jiaxin Shi, Kehang Han, Zhe Wang, Arnaud Doucet, and Michalis K Titsias. Simplified and gener- alized masked diffusion for discrete data. arXiv preprint arXiv:2406.04322 2024.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwieser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
David Silver, Thomas Hubert, Julian Schrittwieser, loannis Antonoglou, Matthew Lai, Arthur Guez, Mare Lanctot, Laurent Sifre, Dharshan Kumaran, Thore Graepel, Timothy P. Lillicrap, Karen Simonyan, and Demis Hassabis. Mastering chess and shogi by self-play with a general reinforce- ment learning algorithm. arXivv1712.01815 2017a.
David Silver, Julian Schrittwieser, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lueas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-359, 2017b.
Jascha Sohl-Dickstein. Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-2263 PMLR, 2015.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' d'Alché-Buu -Bch Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 11895-11907, 2019.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: En- hanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting, ACM Sigart Bulletin, 264):160-163, 1991.
Leela Chess Zero Team. Leela chess zero, 2018. URL https: htttps://lczero.0rg
Sebastian Thrun. Learning to play the game of chess. In NIPS, 1994.
Shubham Toshniwal, Sam Wiseman, Karen Livescu, and Kevin Gimpel. Chess as a testbed for language model state tracking. In AAAI, 2022.
14

---

Published as a conference paper at ICLR 2025
Trieu H. Trinh, Yuhuai Wu, Quoe V Le, He He, and Thang Luong. Solving olympiad geometry without human demonstrations. Nature, 625(7995):476445 482, 2024.
Karthik Valmeekam, Alberto Olmo, Sarath Sreedharan, and Subbarao Kambhampati. Large lan- guage models still can't 't plan (a benchmark for Ilms on planning and reasoning about change). In NeurlPS 2022 Foundation Models for Decision Making Workshop, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp- 5998-6008 2017.
Tingwu Wang, Xuchan Bao, Ignasi Clavera, Jerrick Hoang, Yeming Wen, Eric Langlois, Shunshi Zhang, Guodong Zhang, Pieter Abbeel, and Jimmy Ba. Benchmarking model-based reinforce- ment learning. arXiv preprint
2019.
Chenjun Xiao, Yifan Wu, Chen Ma, Dale Schuurmans, and Martin Müller. Learning to combat compounding-ervon in model-based reinforcement learning. arXiv preprint arXiiv11122 1226, 2019.
Wenhan Xiong, Jingyu Liu, Igor Molybog, Hejia Zhang, Prajjwal Bhargava, Rui Hou, Louis Martin, Rashi Rungta, Karthik Abinav Sankararaman, Barlas Oguz, et al. Effective long-context scaling of foundation models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 1: Long Papers), pp. 4643-4663 2024.
Shunyu Yao, Dian Yu, Jeffrey Zhao Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. Ad- vances Neural Information Processing Systems, 36, 2024.
Jiacheng Ye, Jiahui Gao, Shansan Gong, Lin Zheng. Xin Jiang, Zhenguo Li, and Lingpeng Kong. Beyond autoregression: Discrete diffusion for complex reasoning and planning. arXiv preprint arXiv:2410.14157, 2024a.
Jiacheng Ye, Shansan Gong, Liheng Chen, Lin Zheng, Jiahui Gao, Han Shi, Chuan Wu, Zhenguo Li, Wei Bi, and Lingpeng Kong. Diffusion of thoughts: Chain-of-thought reasoning in diffusion language models. arXiv preprint arXiv:2402.07754 2024b.
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Picter Abbeel, and Yang Gao. Mastering atari games with limited data, Advances in Neural Information Processing Systems, 34:25476-25488, 2021.
Lunjun Zhang, Yuwen Xiong, 7e Yang, Sergio Casas, Rui Hu, and Raquel Urtasun. Learn- ing unsupervised world models for autonomous driving via discrete diffusion. arXiv preprint arXiv:2311.00017 2023a.
Yizhe Zhang, Jiatao Gu, Zhuofeng Wu, Shuangfei Zhai, Joshua M. Susskind, and Navdeep Jaitly. PLANNER: Generating diversified paragraph via latent language diffusion model. In Thirty- seventh Conference on Neural Information Processing Systems, 2023b.
Zirui Zhao, Wee Sun Lee, and David Hsu. Large language models commonsense knowledge for large-scale task planning. Advances in Neural Information Processing Systems, 36, 2024.
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. ArXiv preprint, abs/2302.05737, 2023.
Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang, Yong Yu, and Weinan Zhang. Diffusion models for reinforcement learning: A survey. arXiv preprint 23X0022311.01223, 2023.
15

---

Published as a conference paper at ICLR 2025
A DETAILS ABOUT MCTS-ENHANCED POLICY
This baseline is fully aligned with the approach used in AlphaZero (Silver et al., 2017a). The one- step policy directly predicts the next action, while the MCTS-enhances Policy constructs search tree that simulates the future to enhance the evaluation of potential next actions. Each node in the search tree contains edges (s, a) for all legal actions a E A(s). Each edge stores a set of statistics,
(N(s, a). ((,, Q(s, Q(..), P(s,a))
(4)
where (s, a) is the visit count, (ss, a) is the total action-value, Q(s, a) is the mean action-value, and P(s, a) the prior probability of selecting that edge. The algorithm proceeds by iterating over the former three phases below and then selects move play:
Selection. The algorithm begins at the root node and traverses the tree, selecting child nodes based on strategies to maximize the exploration termediate node, action is selected according to the statisties in the search tree, a argmax (Q(se, a) + (s. a)), using a variant of the PUCT algorithm,
promising paths, Specifically, at each in-
VE, (s,
U(s,a) - Cpoa a)
(5)
N(s,0) a)
where Cpuct is constant determining the level of exploration; this search control strategy initially prefers actions with high prior probability and low visit count, but asymptotically prefers actions with high action-value.
Expansion and evaluation. Upon reaching lcaf node, if it does not represent a terminal state (i.e., the end of the game), one or more new child nodes are expanded and evaluated by the policy and value model. The leaf node s is added to queue for neural network evaluation, and p Pa (.). The leaf node is expanded and each edge (s1, a) is initialized to (N(s,. 0, W (SL, 0) 0, Q(st, a) 0, 0, P(st, a) }; the value v then backed up.
(8()))
Backup. The edge statistics are updated in backward pass through each step t S L. The visit counts are incremented, N(s,, a.) = N(,,, a,) + 1, and the action-value is updated to the mean value, W(s,, ax) = (s),,a;) + v, Q(s,
W(snar)
N(a...)
Play. After iteratively cycling through the above phases, a move is selected to play in the root position 8o the end of the search based the statistical information, e.g., proportional to its exponentiated visit count, <(assu) rameter that controls the level of exploration. The search tree is reused at subsequent time-steps: the child node corresponding to the played action becomes the new root node; the subtree below this child is retained along with all its statistics, while the remainder of the tree is discarded.
N(su, // EE, N N(5g, b)2/T, where is a temperature pa-
B DERIVATIONS
B. DISCRETE DIFFUSION
In this section, we provide a detailed derivation of the representation for distributions used in the objective (1), which we bring here for a better illustration:
T
Lub = Eq(mm)) DKL (kulo(arl(o)ll((rr))
+
E ga (mn)) [KL (g(2, 2e xo
Ipe (2, x [x)]]]
lsT
1=2
Ext (ma) [log po c (z) i)]
16

---

Published as a conference paper at ICLR 2025
where Ly is constant when fixed prior p(ær) is employed. In discrete diffusion, both the forward and backward distribution are defined as categorical distribution, e.g., ()) (z1-1) = Cat (zy; = Q, xt-1) and po(x,,1 | - 1 (z,; 0)) (Hoogeboom et al., 2021). where Qe is pre-
defined K x K transition matrix and K is the size of categories.
The posterior g( (o) Starting from æo. we obtain the following 1-step marginal and pos-
terior at time - 1:
(x11(20) Cat x1; p- Q₂ 20). with Q, Q1Q2.... Q2
q( 2 2o ( ( (00) = Cat It-1:P = Q. 2 Q Q2-100
q(z,-11 x²)
(6)
q( (o)
xt Q2 200
where (10 = g 1+)-) due to the Markov property of the forward process. The KL divergence between q and po be computed by simply summing over all possible values of each random variable. The cumulative products Q which can be computed in closed form or precomputed for i depending on the choice Q, may be prohibitive for large T and number of categories. Therefore, two commonly used forms of Q are introduced by Hoogeboom et al. (2021) and Austin et al. (2021), ) which ensures Q2 can still be computed efficiently, allowing the framework to scale to a larger number of categories.
Multinominal diffusion The transition matrix initially proposed for the binary scenario by Sohl- Dickstein et al. (2015) and later expanded to categorical by Hoogeboom et al. (2021) can be repre- sented as K x K matrix:
1
K-1 Br K
if i=j i
3c
This transition matrix can also be written as - (1-B)) 8,11 / where is a column vector of all ones. The transition matrices Q can be computed closed form. Denote the vector represents the uniform noise distribution Jnoise 1/K. In each step, transition another token with probability 3r and stay the same with probability - 8c. After / steps, the only operative quantity is the probability of not yet having transitioned to another token, given by = 1I: BB))
Therefore, derive:
Q, = o
where setting Aooise 1/K gives the Q, for multinominal diffusion.
Absorbing diffusion For diffusion models with an absorbing state m, the following matrix is introduced by Austin et al. (2021):
if i i=j= = m
1 B
if i i= j +
8 if j m. + m
The transition matrix can also be written as (1 - 3,)1 + 3,1eT where em is a vector with a one on the absorbing state and zeros elsewhere. Since is absorbing state, the corruption process converges not to a uniform distribution but to the point-mass distribution on m. For text generation, m is the [MASK] token and this leads BERT-like training objective (Devlin al., 2019), while masks tokens according to some schedule and learns to denoise them iteratively. Similar as in multinomial diffusion, we set Quoise = em for absorbing diffusion, where en is a one-hot vector on the [MASK] token, and obtain Q, based on (7)
17

---

Published as a conference paper at ICLR 2025
B.2 A SIMPLIFIED OBJECTIVE
The categorical distribution parameterized by for each variable that follows g(00,- [2,, 20) based on (6) given as:
p QiTa @ Q
2, Q, 2o
[1 ae 010 1] @ (0 + (1 a1 1 Inoise
a 2o + (1 - (1) 0 Qnoise
(1 11) +(1 3,) 1111--4--1 -4-1) Gnoise B,0y 1x00 (1011---111 -----)10000 (1111-11)) 2 Quoise
0,2, 2D + (1 - x quoise
(1 10 + (1 - B.) 10]1-04-111
04-00-111 )o a 31 + 1:0 Do + 1101--441111 1 Unoise
a Do + (1 0z.
where Ow: = Quoise wx) represents the probability of noise drawn from Qnoise being equal to Ct. Note D2 @ 2p 0 if at * wo otherwise 1. Thus the computation of that parameterize ( (o) breaks down into cases:
T2 + (1 T7t)) Quoise:
if az - 200
p
ArDD + (1 Xe) Inoise (2tz,
if at + Z0;
where 7)2 == 1 8e(1 -0-- Jeuue (u==))) de = ae-1-at and Imoise (xxc) = (1 (1-Be) (Be 3e denotes
+c+(1-au (euise )
I-ax
a noise distribution that interpolates between and Inoise.
Since we set po (2 = q( q xt fax;; 0)), the KL divergence between g( 2 xo) and po( becomes O when
In the case of absorbing diffusion, at Quoise
if
24 * 2y and Guoise (xte) qmoise- p has probability ) on index zp and 1 X on the absorbing state. The model f 0) has cero-probability the absorbing state as it never predicts the mask token. Therefore, Po xt,-1 (x))) also has d probability on the absorbing state. Putting them together, we derive the KL divergence as:
A
X + (1 Ar) log X
2o (0 = L De log
f(x;; 0) 20
- 2n log f (ax;; 0) + C,
where 1 / is 1 x4, * 2g otherwise O, and C' constant. Moreover, given ag by definition and therefore Ao = 1, To in (1) can also be written into the final formulation:
T
L,b = Eg((
X Ae (xa) 1 To log fmz;; 0)
t=1
For Xy that represents a sequence of random variables computed losses for each token, arriving the final expression for the whole sequence:
(æo.1, No,,), we can add all
T
Lub (gy((
ΣΣ Xe
(x0,,,) txo.. X,,n log fxt, :): 0).
For multinomial diffusion, we follow Zheng et al. (2023) to adopt reparameterized form, which results in the above formulation as well.
C ADDITIONAL EXPERIMENTS
C.I DYNAMIC SEARCH DEPTH IN DIFFUSEARCH
In explicit search algorithms, the search depth is predefined either through exact parameter in depth-first first search, or a related parameter such as the number of simulations as MCTS. In DIF- FUSEARCH, the deeper search is achieved by extending the context length of the input. In this
18

---

Published as a conference paper at ICLR 2025
Table 6: Action accuracy when increasing implicit search depth by extending context length. We compare training separate models on D with different horizon h and single model on D with 4. For the single model, predicting future steps equal to or larger than 4 requires the model's extrapolation ability.
Future Step
Length
Separate Model
Single Model
0
23.36
34.71
1
168
37.35
38.11
2
248
38.82
37.41
3
328
41.31
36.99
4
408
39.34
36.78
5
488
40.87
36.02
6
568
41.04
34.56
7
648
41.69
32.72
section, we present the results of training single model for dynamic search depth, compared with separate models in the previous sections. We convert the learned position embedding to RoPE (Su et al., 2024), enabling the utilization of a context length beyond what was encountered during train- ing inference time. During training, horizon h is randomly chosen from the interval [1, 4] for each data to expose the model to various input lengths. As shown in Table 6, the single model sur- passes the one-step policy with a lookahead of up to 5 future steps, exceeding the training stage's future step of 3. Nonetheless, a diminishing trend emerges we escalate the search depth, possibly attributed to the constrained training data and limited context extension capability of the current RoPE-based model. We leave more effective context extension beyond training stage for implicit search to future work.
C.2 SCALING BEHAVIOR WITH MORE DATA
Table 7: Detailed action accuracy with increasing mode size on 10k and 100k games.
Model Layers
Transformer S-A
DIEFUEEARCH S-ASA
10k games (660k records)
1
21.92
11.97
2
23.28
26.61
4
23.05
36.49
8
22.10
41.31
16
21.55
42.87
100k games (6.6M records)
1
29.62
11.32
2
35.40
31.27
4
36,93
42.57
8
36.58
48.66
16
35.03
51.89
Improved next-action accuracy In Table 7, we show the comparison of Transformer S-A and DIFFUSEARCU ASA when scaling data and model size. We can see the performance of DIF- FUSEARCH consistently improves with more data and model layers, while that of Transformer S-A converges with layers with 10k games and layers with 100k games. Further increasing model size is still useful for DIFFUSEARCH under both data-limited (e.g., 10k games) and relatively data- sufficient (e.g.. 100k games) scenarios.
Improved future accuracy In Table 8, show the quality of predicted futures for DIFFUSE- ARCH with horizon h = 4. We find when we scale data to 100k games (6.6M records), almost all the future actions valid (i.e., legal), future states valid (i.e.. the predicted states tokens correctly represent valid board state), and the action-state transition dynamics are well learned. Moreover, the best action percentage (i.e., action accuracy) also improves greatly compared that
19

---

Published as a conference paper at ICLR 2025
Table Detailed percentage of future predictions with increasing mode size on 10k and 100k games for DIFFUSEARCH with horizon h 4. Best a; percentage when future step = 0 is equivalent to action accuracy.
Future Step
Valid ag
Best a;
Valid s/
Match a;-1-S;
10k games (660k records)
0
98.40
41.31
100.00
1
79.33
20.72
97.35
37.22
2
50.40
4.60
53.59
6.74
3
50.07
3.00
51.26
3.30
100k games (6.6M records)
0
99.85
48.66
100.00
1
99.72
32.52
99.89
99.12
2
99.67
19.67
99.88
99.13
3
99.17
13.85
99.92
93.71
the 10k games setting. This demonstrates the potential accurate future world modeling through model scaling.
20

---

Published as a conference paper at ICLR 2025
Table 9: Example of training example for each training paradigm. We show horizon for illustration.
Paradigm
Input
Outpet
zbokhhurpppp. ppp M...oooPPP...... Mec.oe........... P............................... P.........PP ...PP
S_A
PE.PPPRNBOKB C.PPPPRNBOO..BBKKKK0.2....
d2d4
=bskhhrrppp-ppp ................................ ................................ P..P......PP WIN[56]
S_V
PRPPPRNBOOK RWA.ANBKKKKKKKKKKKKOOOOOO RwKQkg.
cbokbuupppoppp Jl. R....o P...... P........PPP ...PP WIN[6I]
SA-V
PAPPPRNBOKB RwKQkg-222 2.3.. d264
wbgchmmppp--ppp...... .....o P............. P.........I.
S-AA
ROPPPRNROKRRRKOOO 2.3.
d264 c564 f344 g816
=bglbbrropppppp pww.m P............................... P. PP d2d4 rhajkborpeep-ppe pop............................. .... .P...PP PP......NN. .N.... PPP. PPPRNROKD R5KQkq-O.. 0.3.. e5d4
rhagbbrrpppo ppp............................. www............................. DP. ...........P.. PPP. PPPRNBOKR K.Kkk-00.44 F34
PEPPPRNBOKB RwKQQk...3.
xbggbarpppo.pp.. wwwpyppppppp...........
NP. PPP. PPPRNBOKK RbKQQq- 0.4., g816
thgpharppeo-popp..... P............................... Po P......N. PPPPPPRRRRRRRR.. RKKQQq-0.. rbakbarpppoopgg0 ptt.rpppr.ppp...... ppp............................. p......N.
S-ASS
PEPPPRNBOKB RwKQkq-2.3.
:bolbarpppo.ppp ................................
PPP.....PPP N.
N. PPP. PPP.NERK RwKQkg-0A4 0.4..
:bghbarpppe-ppe ppp... ip........
NP. ........ PPE.PPPRNBOKE RbKQkkKKKk 0.4..
=.bgbhurpppp.pp Rw
wwwppppppppppp.... www.ppp.pppp.......pp.. Mec.eep...... P........PP N.PP
S-AVAV
a264 WIN[6I] cddi WIN[69] 1344 WIN[6UJ gsis WIN[68]
WIN[6]] Thgkberpop--ppP ppp.ppppppp................ m..... P.... 0.3.. c5d4 WIN[69]
WWW.PPPPRNOKK.RKKCC RbK
S-AVSAV
cbokburpopp.og
PRPPPPRNROK RwKQkg-22.
N..PP
zhokborpppe ppp...... app............................. pP. p.......NN......
RwKKkg-
WIN[60] :bglsburper--pp
NP. PPP. PPPRNBOKB RUPPRBBOKBRRKKQQQ...OO
g8f6 WIN[68]
C.3 EXAMPLE OF TRAINING INSTANCE
We show an example of each training paradigm Table 9.
C.4 ADDITIONAL CASES
In Figure 5, we provide the predictions of Transformer (S-A) and DIFFUŠFARCH on more challeng- ing puzzles. We also show the prediction of all models in Figure 6, where all models are trained on 10k games and 100 MCTS simulations are used for Transformer with MCTS.
21

---

Published as a conference paper at ICLR 2025
Transformer (S-A)
DiffuSearch
Transformer (S-A)
DiffuScarch
+
MAI
Transformer (S-A)
DiffuSearch
Transformer (S-A)
DiffuSearch
VAA
LAM
Transformer
DiffuSearch
Transformer (S-A)
DiffuSearch
Figure 5: Additional prediction cases on challenging puzzles.
Transformer (S-A)
Transformer (S-V)
Transfurmer (SA-V)
Transformer (MCTS)
DilfuSearch
Transformer (S-A)
Transformer (5-V)
Transformer (SA-V)
Transtormer (MCTS)
Figure 6: Additional cases challenging puzzles compared with all baselines.
22


Reviews:
[{"review_id": "YWCuwzndTv", "reviewer": "Reviewer_SREQ", "summary": "The paper proposes to address Chess playing with an implicit search using diffusion modeling. The paper is compared to Transformers networks and MCTS. It reaches an Elo of 1728 when trained on data from Stockfish.", "questions": "What would you propose to improve the Elo of your program?", "limitations": "", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The paper compares the proposed approach to other approaches", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "t6oHyWz7i3", "reviewer": "Reviewer_KzTF", "summary": "This work aims to explore the possibility of planning using LLM without explicit search (e.g., MCTS), i.e., employing implicit search. Specifically, this paper proposes a method named DiffuSearch that looks into the future world by diffusion modeling. DiffuSearch considers the bidirectional self-attention  architecture and the multi-step diffusion generative process. This work focuses on a study of a specific board game--chess. The numerical experiments demonstrate the efficacy of the proposed DiffuSearch approach.", "questions": "Please refer to the \"Weaknesses\" section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "Overall, the idea of this work is novel and intuitive.\n\nThe paper is well-written and easy-to-follow.\n\nThe technical contents like the theorems as well as proofs are well-organized and sound.\n\nThe experimental setup is detailed and clear. The empirical results substantiate that DiffuSearch outperforms the existing baselines.\n\nThe demonstration plots like Figures 4 and 5 are very clear and intuitive.", "weaknesses": "Line 83. The link of source code is invalid. I hope to see the code during the rebuttal. The authors may consider submitting via a zip file or providing a valid link to an anonymous repo.\n\n\nI would expect more detailed explanation of Figure 1 in both the main texts and the caption of Figure 1. As the comparison between the explicit and implicit searches is the main idea of this work. For example, the authors could explain the difference of explicit and implicit searches by describing Figure 1 more carefully. The discussion of the structures in Figure 1 is not enough. I felt that the difference between the explicit and implicit that authors mentioned in the texts is not well-connected to Figure 1.\n\nThe 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?\n\n\nIn 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? If it is feasible, I would like to see if it converges when you run the same depth as the baseline method, i.e., depth = 24 or 25.\n\n\nI 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)? If it is feasible, the authors may show at least one scenario for all of these three baseline comparisons. If not feasible, can the authors provide some explanations/reasons?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "hpIiMdWbQ0", "reviewer": "Reviewer_XeCm", "summary": "This paper trains a transformer to imitate actions from a chess engine, while using discrete diffusion modeling to incentivize a form of implicit search during action selection. The diffusion modeling distills forward and backward prediction into the policy network, whereas the baseline transformer with MCTS involves only forward prediction. The paper includes several comparisons and ablations, and claims that diffusion modeling improves performance.", "questions": "1. Can you provide more detail on precisely how MCTS is being combined with the policy network?\n\n2. The Best $a_i$ and Match $a_{i-1}-s_i$ lines in Fig. 2 (left) are not discussed in the text. It seems like the model is terribly inaccurate for the actions it actually takes. How does it do so well?\n\n3. In Table 8 (Appendix B), if $s'$ matches $f(s,a)$ with 99% accuracy, why does Best $a_i$ accuracy drop by ~1/3 after each step?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper investigates whether diffusion modeling can be helpful for emulating search using a feedforward network. This is an interesting question, as transformers have generally struggled thus far to solve problems requiring search. Normally, the solution is to add explicit search in the form of MCTS using the outputs of the transformer. This paper tries to improve the policy network instead, and indeed provides some evidence that transformers can simulate search using search with a single forward pass.", "weaknesses": "1. I wish it were more clear what exactly the paper is arguing. The paper seems to provide evidence for the following claim: \"If we must use transformers in tasks that require search, it is more efficient and effective to train them via diffusion to do implicit search than it is to add explicit search via MCTS.\" However, it seems as though the paper is instead arguing that implicit search is better than explicit search. There doesn't seem to be evidence for this, as the method still relies on a dataset that comes from a Stockfish oracle that uses explicit search (and which vastly outperforms the diffusion model). If the paper is arguing for the former, weaker claim, then I think it does an fine job of this, provided the intro/discussion/etc. are updated to make this more clear.\n\n2. 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. How does the paper combine the policy model with MCTS to do search for each different \"future paradigm\"? 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.\n\n3. The text is often rather imprecise, and in some places even a bit misleading. For instance:\n    1. AlphaGo pre-dates the transformer, but the abstract makes it seem like people were working on LLMs prior to AlphaGo.\n    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.\n    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.\n    4. The problem setting defines the value function with respect to a single policy, despite the fact that the players generally have different policies.\n    5. In Sec 4.5 (line 288), S-AVAV does not lead to a significant performance improvement, though S-ASS does.\n    6. It is unclear what the names in Table 4 actually mean. The process is described rather vaugely in the adjacent paragraph.\n    7. The scaling trends in lines 411-418 are presented as scaling \"laws\", but it is far from clear that these results will have that level of robustness beyond that single experiment in this paper.\n    8. Line 429 suggests that the method correctly values \"the long-term positional benefits of opening lines for its rooks.\" However, the model is playing as white, has only one rook, and sacrifices it in the first turn, preventing any such long-term benefits.\n    9. Section 5.2 (line 468), suggests that diffusion models might be \"a way to substitute for conventional handcrafted search algorithms,\" but offers no evidence of this claim, since the oracle that provides the training data for the model uses exactly that sort of conventional handcrafted search algorithm.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "U2srL9WAvJ", "reviewer": "Reviewer_LQsg", "summary": "DiffuSearch is a novel approach to enhancing the planning abilities of Large Language Models (LLMs) without relying on explicit search techniques like Monte Carlo Tree Search (MCTS). Developed in response to the limitations of current next-token prediction models in long-term planning, DiffuSearch uses discrete diffusion modeling to perform implicit search by predicting and utilizing future states. The method is implemented and tested on the game of chess, where explicit search has traditionally been essential. In extensive experiments, DiffuSearch outperforms both searchless and explicit search-enhanced policies, demonstrating a 19.2% improvement over one-step policies and a 14% improvement over MCTS-enhanced policies in action accuracy. Additionally, it shows a 30% enhancement in puzzle-solving abilities and a significant 540 Elo increase in game-playing strength compared to explicit search methods. These results suggest that DiffuSearch’s approach of internalizing a world model within the policy, without intermediate components, could be a promising alternative to traditional explicit search techniques in AI problem-solving.", "questions": "Line 060: It is unclear if the Ha & Schmidhuber citation is correct or necessary here. Work has existed prior to this paper[2] and there are many more works that use it. This is a fundamental concept and probably does not need a citation unless it is citing a textbook. It should be removed.\n\nIs the performance conditioned on the amount of train time and test time compute? What is the trade-off in FLOPs at train time versus test time?\n\nLine 083: Code link is broken. Make sure to correct this for publication.\n\n\n[2] Triggs, B., and Cameron, S. (1991). ““The Oxford robot world model,”,” in Expert systems and robotics (Springer Berlin Heidelberg), 275–284.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. Superior performance: DiffuSearch significantly outperforms baseline models, as evidenced in Table 2. It demonstrates a substantial improvement over both searchless and explicit search-enhanced policies, showing the potential of using diffusion as a model for search in chess.\n2. Novel methodology: DiffuSearch incorporates future states and actions through discrete diffusion modeling, enabling it to leverage future information for improved next action prediction without relying on explicit search. This approach offers a new perspective on implicit search in AI.\n3. Sample efficiency: Despite using 20 times fewer data records than some baseline models (e.g., SA-V), DiffuSearch demonstrates superior performance with approximately 10% higher action accuracy.", "weaknesses": "1. 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.\n2. 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.\n3. 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.\n[1] Ruoss, A., Delétang, G., Medapati, S., Grau-Moya, J., Wenliang, L. K., Catt, E., ... & Genewein, T. (2024). Grandmaster-level chess without search. arXiv preprint arXiv:2402.04494.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)