Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning
Authors: Jiacheng Ye
Jiahui Gao
Shansan Gong
Lin Zheng
Xin Jiang
Zhenguo Li
Lingpeng Kong
Venue: ICLR 2025 Poster (2025)
Keywords: diffusion model
autoregressive model
complex reasoning
Primary area: foundation or frontier models, including LLMs
TL;DR: We introduce a discrete diffusion model that significantly outperforms autoregressive models in solving complex reasoning tasks.

Abstract: Autoregressive language models, despite their impressive capabilities, struggle with complex reasoning and long-term planning tasks. We introduce discrete diffusion models as a novel solution to these challenges. Through the lens of subgoal imbalance, we demonstrate how diffusion models effectively learn difficult subgoals that elude autoregressive approaches. We propose Multi-Granularity Diffusion Modeling (MGDM), which prioritizes subgoals based on difficulty during learning. On complex tasks like Countdown, Sudoku, and Boolean Satisfiability Problems, MGDM significantly outperforms autoregressive models without using search techniques. For instance, MGDM achieves 91.5\% and 100\% accuracy on Countdown and Sudoku, respectively, compared to 45.8\% and 20.7\% for autoregressive models. Our work highlights the potential of diffusion-based approaches in advancing AI capabilities for sophisticated language understanding and problem-solving tasks. All associated codes are available at \href{https://github.com/HKUNLP/diffusion-vs-ar}{https://github.com/HKUNLP/diffusion-vs-ar}.

Published as a conference paper at ICLR 2025
BEYOND AUTOREGRESSION: DISCRETE DIFFUSION FOR COMPLEX REASONING AND PLANNING
Jiacheng Jiahui Gao? Shansan Gong,, Lin Zheng'
Xin Jiang", Zhenguo Li", Lingpeng Kong' The University of Hong Kong 2 Huawei Noah's Ark Lab carscnye@coneect. hwt.connnnnc....hkkkk.... hku.hk
ABSTRACT
Autoregressive language models, despite their impressive capabilities, struggle with complex reasoning and long-term planning tasks. We introduce discrete dif- fusion models as a novel solution these challenges. Through the lens of subgoal imbalance, we demonstrate how diffuusin models effectively learn difficult sub- goals that elude autoregressive approaches. We propose Multi-Granularity Dif- fusion Modeling (MGDM), which prioritizes subgoals based on difficulty dur- ing learning. On complex tasks like Countdown, Sudoku, and Boolean Sat- isfiability Problems, MGDM significantly outperforms autoregressive models without using search techniques. For instance, MGDM achieves 91.5% and 100% accuracy Countdown and Sudoku, respectively, compared to 45,8% and 20.7% for autoregressive models. Our work highlights the potential of diffusion- based approaches in advancing Al capabilities for sophisticated language derstanding and problem-solving tasks. All associated codes are available at http:///tttll..oo///UUUU.PP///fO
INTRODUCTION
In recent years, autoregressive language models (LMs; Bengio et al. 2000) have dominated the landscape of natural language processing and artificial intelligence. Empowered by scaling laws (Kaplan et al., 2020), these models have demonstrated impressive performance across various ap- plications (OpenAI, 2022; Achiam et al., 2023; Anthropic, 2023; Team et al., 2023, inter alia). However, this apparent success masks significant limitations that are becoming increasingly evi- dent. Autoregressive models inherently struggle with tasks requiring complex reasoning, long-term planning, and maintaining global coherence (Bubeck et al., 2023; Valmeekam et al., 2023; 2024; Dziri et al., 2024; Kambhampati et al., 2024). These shortcomings represent substantial challenges in developing Al systems capable of robust problem-solving and adaptable cognition (Wu et al., 2022; Zhao et al., 2023; Trinh et al., 2024; Yao et al., 2023; Shinn et al., 2024, inter alia). While autoregressive approaches have driven considerable progress, their limitations suggest that they may not be the optimal solution for all aspects of machine intelligence. As the field evolves, it becomes increasingly important to explore alternative paradigms that can address these inherent drawbacks and potentially offer new avenues for advancement in AI capabilities.
In response these limitations, recent research has focused on addressing the inherent constraints of autoregressive models. Various strategies have been explored, including the integration of search algorithms at inference (Yao al., 2024; Besta et al., 2024) and the incorporation of backtracking supervision during training (Lehnert et al., 2024; Gandhi et al., 2024). However, these approaches are not without their own drawbacks: the former often incurs significant computational costs, while the latter frequently results in verbose inputs and suboptimal performance.
To address this challenge, we argue for fundamentally different modeling approach: discrete diffu- sion models. While most contemporary language models are autoregressive, diffusion-based models have become predominant in image (Dhariwal & Nichol, 2021; Rombach et al., 2022; Peebles & Xie, 2023) and video domains (Ho al., 2022; Wu et al., 2023a; Brooks et al., 2024). Diffusion models are also gaining traction in various other applications, such as protein desiging (Xu et al., 2022; Hoogehoom et al., 2022b; Corso et al., 2023) and planning in reinforcement learning (Jan-

---

Published as a conference paper at ICLR 2025
ner et al., 2022; Ajay al., 2022; Chi al., 2023). In this work, reveal that discrete diffusion models demonstrate significantly superior performance compared to the autoregressive counterparts, particularly in tasks requiring complex planning and reasoning.
To substantiate this argument, we first examine the problem through the lens of subgoal imbalance ($3.1). We present both theoretical and empirical evidence via a synthetic planning task (Figure 1) to illustrate why autoregressive models struggle with these types problems, often achieving near- random performance. In contrast, demonstrate how diffusion models effectively learn the sub- goals that challenge autoregressive models ($3.2). The key insight lies in the training objective of diffusion models. where difficult subgoals are decomposed into a diverse range of interrelated views within a multi-view learning framework (Xu et al., 2013), Each of these views is more manageable, resulting an overall easier and more effective learning process.
Building upon these insights, we propose natural extension to current discrete diffusion models, which we term multi- guatiigranullrrity diffusion modeling (MGDM; $3.3). This approach prioritizes different subgoals based on their difficulty during the learning process, leading to more effective learning outcomes and faster convergence.
In our experimental evaluation ($4), we focus on substantially more complex problem-solving tasks, such as Countdown (Gandhi et al., 2024) and Sudoku (Garns, 1979). These tasks demand extensive planning over large number of combinations and pose challenges even for commercial Large Lan- guage Models (e.g.. GPT-4 Achiam et al. 2023). Notably, without employing any search techniques, MGDM achieves 91.5% and 100% accuracy on Countdown and Sudoku respectively, while its toregressive counterpart only solves 45.8% and 20.7% of the problems. Additionally, we conduct experiments on the Boolean Satisfiability Problem (SAT), NP-complete problem (Cook, 1971) that represents a wide range of constraint satisfaction problems. Our model exhibits superior per- formance solving SAT problems with higher accuracy compared the autoregressive alternative, particularly when dealing with an increased number of variables and constraints. Through this sys- tematic exploration, we aim demonstrate the potential advantages of diffusion-basee approaches in addressing sophisticated language understanding and generation challenges.
BACKGROUND
2.1 AUTO-REGRESSIVE MODELING
Let == (æI...., æN denote sequence drawn from data distribution q(x). For decades, it has been common to factorize the joint probabilities of a sequence of tokens as the product of conditional probabilities (Jelinek, 1980; Bengio al., 2000):
N
Po (x) = Po (xx1) II (2n 11:n-1),
(1)
n=
where 0 parameterizes the model distribution and Tin-1 = Ei..... In-1- In order to optimize the generative model po (x) to fit the data distribution q(x), we optimize the negative log-likelihood:
N
LAR - gx(x log pa
By(x))
tog -1 1).
(2)
2.2 DISCRETE DIFFUSION MODELING
Discrete diffusion models (Sohl-Dickstein al., 2015; Hoogeboom et al., 2021; Austin et al., 2021) are a class of latent variable models characterized by a forward noising process and learned reverse denoising process. The forward process (xx) = IT= ((x, xx corrupts the original data Xy x into a sequence of increasingly noisy latent variables X1.7 X).... .x.. The backward process learns to gradually denoise the latent variables to the data distribution given by:
Px(x)=
Σ p(xr) Xi:
I po (x, (x/11/x))
(3)
2

---

Published as a conference paper at ICLR 2025
Input: 4.7 5,8 3,11 /3,
10f
shuffled odges
start.goul
Output: 3,1 1.7 17,5 5,8 8,9
Accuracy
AR
Data Size
Reverse AR
109
Teacherless Diffusion
Hard Subgoal
25
10
0
3
Planning Distance
Planning Distance
10
Planning Distance 3
Figure 2: (Left) Accuracy of different method given 50k training data. (Right) Minimum data size required to solve (i.e., accuracy above 90%) subgoal at each planning distance.
Figure 1: The planning task.
Due to the intractable marginalization, we typically optimize variational upper bound the neg- ative log-likelihood:
T
LDM Eg(xo)
DKL (x(xT) (xo) [x(xT)]
+
Σ
E Ex(x xo)
DKL (x ( po
lT
t=2
- - (xo) [log po (x [x1]]
(4)
Le
where LT is a constant when one uses a fixed prior (. By defining both the forward and backward distribution categorical distribution, e.g., q(x, (x-1) Cat(x;; Q x,-1) where Qi is pre-defined K x K transition matrix and K is the size of categories, and ( (x,) = q(x,11 Nx, f (x,;; 0)), the forward process posterior (xx,11 (xxXx) xo) and each KL term can be calcu- lated analytically (Hoogeboom et al., 2021; Austin et al., 2021).
3 SUBGOAL IMBALANCE AND MULTI-GRANULARITT DIFFUSION MODELS
In this section, we employ a motivation example ($3.1) 1) to elucidate the challenges faced by autore- gressive models in specific scenarios. Through this analysis, introduce the concept of subgoal imbalance-wwerein some subgoals are inherently more difficult than others-which offers insights into these diffieulties. We then extend our discussion $3.2 to examine how diffusion models can more effectively address and learn these hard subgoals, effectively overcoming the limitations observed in autoregressive approaches. We finally propose Multi Sratti-Grrallrttyy Diffusion Modeling (MGDM; $3.3) as natural extension of discrete diffusion models to better address these challenges and improve performance on complex tasks requiring planning and reasoning.
3.1 SUBGOAL IMBALANCE IN AUTOREGRESSIVE AND DIFFUSION MODELING
We designed simple planning task to serve as running example. Consider the example in Figure I, where the input for the task consists At the end of the input sequence, the start and goal nodes specified to indicate the path the model needs to find. The objective of this task is to identify the correct path in the graph and output its constituent edges. The complexity of this problem arises from distracting factors (highlighted in orange) that potentially mislead the path selection. For instance, at node 7, with the goal being node the model must plan over distance 3 nodes determine that the correct next choice should be node rather than O. We define this span as the Planning Distance (PD), parameter adjustable in our synthetic task data. Intuitively, as the PD increases, the model faces greater difficulty in learning to determine the correct subsequent node. We formalize this intuition as subgoal imbalance.
set of shuffled edges from the graph shown below.
Proposition 1 (Subgoal imbalance due to the unknown data distribution qxx)) Given the data sampled from unknown data distribution q(x), the difficulty of learning each subgoal on can differ significantly based on how we parametrize the model distribution, and some subgoals may require substantially more data to learn or may even be infeasible learn.
3

---

Published as a conference paper at ICLR 2025
Subgoal imbalance in autoregressive modeling. Given the data x sampled from an unknown data distribution q(x), the naive autoregressive modeling parametrizes the model distribution pe (x) as po ( II Po ( I1:n-1), the difficulty of learning each subgoal In, can differ significantly as given only the left context, and some subgoals may require substantially more data to learn or may even be infeasible to learn.
Setup. We synthesize the data with only one distracting path. We randomize node numbers in [0, 10j and the intersection positions in 50, We further designed this task to be symmetric, en- suring that simply training with reversed output, as suggested by Bachmann & Nagarajan (2024), cannot solve subgoals with all PDs. For comparison, we include Auto-regressive (AR), reverse AR (Bachmann & Nagarajan. 2024), and teacherless training (Monea et al., 2023; Bachmann & Nagarajan, 2024), which can be seen as a lookahead method that produce all target tokens from the source input, and our proposed diffusion model (detailed in $3.2). For all the models, we keep the model architecture fixed as the same 3-layer Transformer with approximately 6M parameters. More details can be found in Appendix SC.
Discussion. We examine the performance of all the models in two scenarios. In the first scenario, generate fixed number of 50k instances with mixed planning distance. We plot the accuracy on the held-out evaluation set for each model in the left figure of Figure 2. Our findings indicate that autoregressive models (AR and Reverse AR) are only effective in solving cases where the PD equals O or 1 (or equivalently, 5 and 4 in the reverse setting). Due to the aforementioned subgoal imbalance phenomenon, when PD is less than 2, the task barely involves any planning, allowing models to sim- ply copy from the input with ease. However, for larger PDs, AR models barely outperform random guessing. Teacherless training fails to adequately fit the training data, resulting in the production of illegal paths. In contrast, our diffusion model achieves perfect accuracy across all PD values.
In the second scenario, we investigate whether the challenging subgoals can be naturally resolved through data or model scaling, akin to the success observed in large language models (Kaplan al., 2020; Wei et al., 2022a). To investigate this question, we gradually increase the size of the dataset for each model with different PDs and plot the minimum data size required to solve the subgoal in the right figure of Figure 2. We find that the autoregressive models (AR and Reverse AR) can learn the easy cases of PD equal to O and (or equivalently, 5 and 4 in the reverse setting) with only 10k data points. However, exponentially larger amounts of data are required to address in- creasingly challenging subgoals. Both teacherless training and diffusion models exhibit similar U-shaped curve in their performance. This similarity can be attributed to the fact that teacherless training can be conceptualized as special case of diffusion without an iterative noising and denois- ing process. In these models, solving edge PDs necessitates slightly more data. We hypothesize that this phenomenon occurs because the distance to other positions shorter from the middle position (i.e., higher closeness centrality), thus providing the middle position with more nearby tokens to aid prediction. Overall, autoregressive models require significantly more data to address all PDs compared to diffusion models, highlighting their relative data inefficiency.
In addition to our previous experiments, we conducted a series of tests to examine the effect of increasing the parameter count in autoregressive models while maintaining fixed dataset size of 50,000 instances. Our findings reveal that scaling the original million parameter model to 85 million, 303 million, and 1.5 billion parameters fails to resolve all PDs. Only upon fine-tuning substantially larger model, specifically the LLaMA 7B model (Touvron et al., 2023), did we observe successful resolution of all PD subgoals.
3.2 EFFECTIVE HARD SUBGOAL LEARNING IN DIFFUSION MODELING
These experiments collectively indicate that diffusion models are significantly more effective in learning challenging subgoals arising from subgoal imbalance. To elucidate why diffusion models exhibit this superior capability, we first establish a connection between autoregressive (AR) models and diffusion models by reformulating Equation (4). Instead of evaluating the KL divergence be- tween two complicated categoricals (Hoogeboom et al., 2021), we consider discrete diffusion with absorbing state and simplify it as the weighted cross-entropy losses (Austin et al., 2021; Zheng al.,
4

---

Published as a conference paper at ICLR 2025
2023; Shi et al., 2024; Sahoo et al., 2024):
N
DKL. DKI.g((x-11xx4 Xo) (po (x)- (((
xx Ko. X0,n log [x 0)n,
(5)
where w(t) -= (0, is a time-dependent reweighting term which places higher weight 1-oz when l. approaching 0. We then rewrite Equation (4) as:
N
T
LDM = (x0o)
w (t) yx(x No) u X1, 0),
(6)
n=l
log puun (2nn (m)gn))
where u(xu, n) - -1x, -x Nnn log fxx(;) n is the cross entropy loss token n.
We can systematically compare the losses of au- toregressive (AR) and diffusion models (DM), specifically log (aR F1:n-1) and log (DM | Ligh),, pressed in Equations (2) and (6), respectively. In Figure 3, we examine a specific hard subgoal with Planning Distance (PD) equals 3 in both model types. The loss levels of AR and diffu- sion models are depicted using blue and red lines, respectively. The overall loss log (TM ((tt)) in the diffusion model remains relatively low compared its autoregressive counter- part - log PAR (r.,. (ar, formance of the diffusion model on these challenging subgoals in our experiments.
10a
log pase An (1m-1))
Pond (2ra ((rg
log Pond (3ta (orse-)))
10-2
Loss
10-4
#1.n-1), corroborating the superior per-
2
10
Timestep
Figure 3: Loss for a specific hard subgoal, i.e., PD=3, Diffusion and AR modeling. We also show the unweighted loss u(xo: xx: n; O) at different timestep and context Xr in diffusion modeling.
Further analysis of the unweighted loss u(Xo, Xt, rı; 0) in the diffusion model, based on 1,000 samples of x; g(x, xu), reveals a clear trend: as the number of timesteps increases, resulting in more noise in X,, objectives in smaller timesteps (i.e., recovery from less noisy data) become significantly eas- ier to learn. From multi-view learning perspective (Xu al., 2013), each x, can be interpreted as a distinct view of Xo. where each view provides different information about xo- In the diffusion process, by exploring the consistency and complementary properties of different views offered by diverse range of interrelated objectives u(Xp. x,, n: 0), our findings suggest that objectives chal- lenging to learn in AR models become more effective, promising, and exhibit better generalization in diffusion models.
This phenomenon particularly evident when examining scenarios where mask noise applied to positions after the hard token, i.e.. X, 21:a-1, where the diffusion model learns the hard subgoal similarly to AR models. We plot this loss as log PDM -1m--1 in the figure. Unlike in the AR model, where this learning is consistently difficult, in diffusion models, this challenging subgoal is addressed at a much more manageable level during the learning process.
3.3 MULTI-GRANULARIIN DIFFUSION MODELING
These observations provide valuable insights, i.e., diffusion modeling builds diverse range of interrelated views from the data X to handle a challenging subgoal. To handle multiple challeng- ing subgoals in real data, we should prioritize different subgoals based on their difficulty during the learning process to achieve more effective learning outcomes and faster convergence, and this naturally translates prioritizing difficult views as the learning subgoal depends on learning interrelated views related to it. Building on this, we propose the multi- malliiggannllrrty diffusion model as natural extension the discrete diffusion model.
In practice, to optimize Equation (6), we typically employ Monte Carlo sampling, which results in:
N
LDM
w())uxX, nz, 0).
(7)
R=
5

---

Published as a conference paper at ICLR 2025
For sequence of length N, the probability of sampling the same X, in AR is 1. However, in diffusion, this probability reduces 1/CK due to the randomness sampling x,, potentially reducing the training efficiency of diffusion models. We note that Equation (7) employs a sequence- level reweighting term w(t) to indicate the importance of Xs. However, individual tokens within the sequence, given their imbalanced difficulties, not properly reweighted. To address this, propose multi taatii-rrulaartyy diffusion modeling (MGDM), which introduces an additional token-level reweighting mechanism to enhance training efficiency:
LMDDM ΣΣ Σ
((t)o(Xe,)u(xx) Xx, 0)
(8)
where ((xt,n) = a(1 exp(-u(-))) the adaptive token-level reweighting term. Setting 3 reduces the relative loss for easy tokens while emphasizing harder tokens, and is used to control the relative reweighting magnitude. For inference, we employ an easy-first TopK decoding strat- egy, which has demonstrated superior performance compared to the random decoding method used by Austin et al. (2021). This finding aligns with similar observations documented in prior stud- ies (Savinov al., 2021; Zheng al., 2023). We provide detailed derivation and algorithm of the training and inference process in Appendix SA and &B, respectively.
4 EXPERIMENTS
In Section $3.1 we show our model works well on a straightforward planning task with only one hard subgoal. However, it is important note that real-world scenarios often involve instances with multiple challenging subgoals. In this section, we aim to assess the performance of our model in tackling three considerably more complex problem-solving tasks that necessitate deliberate plan- ning. Detailed experimental setup can be found in Appendix $C.
4.1 COUNTDOWN
Countdown (Countdown, 2024) is a mathematical reasoning challenge and is a generalized version of the game of 24, which even advanced models such as GPT-4 struggle with (Yao et al., 2024). The goal of Countdown is to use the given numbers and arithmetic operations (+ - */) to obtain target number. For example, given 4 numbers "97,38,3,17" and target number "14", sttep-by-stey solution is 997-38-59,59- 19-
Table Results the Countdown (CD) task with increasing complexity.
(2024) to generate 500k problems with target numbers ranging from 10 to 100 and randomly hold out 10% of the targets for 'futtof-distriiuttinn evaluation.
Setup. We follow Gandhi
Params
CD
CD
CD 5
Autoregressive
consider three subtasks with increasing complexity by varying the number of input digits in {3,4,5]. Given that search- augmented prompting approaches (Yao et al., 2024) have recently been employed to address the limitations of AR, also compare with such approaches by training on Countdown 4 and evaluating on the same game of 24 test as Yao et al. (2024).
6M
94.1
GPT-2 Scratch
85M
95.9
303M
96.4
41.3
Stream-of Saan--o-Seardee
250M
54.2
7B
95.7
41.1
6.7
LLaMA
13B
96.5
51.1
7.4
Diffusion
VDM
85M
99.1
73.4
16.3
D3PM RDM
85M
99.4
83.1
27.6
85M
99.5
87.0
45.8
6M
98.1
52.0
27.0
Baselines. Our primary comparison involves autoregressive models trained from scratch, employing the GPT-2 architecture (Radford et al., 2019) with parameter sizes ranging from 6M, 85M, and 303M (denoted as GPT-2 Scratch). We also include larger pre-trained AR models LLaMA (Touvron et al., 2023) with sizes 7B and 13B. These models are fine-tuned using the same dataset. In addition, Search (Gandhi et al., 2024), which augments the dataset with search trajectory such that the AR model be taught search. Furthermore, compare with several existing diffusion models,
MGDM (Ours)
85M
99.5
91.5
46.6
303M
99.9
88.3
39.0
compare with Stream-of-

---

Published as a conference paper at ICLR 2025
both continuous models VDM (Kingma al., 2021) and discrete models D3PM Austin et al. (2021) and RDM (Zheng et al., 2023). By default, we use the absorbing noise for discrete diffusion as it significantly outperforms the multinomial one (Austin et al., 2021; Zheng et al., 2023). Finally, we consider in-context learning (Brown et al., 2020; Wu et al., 2023c; Ye et al., 2023a) based on GPT-4 including vanilla input-output (IO), chain-of-thought (CoT; Wei et al. 2022b), CoT with Self-consistency (CoT-SC; Wang et al. 2023) and Tree-of-thought (ToT; Yao et al. 2024). We use 5 in-context examples following Yao et al. 2024.
Results on Countdown. As shown in Table diffusion-basel approaches demonstrate superior performance across all three Countdown tasks compared to autoregressive models, especially as the complexity of the tasks increases. We have several key findings based on the result. Firstly, the 6M diffusion model outperforms both the 303M GPT-2 model trained from scratch and the pretrained 13B LLaMA model, indicating that the modeling approach sometimes outweighs the sheer number of parameters. Secondly, while training with search trajectory supervision provide some benefits, its effectiveness is limited. Importantly, training the entire search trajectory as sequence poses additional challenges due to its long length, such as in the case of Countdown 5 where the search trajectories can span 60,000 tokens. Lastly, our model surpasses all previous diffusion models, demonstrating the efficacy of the multigranularity loss.
tereammoff-eaaree
does
Results on Game of 24. As shown in Table 2, the performance of the GPT-4 with IO, CoT, and CoT-SC prompting methods from Yao et al. (2024) is unsatisfactory for the given task, with only accuracy below 10%. The introduction of ToT, which in- corporates search algorithm designed by human experts into the decoding process, significantly enhances the performance of GPT-4. This integration allows the AR model to backtrack as needed, resulting in notable improvements. However, this paradigm requires the assessment of intermediate steps using LLM, resulting in considerable computational costs due to the need for multiple LLM calls. We list the token cost in Table 2 with more details in Appendix D. ToT consumes 186 times more tokens than MGDM, showcasing the 'internal' search ca- pability by promoting global consistency in diffusion modeling. In summary, our model, despite having a parameter size of only 85M, significantly outperforms both the AR task-specific model of the same size (GPT-2 Scratch) in terms of performance and the larger general pre-trained model (GPT-4) in computation cost, indicating is challenging for model scaling and decoding strategies substitute the advantages of modeling paradigm.
Table 2: Accuracy and token cost on game of 24.
Acc.
# Token
Prompting GPT-4IO IO
7.3
x28
GPT-4 CoT
4.0
x61
GPT-4 CoT-SC
9.0
x241
GPT-4 ToT
74.0
x x186
Supervised training
GPT-2 Scratch
18.8
x1
MGDM
76.0
xl
4.2 SUDOKU
100
100
GPT-2 Scratch
12.5
GPT-2 Scratch (85M)
Sudoku is a classical logic- based number placement puzzle that has gained popularity due its rigorous intellectual de- mands. The goal of Sudoku is to meticulously fill a 9 grid with numerical digits, ensuring that every column, row, and x 3 subgrid contains all the num- bers from to
Accuracy
80
GPT-2 Scratch (303M)
Accuracy
90
10.0
Accuracy
60
LIaMA (7R)
LLaMA (13B)
7.5
40
MGDM (6M)
GPT-2 Scratch MGDM
5.0
20
10a
10
70
5
Parameter Size (M)
Number Variables
Figure 4: (Left) Accuracy on Sudoku. (Right) Aceuracy on
boolean satisfiability problem with increasing difficulty.
Setup. We collect one million solved games from Park (2016) and use the first 100k as our training set and the subsequent 1k as the testing set. We employ the digit 0 to represent the vacant position that needs to be filled. We then transform the 9 x grid into a sequence of 81 digits, which serves as the model input. To illustrate, an example input appears as "080050000...60300000 for brevity), while the corresponding output is represented as tokenization, we treat each digit as separate token.
(omitted
During
7

---

Published as a conference paper at ICLR 2025
Random
TopK
0.75
reweighting
82.1
87.3
Accuracy
Original grnuence-reeighhtnng
83.1
88.5
0.50
noken-reweighte (a=0.25, 9=1)
84.9
90.4
token-reweighting (a=1, (=1)
82.4
89.3
AR
nekke-reweighting (o=0.25, (==)
82.4
87.9
0.25
No Rew
Seq Rew
Linear sequence-reweighhing
79.6
87.0
Seq
reweightng (a=0.25, 8=1)
83.2
88.0
0.00
token- reweighinng (a=1, (=1)
86.7
90.4
200000
token-reweighting (o=0.25, (=2)
85.6
91.5
Training Steps
Figure 5: Evaluation accuracy Table 3: Ablation training reweighting strategies throughout the training process for AR and MGDM with differ- ent reweighting strategies.
and inference decoding methods.
Result. We show the results in the left figure of Figure 4. As the size of the AR model increases, the performance remains unsatisfactory. For instance, the LLaMA model achieves performance of only 32.9 with 13B parameters. contrast, our model, which has only 6M parameters, able to perfectly solve all the problems, demonstrating the significant advantage brought by the modeling architecture.
4.3 BOOLEAN SATISFIABILITY PROBLEM
The Boolean satisfiability problem, commonly known as SAT, a foundational problem com- puter science that has been rigorously proven to be NP-complete (Cook, 1971). This challenging combinatorial problem is attractive as broad range search problems from domains such as soft- verification, test pattern generation, planning. scheduling, and combinatories be solved by reducing to an appropriate SAT problem (Gomes et al., 2008). The goal of SAT is to determine whether given Boolean formula represented in conjunctive normal form (CNF) can be assigned a set of values (0 or 1) to its variables, such that the formula evaluates to true (1). An ex- ample formula with three variables can be (x1 V-I2) (-x1 V233Va)) = and an corresponding
all routinely
assignment is 1 0, x2 = 0, X3
Setup. Given the number of variables n and clauses m, we iteratively sample A: = 3 variables (and their polarities) uniformly at random until clauses obtained. To ensure that get relatively hard instances of SAT, we take advantage of the well-studied family of random k-SAT problem (Ding et al., 2015) and set the m to be close to m = 4.258n + 58.26m-22 -2/3 given n, as has been observed that SAT solvers slow determine the satisfiability of a formula when m is near the threshold (Crawford Auton, 1996). We consider increasing numbers of variables from (5.7.9) and generate 50k training data for 5, and 100k for n 9, as well as additional Ik testing data for each n.
As shown in the right figure of Figure 4, MGDM performs well in solving scenarios with five variables, while the AR model falls slightly short. As the number of variables increases, both our model and the AR model experience a certain degree of decrease in accuracy. However, the per- formance gap between the two models widens the difficulty of the task increases. This indicates that our diffusion model exhibits more pronounced advantage in handling more challenging tasks than the AR counterpart.
Result.
4.4 ANALYSIS
On the Effect of Training and Decoding Strategies. As listed in Table 3, we find that changing the sequence-reweighhing strategies has only led to slight improvement in performance. However, when a suitable parameter is selected for roken-reweighting, a more significant improvement can be observed. Additionally, the easy first decoding (TopK) outperforms the random one, which aligns with previous findings (Ghazvininejad et al., 2019; Zheng et al., 2023). We compare the evaluation accuracy along the training process in Figure 5. By aligning the AR training steps with the diffu-
8

---

Published as a conference paper at ICLR 2025
AR converges rapidly, with the performance tends to plateau afterward. The utilization of our multi-gran gullti-raaularity loss, which incorporates sequence and token reweighting, demonstrates superior performance, particularly during the middle stages of training. This implies that the inclusion such loss function contributes to enhanced convergence during the training process.
sion process,
On Decoding Speed. We assess the
trade-off between accuracy and de-
(o) Accuracy vs Speed on Countdown 4 and 5 (b) Planning and Colculation Error at each stop
coding speed by comparing the per-
90
formance of the AR model (GPT-2
40
0.4
Scratch 85M) and MGDM (85M).
0.4
The speed metric is determined by
30
AR MGDM
AR
the number of samples processed us-
MCDM
MCDM
ing batch size of 1 on the NVIDIA
20
0.2
0.2
GeForce V100 GPU. As shown in
AR
10
Figure 6(a), MGDM can flexibly control the trade-off between accu- racy and decoding speed by vary- ing the diffusion timesteps. Notably, by employing just one diffusion step, MGDM demonstrates a remarkable 10x improvement in speed compared to AR, while maintaining superior ac- curacy with 75% and 12.7% com- pared to 45.8% and 5.1% of AR on Countdown 4 and 5, respectively. We observed that the slope of countdown is smaller compared to countdown 5 in the trade-off. This suggests that for tasks with lower complexity. diffusion demonstrates more noticeable speed advantage by setting a smaller diffusion step. addition, also indicates that sacrificing some efficiency for performance im- provements becomes particularly evident when dealing with more intricate tasks.
MGDM
0.0
0.0
10
Eq]
Eq2
Eq3
Ans. Overall
Figure 6: (a) Accuracy and speed (samples per second)
trade-off by varying the diffusion timesteps on Countdown 4 (left y-axis) and 5 (right y-axis) Ratio of planning (left y-axis) and calculation (right y-axis) error each reasoning step on Countdown 4.
Error Analysis: The Regretful Compromise. To gain a deeper understanding of eтror patterns in AR and MGDM, we conducted an error analysis on Countdown 4. For instance, given the input "7,38,3,1" " and the target number 14, a correct solution would be We divide the solution into four parts: equation I, equation 2, equation 3, and answer checking. First, from calculation perspective, we assess the enror ratio in each equation by comparing the left-hand side and the right-hant side, regardless of whether the correct number was chosen. As shown in Figure 6(b). the bar plot demonstrates that the majority of calculation errors for AR are concentrated in Equation 3. For example, given input "16,4,40,51" and target 87, the prediction of AR is "51-40=11, 16, 46 11+44-87" while the correct solution is 116/444400-51111191444333
The first two equations are calculated correctly, but the last one is incorrectly forced to equal 87. Moreover, in many cases, the error rate for the last equation in AR (48.9%) is significantly higher than that of the previous equations (0.2% for the first and 7.2% for the second). This indicates that the model's lack of planning capability results its realization of being unable achieve the goal only at the end. By that point, it cannot correct the previous errors, leading it "regretfully" to rely on incorrect equations reach the goal (e.g., the number 87 the above case). This significantly increases the number of calculation errors in the third equation. We call this phenomenon "The Regretful Compromise'. The reason for this is that the AR model made incorrect choices of numbers or operations in previous equations. This is demonstrated by examining the step at which the models fail the task, as depicted Figure 6(b). It is evident that there is notable frequency of planning errors in the first equation for the AR model, where the number of errors is significantly higher compared to our model. This highlights the limitations of the left-to-right decoding approach in AR, which adversely affects its planning ability.

---

Published as a conference paper at ICLR 2025
5 RELATED WORK
5.1 AUTOREGRESSIVE MODELING
Starting from Bengio et al. (2000) and later Sutskever al. (2011), the autoregressive modeling paradigm, where the prediction of token only depends on the preceding context, is widely adopted in modeling language, until recently (OpenAI, 2022; Achiam et al., 2023; Anthropic, 2023; Team et al., 2023; Touvron et al., 2023; Jiang et al., 2023; Bai et al., 2023, inter alia). Theoretically, the autoregressive Transformers have limited expressive power, but their capabilities can be expanded given sufficient chain-of-thoughh intermediate steps (Wei et al., 2022b; Merrill & Sabharwal, 2023; Malach, 2023). However, Lin et al. (2021) demonstrates that the expressing of some next-tokens requires super-polynomial computational resources and is NP-hard to approximate. Numerous ad- vancements have been made upon the AR paradigm to compensate for modeling deficiencies, such training (Bavarian ct al., 2022), future-token prediction (Qi et al., 2020; Glocckle et al., 2024), lookahead attention (Du et al., 2023) during the training stage, as well as search-augnen -uugmeted decoding (Lu et al., 2022; Xie et al., 2023; Yao et al., 2024, inter alia) during inference. In practice, autoregressive next-token pre- dictors are shown be ill-suited for planning tasks (Bubeck al., 2023; Valmeekam et al., 2023; 2024; Dziri et al., 2024; Kambhampati et al., 2024). Besides, Bachmann & Nagarajan (2024); Lin al. (2024) find not all tokens are equal and some tokens are hard to learn in the AR pretraining stage, implying the introduced subgoal imbalance phenomenon also exists in the general text corpus.
reverse training (Lee et al., 2023; Golovneva et al., (024).fil-in-the-miide
5.2 NON-AUTOREGGRESSIVE MODELING
The non-autoregressive (NAR) generation method, which produces all target tokens simultancously given the source context, is first proposed by (Gu et al., 2017) in the text field for machine transla- tion, primarily due to the efficiency consideration. While series of advancements have been made afterward (Lee et al., 2018; Gu et al., 2019; Ghazvininejad et al., 2019; Qian et al., 2021; Huang et al., 2022, inter alia). traditional NAR models still tend to underperform AR models in terms of generation quality (Xiao et al., 2023). Diffusion models (Sohl-Diekstein et al., 2015; Ho et al., 2020), powerful class of generative models known for their impressive ities (Dhariwal & Nichol, 2021). have recently been applied the field of text generation (Hooge- boom et al., 2021; Austin et al., 2021; Li et al., 2022; Camphell et al., 2022; Dieleman et al., 2022; Chen et al., 2023; Ye al., 2023b; Lovelace et al., 2024), reinforcement learning (Janner et al., 2022; Chi et al., 2023) and protein design (Xu et al., 2022; Hoogehoom et al., 2022b; Corso et al., 2023). In essence, diffusion models perform a multi-step denoising process to progressively convert a random noise into a 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 et al., 2006). For text, the diffusion model can be seen as an extension of the traditional iterative NAR models (Gong et al., 2022) and has been shown to approach or outperform AR models text perplexity (Han et al., 2023; Lou et al., 2023; Gulra- jani & Hashimoto, 2024), diversity (Gong et al., 2022; 2023; Zhang et al., 2023) as well as various seq-to-seq tasks (Wu et al.. 2023b; Zheng et al., 2023; Ye et al., 2024). In this work. we compare diffusion with AR from a perspective of subgoal imbalance and demonstrates the effectiveness of diffusion in tasks requiring complex reasoning and planning.
soage--eneration capabil-
CONCLUSION
This paper presents extensive analysis of the limitations of auto-regressive (AR) language models when applied to planning tasks that involve deliberate planning, both in controlled settings and real- world contexts. Based on an advanced understanding, we propose an improved diffusion model, named MGDM, that performs significantly better than AR and previous diffusion models on vari- sophisticated planning tasks. Our findings underscore the necessity reevaluate the sequence modeling paradigm for modern large language models, especially in tackling challenging problem- solving tasks.
10

---

Published as a conference paper at ICLR 2025
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.KKU714/21.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.086 087774, 2023.
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B Tenenbaum, Tommi S Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision-making? In NeurlPS 2022 Founda- tion Models for Decision Making Workshop, 2022.
Anthropic. Introducing Claude, 2023.
URL
Introducing-cl
Jacob Austin, Daniel D. Johnson, Jonathan Ho, Daniel Tarlow, and Rianne 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, NeurlPS 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.06663, 2024.
Jinze Bai, Shuai Bai, Yunfei Chu, Zeyu Cui, Kai Dang, Xiaodong Deng. Yang Fan, Wenbin Ge, Yu Han, Fei Huang, et al. Qwen technical report. arXiv preprint arXiv:2309.16699 2023.
Mohammad Bavarian, Heewoo Jun, Nikolas Tezak, John Schulman, Christine McLeavey, Jerry Tworek, and Mark Chen. Efficient training of language models to fill in the middle. arXiv preprint arXiv:2207. 14255, 2022.
Yoshua Bengio, Réjean Ducharme, and Pascal Vincent. A neural probabilistic language model. Advances in neural information processing systems, 13, 2000.
Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Michal Podstawski, Lukas Gian- inazzi, Joanna Gajda, Tomasz Lehmann, Hubert Niewiadomski, Piotr Nyczyk, et al. Graph of thoughts: Solving elaborate problems with large language models. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 38, pp. 17682-17690, 2024.
Tim Brooks, Bill Peebles, Connor Holmes, Will DePue, Yufei Guo, Li Jing, David Schnurr, Joc Taylor, Troy Luhman, Eric Luhman, Clarence Ng. Ricky Wang, and Aditya Ramesh. Video generation models as world simulators. 2024. URL pettps::/opeaaate ctt. com/research/ www.ldddmmmlator...mmmmm
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Sébastien Bubeck, Varun Chandrasekaran, Ronen Eldan, Johannes Gehrke, Eric Horvitz, Ece Ka- mar, Peter Lee, Yin Tat Lee, Yuanzhi Li, Scott Lundberg, intelligence: Early experiments with gpt-4. arXiv preprint 23.. 2023.
Sparks of artificial general
Andrew Campbell, Joe Benton, Valentin De Bortoli, Thomas Rainforth, George Deligiannidis, and Arnaud Doucet. A continuous time framework for discrete denoising models. Advances Neural Information Processing Systems, 15:28266-28277 2022.
Ting Chen, Ruixiang ZHANG, and Geoffrey Hinton. Analog bits: Generating discrete data using diffusion models with self celf-condiiinnnn. In The Eleventh International Conference on Learning Representations, 2023.
11

---

Published as a conference paper at ICLR 2025
Cheng Chi, Styuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shu- ran Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint 0rXiv:23033.4137,, 2023.
Stephen A Cook. The complexity of theorem-proving procedures, 1971.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi S Jaakkola. Diffdock: Diffusion steps, twists, and turns for molecular docking. In The Eleventh International Conference on Learning Representations, 2023.
Countdown. Countdown (game show), 2024. URL https: httpp:://en...iipeaaaa..oo/////p Countdoww_(ggame (game __ome____ww.....
James M Crawford and Larry D Auton. Experimental results on the crossover point in random 3-sat. Artificial intelligence, 81(1-2):31-57, 1996.
Jacob Devlin, Ming-Wei 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 I (Long and Shori Pupers), pp. 4171-41 1416. Minneapolis, Minnesota, June 2019. Association for Com- putational Linguistics. doi: 10. 18533/ URL https: //clclllcceeooe h/////aaaattooooo....oooo orgg/ N1 919-144 91133.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis, Advances in neural information processing systems,
2021.
Sander Dieleman, Laurent Sartran, Arman Roshannai, Nikolay Savinov, Yaroslav Ganin, Pierre H Richemond, Arnaud Doucet, Robin Strudel, Chris Dyer, Conor Durkan, et al. Continuous diffu- sion for categorical data. arXiv preprint arXiv:221 15089, 2022.
Jian Ding, Allan Sly, and Nike Sun. Proof satisfiability conjecture for large k. In Proceedings of the seventh annual ACM symposium Theory of computing, pp. 59 68, 2015.
Li Du, Hongyuan Mei, and Jason Eisner. Autoregressive modeling with lookahead attention. arXiv preprint arXiv:2305.12223 2023.
Nouha Dziri, Ximing Lu, Melanie Selar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Sean Welleck, Peter West, Chandra Bhagavatula, Ronan Le Bras, et al. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 36, 2024.
Kanishk Gandhi, Denise Lee, Gabriel Grand, Muxin Liu, Winson Cheng, Archit Sharma, and Noah D Goodman. Stream of search (sos): Learning to search in language. arXiv preprint arXiv:2404.03683, 2024.
Howard Garns. Sudoku, 1979. URL
Marjan Ghazvininejad, Omer Levy, Yinhan Liu, and Luke Zettlemoyer. Mask-predict: Parallel decoding of conditional masked language models. arXiv preprint arxiv:1904.09324, 2019,
Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Syn- naeve. Better & faster large language models via multi-token prediction. arXiv preprint arXiv:2404. 19737, 2024.
Olga Golovneva, Zeyuan Allen-Zhu, Jason Weston, and Sainbayar Sukhbaatar. Reverse training to nurse the reversal curse. arXiv preprint arXiv:2403. 2...::2403.339999 13799, 2024.
Carla P Gomes, Henry Kautz, Ashish Sabharwal, and Bart Selman. Satisfiability solvers. Founda- tions of Artificial Intelligence, 3:89-134, 2008.
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong. Diffuseq: Sequence to sequence text generation with diffusion models. The Eleventh International Conference on Learning Representations, 2022.
12

---

Published as a conference paper at ICLR 2025
Shansan Gong, Mukai Li, Jiangtao Feng, Zhiyong Wu, and Lingpeng Kong. Diffu v2: Bridging discrete and continuous text spaces for accelerated Seq2Seq diffusion models. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. 9868-9875 Singapore, December 2023. Association for Computational Lin- guisties. doi: 10. http:////22.2...nnnneneeeeem 2023. www.finningg. emnip.
URL wttttssss/ccaaaaeeeeeteeeeeeeeae /tt////caaaathhoaa/oo/o ogy .org/
Jiatao Gu, James Bradbury, Caiming Xiong, Victor OK Li, and Richard Socher. Non-autoregrressive neural machine translation. arXiv preprini arXiv:/711.0000 1.02281, 2017.
Jiatao Gu, Changhan Wang, and Junbo Zhao. Levenshtein transformer. Advances in neural infor- mation processing systems, 32, 2019.
Ishaan Gulrajani and Tatsunori B Hashimoto. Likelihood diffusion language models. Ad- vances in Neural Information Processing Systems, 36, 2024.
Xiaochuang Han, Sachin Kumar, and Yulia Tsvetkov. Ssd-Im: Semi-autoregressive simplex-based diffusion language model for text generation and modular control. In Proceedings of the 61 Annual Meeting of the Association for Computational Linguistics (Volume Long Papers), pp. 11575-11596, 2023.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell,
Balcan, and Hsuan-Tien Lin (eds.). Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633- 8646, 2022.
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 Dauphin, Perey Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Process- ing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual, pp. 12454-12465, 2021.
Emiel Hoogeboom, Alexey Gritsenko, Jasmijn Bastings, Ben Poole, Rianne van den Berg, and Tim Salimans. Autoregressive diffusion models. In International Conference on Learning Rep- resentations, 2022a.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference machine learning, pp. 8867- 8887. PMLR, 2022b.
Edward Hu, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al. Lora: Low-rank adaptation of large language models. In International Conference Learning Representations, 2021.
Fei Huang, Hao Zhou, Yang Liu, Hang Li, and Minlie Huang. Directed acyclic transformer for non-autoregressive machine translation. In International Conference on Machine Learning, pp- 9410-9428. PMLR, 2022.
Aapo Hyvärinen and Peter Dayan, Estimation of non non-noomalized statistical models by score match- ing. Journal of Machine Learning Research, 6(4), 2005.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.0999), 2022.
Frederick Jelinck. Interpolated estimation of markov source parameters from sparse data. In Proc. Workshop Pattern Recognition in Practice, 1980, 1980.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, et al. Mistral 7b. arXiv preprint arXiv:2310006825, 2023.
13

---

Published as a conference paper at ICLR 2025
Daniel D Johnson, Jacob Austin, Rianne den Berg, and Daniel Tarlow. Beyond in in-place corrup- tion: Insertion and deletion in denoising probabilistic models. In ICML Workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models, 2022.
Subbarao Kambhampati, Karthik Valmeekam, Lin Guan, Kaya Stechly, Mudit Verma, Siddhant Bhambri, Lucas Saldyt, and Anil Murthy, Llms can't plan, but can help planning in Ilm-modulo frameworks. arXiv preprint 27Xiv:2402.01817, 2024.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.0836. 2020.
Diederik Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational diffusion models. Ad- vances in neural information processing systems, 34:21696-21707, 2021.
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, and Fujie Huang. A tutorial on energy-based learning. Predicting structured data, 1(0), 2006.
Jason Lee, Elman Mansimov, and Kyunghyun Cho. Deterministic non-autoregressive neural se-
quence modeling by iterative refinement. arXiv preprint arXiv:1802.06901, 2018.
Nayoung Lee, Kartik Sreenivasan, Jason D Lee, Kangwook Lee, and Dimitris Papailiopoulos.
Teaching arithmetic small transformers. arXiv preprint arXiv:2307.03350 2023.
Lucas Lehnert, Sainbayar Sukhbaatar, Paul Movay, Michael Rabbat, and Yuandong Tian. Be- yond a=: Better planning with transformers via search dynamics bootstrapping. arXiv preprint arXiv:2402. 14083, 2024.
Xiang Lisa Li, John Thickstun, Ishaan Guirajani, Percy Liang, and Tatsunori B Hashimoto. Diffusion-Im improves controllable text generation. In Conference on Neural Information Pro- cessing Systems, NeurIPS, 2022.
Chu-Cheng Lin, Aaron Jacch, Xin Li, Matthew R Gormley, and Jason Eisner. Limitations of au- toregressive models and their alternatives. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, pp. 5147-5173, 2021.
Zhenghao Lin, Zhibin Gou, Yeyun Gong, Xiao Liu, Yelong Shen, Ruochen Xu, Chen Lin, Yujiu Yang. Jian Jiao, Nan Duan, et al. Rho-l: Not all tokens are what you need. arXiv preprint arXiv:2404.000 07965, 2024.
Aaron Lou, Chenlin Meng, and Stefano Ermon. Discrete diffusion language modeling by estimating the ratios of the data distribution. ArXiv preprint, abs/2310. 2023.
Justin Lovelace, Varsha Kishore, Chao Wan, Eliot Shekhtman, and Kilian Q Weinberger. Latent diffusion for language generation. Advances in Neural Information Processing Systems, 36, 2024.
Ximing Lu, Sean Welleck, Peter West, Liwei Jiang, Jungo Kasai, Daniel Khashabi, Ronan Le Bras, Lianhui Qin, Youngjae Yu, Rowan Zellers, et al. Neurologic a esque decoding: Constrained text generation with lookahead heuristics. In Proceedings of the 2022 Conference of the North Amer- ican Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 780-799 2022.
Eran Malach. Auto-regressive next-token predictors are universal learners. arXiv:2309.06977 2023.
arXiv preprint
William Merrill and Ashish Sabharwal. The expresssive power of transformers with chain of thought. arXiv preprint arXiv:2310.07923 2023.
Giovanni Monea, Armand Joulin, and Edouard Grave. Pass: Parallel speculative sampling. arXiv preprint arXiv:2311.1358/ 2023.
OpenAI. Introducing ChatGPT, 2022. URL https:ttt://ppeee http:///aa..a..cccoooooo//////// httpa/aa.....abolooc//at.t.....
14

---

Published as a conference paper at ICLR 2025
Kyubyong Park. 1 million sudoku games, 2016. URL https://www.tt www.l//..ww........./.. com/
bryanpark/ sudoku.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the JEEE/CVF International Conference on Computer Vision, pp.
2023.
Weizhen Qi, Yu Yan, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang, and Ming Zhou, Prophetnet: Predicting future n-gram for sequence-to-sequence pre-training. arXiv preprint ar XXh1:2001004063 2020.
Lihua Qian, Hao Zhou, Yu Bao, Mingxuan Wang, Lin Qiu, Weinan Zhang, Yong Yu, and Lei Glancing transformer for con-euttregressive neural machine translation. In Proceedings of the 59th Anmual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference Natural Language Processing (Volume 1: Long Papers), pp. 1993-2003 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9 2019.
Machel Reid, Vincent J Hellendoom, and Graham Neubig. Diffuser: Discrete diffusion via edit- based reconstruction. arXiv preprint arXiv:2210. 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer: High- resolution image synthesis with latent diffusion models. In Proceedings the the IEEE/CVF confer- computer vision and pattern recognition, pp. 10684-10695, 2022.
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 S.KIv:2406.075244 2024.
Nikolay Savinov, Junyoung Chung, Mikolaj Binkowski, Erich Elsen, and Aaron van den Oord. Step- unrolled denoising autocncoders for text generation. In International Conference on Learning Representations, 2021.
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.04333 2024.
Noah Shinn, Federico Cassano, Ashwin Gopinath, Karthik Narasimhan, and Shunyu Yao. Reflexion: Language agents with verbal reinforcement learning. Advances in Neural Information Processing Systems, 36, 2024.
Jascha
Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-22 2265.22652222 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' A'Aché-Buc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 32: Annual Conference Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 11895-11907, 2019.
Haoran Sun, Lijun Yu, Bo Dai, Dale Schuurmans, and Hanjun Dai. Score-based continuous-time discrete diffusion models. In The Eleventh International Conference on Learning Representa- fions, 2023.
Ilya Sutskever, James Martens, and Geoffrey E Hinton. Generating text with recurrent neural net- works. In Proceedings of the 28th international conference machine learning (ICML-11), pp- 1017-1024, 2011.
Gemini Team, Rohan Anil, Sebastian Borgeaud. Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint 23.. 2023.
15

---

Published as a conference paper at ICLR 2025
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXhv:2302.13971, 2023.
Trieu H Trinh, Yuhuai Wu, Quoe V Le, He He, and Thang Luong. Solving olympiad geometry without human demonstrations. Nature, 625(7995):476- 482, 2024.
Karthik Valmeekam, Matthew Marquez. Sarath Sreedharan, and Subbarao Kambhampati. On the planning abilities of large language models-a critical investigation. Advances in Neural Informa- fion Processing Systems, 36.75993-76005, 2023.
Karthik Valmeekam, Matthew Marquez. Alberto Olmo, Sarath Sreedharan, and Subbarao Kamb- hampati. Planbench: An extensible benchmark for evaluating large language models on planning and reasoning about change. Advances in Neural Information Processing Systems, 36, 2024.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc V Le, Ed H Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Self-consisteney improves chain of thought reasoning in language models. In The Eleventh International Conference on Learning Representations, 2023.
Jason Wei, Tay, Rishi Bommasani, Colin Raffel, Barret Zoph, Sebastian Borgeaud, Dani Yo- gatama, Maarten Bosma, Denny Zhou, Donald Metzler, al. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022a.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Sfain--f--thoggh prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824 24837, 2022b.
Jay Zhangjie Wu, Yixiao Ge, Xintao Wang. Stan Weixian Lei, Yuchao Gu, Yufci Shi, Wynne Hsu, Ying Shan, Xiaohu Qie, and Mike Zheng Shou. Tune-a-video: One-shot tuning of image diffusion models for text-to-video generation. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 7623-7633, 2023a.
Tong Wu, Zhihao Fan, Xiao Liu, Hai-Tao Zheng, Yeyun Gong, Jian Jiao, Juntao Li, Jian Guo, Nan Duan, Weizhu Chen, et al. Ar-diffusion: Auto-regressive diffusion model for text generation. Advances in Neural Information Processing Systems, 36:39957-39974, 2023b.
Yuhuai Wu, Albert Qiaochu Jiang, Wenda Li, Markus Rabe, Charles Staats, Mateja Jamnik, and Christian Szegedy. Autoformalization with large language models. Advances in Neural Informa- tion Processing Systems, 35:32353-32366 2022.
Zhiyong Wu, Yaoxiang Wang, Jiacheng Ye, and Lingpeng Kong. Self-adaptive in-context learning: An information compression perspective for in-context example selection and ordering. In The slst Annual Meeting Of The Association For Computational Linguistics, 2023c.
Yisheng Xiao, Lijun Wu, Junliang Guo, Juntao Li, Min Zhang, Tao Qin, and Tie-yan Liu. A survey on non-eutorrereessvve generation for neural machine translation and beyond. IEEE Transactions on Paitern Analysis and Machine Intelligence, 2023.
Yuxi Xie, Kenji Kawaguchi, Yiran Zhao, Xu Zhao, Min-Yer Kan, Junxian He, and Qizhe Xie. Decomposition enhances reasoning via self-evaluation guided decoding. arXiv preprint arXiv:2305.00633 2023.
Chang Xu, Dacheng Tao, and Chao Xu. A survey on multi-view learning. arXiv preprint arXiv:1304.5634 2013.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geo- metric diffusion model for molecular conformation generation. In International Conference on Learning Representations, 2022.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime Carbonell, Russ R Salakhutdinov, and Quoe V Le. Xlnet: Generalized autoregressive pretraining for language understanding. In H. Wal- lach, H. Larochelle, A. Beygelzimer, F. d'Alché-Buc, Buch E. Fox, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., nww.iiss..........e. ec wc/pppprr...... httppppfieeeppppeeeeppp /tt.es/pppee/e/1s/////// fi httee//cccceessseeee3ee6eeeeeeee
2019. URL
.ddf.
16

---

Published as a conference paper at ICLR 2025
Shunyu Yao, Jeffrey Zhao, Dian Yu, Nan Du, Izhak Shafran, Karthik R Narasimhan, and Yuan Cao. React: Synergizing reasoning and acting in language models. In The Eleventh International Conference on Learning Representations, 2023.
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, Zhiyong Wu, Jiangtao Feng, Tao Yu, and Lingpeng Kong. Compositional exemplars for in-context learning. In International Conference on Machine Learning, pp. 39818-39833. PMLR, 2023a.
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 preprini 4rXiv:2402.07754. 2024.
Jiasheng Ye, Zaixiang Zheng, Yu Bao, Lihua Qian, and Mingxuan Wang. Dinoiser: Diffused condi- tional sequence learning by manipulating noises. arXiv preprint arXiv:2302./0025, 2023b.
Yizhe Zhang, Jiatao Gu, Zhuofeng Wu, Shuangfei Zhai, Joshua Susskind, and Navdeep Jaitly. PLANNER: Generating diversified paragraph via latent language diffusion model. In Thirty- seventh Conference on Neural Information Processing Systems, 2023.
Xueliang Zhao, Wenda Li, and Lingpeng Kong. Decomposing the enigma: Subgoal-based demon-
stration learning for formal theorem proving. arXiv preprint arXiv:2305.16366. 2023.
Lin Zheng, Jianbo Yuan, Lei Yu, and Lingpeng Kong. A reparameterized discrete diffusion model for text generation. ArXiv preprint, abs/2302.057377, 2023.
17

---

Published as a conference paper at ICLR 2025
A
MORE BACKGROUND AND DERIVATION OF DISCRETE DIFFUSION MODELS
A.1 BACKGROUND
Discrete diffusion probabilistic models are first introduced in Sohl-Dickstein et al. (2015) for bi- nary data, and later extended to categorical data in (Hoogeboom et al., 2021). Austin et al. (2021) provides a general form of discrete diffusion and introduces multiple transition matrices, includ- ing an absorbing variant that draws close connections to masked language models (Devlin et al., 2019). Several subsequent works push this line of research further from various aspects, such as incorporating editing-baste based operations (Johnson al., 2022; Reid al., 2022), casting permuted language models (Yang et al., 2019) as diffusion models (Hoogeboom et al., 2022a), developing continuous-time framework (Campbell et al., 2022), parameterizing the routing mechanism (Zheng et al., 2023), and investigating score functions for learning the reverse process Sun et al. (2023); Lou et al. (2023).
A.2 DERIVATION SETUP
We now provide a detailed derivation of the loss in Equation (6). For clear illustration, we initiate derivation with a single random variable 2g and ultimately link it with the multi-variable sequence Xo- Suppose Ze (Za) is a discrete random variable with K possible categories and represented as one-hot vector. The forward process q( (20) N g(ar, (a:-1) corrupts the original data To into a sequence of increasingly noisy latent variables 1:7 := æ; ........ ET. The learned backward process po (#For) = p( p(xr) I 1 (ar l gradually denoises the latent variables to the data dis- tribution. In discrete diffusion, both the forward and backward distribution are defined as categorical distribution, e.g., (0 (a Cat(#;;; p Q, a1(-1) and pe ( g 1 | f (a 0)), where Qr is a pre-defined transition matrix of size K x K (Hoogeboom et al., 2021; Austin et al., 2021).
A.3 THE MARGINAL AND POSTERIOR
Starting from To, we obtain the following t-step marginal and posterior at time t - 1:
(<x1|=0)) = Cat (xrip= = Q, a00.
with Q, = -Q1Q2... Qu
((1t-11/22 2o)= (0 1, ro (2)))((())))) (=) 1
Q. @ QQ,1
Cat Tt-1:p
(9)
x. Q 3g
where q( (2, mn) g( 11) due the Markov property of the forward process. The KL divergence between y and pa can 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 all t depending categories. Therefore, two commonly used forms of Q are introduced by Hoogeboom et al. (2021) and Austin et al. (2021), which ensures Q, can still be computed efficiently, allowing the framework to scale to a larger number of categories.
the choice Qt. may be prohibitive for large T and number of
A.4 TRANSITION MATRIX
Austin et al. (2021) introduced multiple types of the transition matrix Q., such as uniform (Hooge- boom et al., 2021), absorbing, discretized Gaussian and token embedding distance. The absorbing noise for discrete diffusion has been demonstrated to outperform the others (Austin et al., 2021), where the transition matrix is given by
if = j = m
Qlyy = 1 1
3
if i = } + m
Br
if == m,,mm,t + m
The transition matrix can also be written as (1 3,)I + 8,1e where cm is a vector with a one on the absorbing state and zeros elsewhere. Since m is absorbing state, the corruption process
18

---

Published as a conference paper at ICLR 2025
converges not to uniform distribution but the point-mas, distribution matrices Q Q2.Q..... ... Q: can be computed in closed form. Specifically, we transition to another token with probability 3g and stay the same with probability only operative quantity is the probability of not yet having transitioned to another token, given by Or (1 3:). Therefore, we have Q a. + (1 ax)lem
The transition
3 in each step. After steps, the
A.5 DERIVATION OF ELBO
In order to optimize the generative model pe to fit the data distribution g( we typically minimize a variational upper bound the negative log-likelihood, defined below:
log pø (xo)
= log
pe (o T1....
xT )d ... dxr
po(20: ,.,.....
FT
log
(111, 2T x 1o) dar
((x1. ((......,
Zr [D))
po (00 01 21..... . ØT
log ze (z0)
((00) .... TT (t))
po (20, 21 ................................ (TT)
S (ei.......rr ør 2r/0))))))) (co)
log
a (1...... 2T 170)
log po 100 x poo xrr)
1=2 Po [x (2.)
q p (r
1= g 1 D4, vo)
D1. wo q
( (2)))
= - -E .r. (z(o))
log ( x1) log
log
[ xx,)
po (xTT)
1=2
Eg log (oo
2p)
DKLLO(0T xo)llp(2T]
(const)
gge (20) log po (2z
(ao))
DKI. |x,, wo) [2:)]] + (constt)..
A.6 DERIVATION FOR EQUATION (6)
The categorical distribution ((0; 2o) based on Equation (9) is given as:
q( (x,- Ft. 20
Q 3 Q
x Q Xo
[1 B3) + Bro 1] 3 [a1-120 + (1 1--aa-lemm
a,30 xo (1 (1) em
(1 3,)00 0
- -0- (1-
)x, Gem + +,,00-10 (000000000000000000000000-00)))1 10em
0230 30 +
04 )x, 6m
(1 81)a-10; + (1 3x) (1- a 1) 0 3t + B:01- - -1 em
0, xy
02)
where Ow. = Cmlu xx) represents the probability of noise drawn from em being equal to Tt-
Note 2; @ 20 = 0 if x, * xo otherwise 1. Thus the computation of g( xo) breaks down into two cases:
g(0, [0, v00) = [n,,, + (1 T7) cm: (xeo + (1 - Xx) (m))))
x = xPn
if x: + Co.
19

---

Published as a conference paper at ICLR 2025
where 7/t Be( - -------eemm 13 (u=mt)
+x++11 (w++(aa Jen (u=z:)) de 0111-a and em em))) (1 -3)ax 3.0m denotes 1-a;
noise distribution that interpolates between æy and cm-
Recall the distribution pe (2t-1 (x;) is parameterized by ((0=11 ( 0)), the KL divergence between g ( (e, 2o) and po (o becomes 0 when 2t 2p. In the case of absorbing diffusion, = if 2; * Zo and (m)) = Cm. ( 1/;; To) has probability X on index Zo and A on the absorbing state. The model f 0) has pro-probbbblityy on the absorbing state it never predicts the mask token. Therefore, (0 (x))) also has 1 - X probability on the absorbing state. Putting them together, we derive the KI. divergence as:
DRI.gg(a:-1(01, (0 (2:-1 [x:]] = In a [Ar log
+ (1 Ao) log X
[( 0 2g
- Do log (a:;; 0) + C.
where Laogt is if 24 + z0 otherwise O, and C' is constant. Moreover, given ag by definition and do have:
T
L(xo) = -Bg(2n)
d Ee (2, (m)) te. vn log f(x; 0)
for a single random variable, and
L(xo) Σ Ex(xx.n)
Ae Eyg Ae (x ax.atxo Xn,n log (Xt,ns 0)
for Xo that represents a sequence of random variables xo = (Fo. Do,N), where the ) also represents the reweighting term w(t) in Equation (6).
B ALGORITHMS FOR TRAINING AND INFERENCE
The detailed algorithms for training and inference illustrated in Algorithm 1 and 2. respectively. For conditional training and inference, we split x into xx and freeze the condition part x during training and inference.
Algorithm Training MGDM
Input:
neural network f (:; O). data distribution Pdata (Xo. a custom sequence reweighting term wt)), token reweighting parameters a and timesteps T. Output: model parameters 0. repeat
Draw X0.1:N ~ Pdata (0.11:); Draw & Uniform (I,...,77 T});
Draw X₁ g(x Xx):
for n 1,2,
do
Let wwwx(xx.)))) x1, 0) = x xxo. XT, log [xx; O).:
Let m(x(,n)) = ax1 exp u(xo, n 0);;
end for
Lg - wt) MN v(xe,n )u( (u xx, m:; 0);
Minimize Lo with respect to 0;
until converged
20

---

Published as a conference paper at ICLR 2025
Algorithm 2 Sampling from MGDM
Input: trained network (); mask token id m, timesteps T, temperature Output: generated sample Xo-
for 1,,2... N do Initialize XT, = m:;
end for for t == ,,..,, do
Define indicator C2 Topk (fxx; 0)) with indices in top-t/T values set and others for n ........NN do
Draw No,n ~ Categorical 0)/T); In No,n (1 m
end for end for
Return X0,1:N-
C ADDITIONAL EXPERIMENTAL DETAILS
C.I TASK DETAILS
Table Dataset statistics. Minimal and CD are short for the minimal planning task and Countdown,
respectively.
Minimal
CD3
CD4
CDS
Sudoku
3-SAT 5y
3-SAT 7v
3-SAT 9y
Train Instance
50k
500k
500k
500k
100k
50k
50k
100k
Test Instance
1k
1k
1k
1k
1k
lk
1k
lk
Avg Input Token
11
305
Avg Output Token
21
16
25
Max Input Token
49
12
15
18
81
245
269
305
Max Output Token
23
22
35
52
81
9
13
17
We show the statistics and input-output examples on each dataset in Table and Table 10, respec- tively.
Table 5: Model parameters with varying model size.
Tiny
Base
Medium
Parameters
6M
85M
303M
Num of Layer
3
12
24
Num Head
12
Hidden Dim
384
768
1024
C.2 MGDM IMPLEMENTATION DETAILS
We conduct all the experiments on NVIDIA V100-32G GPUs, and we use 8 GPUs for training and sampling. We mainly consider comparing diffusion and AR models trained from scratch with different model sizes, with arguments for each size listed in Table 5. We use the GPT-2 architecture for both MGDM and AR. We set the learning rate to le-3 for the tiny model and 3e-4 for others, and we set the batch size to 1024 across all the models and tasks. We train MGDM for 1200 epochs on the minimal planning task, 300 epochs on Sudoku, and 600 epochs on other datasets. By default, set the diffusion sampling steps to T 20 for tasks with average output tokens larger than 20, otherwise T 10. We use decoding temperature т 0.5 for all tasks. For all the experiments, have verified the statistical significance by running them multiple times.
21

---

Published as a conference paper at ICLR 2025
C.3 BASELINE IMPLEMENTATION DETAILS
We train the AR model until convergence, and the number of training epochs is set to 200 for the minimal planning task, 300 for SAT, and 40 for others. We keep other parameters, e.g., batch size and learning rate, the same as training MGDM.
For LLaMA (Touvron al., 2023), we use LoRA fine-tuning (Hu al., 2021) with lora rank setting to 16. We use a learning rate of le-4, a batch size of 256, and train for a maximum of 20 epochs to ensure the model has converged. For GPT-4, borrow the numbers from Yao et al. (2024).
For all the diffusion baselines, we use the same transformer architecture as GPT-2 to control the variables. We the training parameters the same as MGDM, e.g., number of training epochs to 600, learning rate to 3e-4, and batch size to 1024. During inference, we set decoding timesteps to 20 for all diffusion models as we didn't observe clear performance improvement when scaling to 1024.
D ADDITIONAL EXPERIMENTS
D.1 TOKEN CONSUMPTION ON GAME OF 24
We show the detailed accuracy and token consumption on the game 24 in Table 6,
Table 6: Detailed accuracy and token consumption on game of 24.
Accuracy
Prompt Tokens
Generate Tokens
GPT:410
7.3
1k
18
GPT-4 CoT
4.0
2.2k
67
GPT-4 CoT-SC
9.0
2.2k
6.7k
GPT-4 ToT
74.0
1.4k
2.5k
GPT2-Seratch
18.8
11
26
MGDM
76.0
11
26
D.2 AR WITH TOKEN REWEIGHTING
We show the accuracy of AR with the same token reweighting mechanism in Equation 8 on the minimal planning task in Table 7. We find that applying token reweighting AR models still cannot solve subgoals with PD larger than I (i.e. with accuracy around 50%), similar to the original AR.
Table 7: Results of AR with token reweighting.
Planning Distance
AR
AR with token reweighting
0
100
100
1
100
100
2
51.1
52.1
3
46.9
51.5
52.0
50.3
5
49.9
51.9
22

---

Published as a conference paper at ICLR 2025
D.3 SCALING BOTH DATA AND MODEL SIZE
As an extension of Table 1, we show the accuracy of AR and MGDM when both data and model size increased in Table 8. We find scaling both data and model size is effective for both AR and MGDM.
Table 8: Results of scaling both model and data size.
AR
MGDM
85M model, 500k data
45.8
91.5
303M model, 500k data
41.3
88.3
303M model, 1M data
53.3
95.6
D.4
CASE STUDY
Table 9:
Example predictions on Countdown 4. For each sub-equation, mark the planning error in red and the calculation error in bold. AR exhibits more calculation errors in the last equation due incorrect planning in the previous steps.
Numbers
Goal
Groundtruth
AR Prediction
MGDM Prediction
64,36,52.42
14
64:52:1236/1223422544
61736222577722222226--11
64-52-12.331122342277514
9,73,99,75
73:73:22.992:18999118-811
99:75::774.144::167441666
75-73-22.992-11999-188883
2.52.20,73
57
52-20-332.722:1167-1166557
2420:4011--5-221444421.2.))
52-20-39..222--667--16633
9.80.4.5
89
9480=889,544118991113
65-61=442-2--2444444-99 9.5=4.1/4=1.804-=311
9:38:=89,.-441188911133
65,2,61,22
96
65:61=4.2"448882488880
65-61=42+-22-244424-9
42,47.9.14
81
97-42=5 48-411:7.2077113441122
47-42-5_1495-7044476-888
41,4,48,20
96
16- 36-21=151150==442255577
420--60,41-40-222442233
21,36,3,42
39
42-16:6,,66-18821111833
42:21221330:-22222222233
We show more prediction cases of the autoregressive model and our model in Table 9.
23

---

Published as a conference paper at ICLR 2025
Table 10: Task details by showing example input and output for each dataset.
Task
Input Example
Output Example
Minimal Planning
2,0000.||11.22202222011122154119
8,2/2.000,11/11,,555.
8.3
Countdown 3
15,44,79,50
44-15-29,7--29=50
Countdown 4
16,33,9
364428-114331133111.144
18=96
Countdown
30,36,82,44431155
44.36=8,,0231122542.
3+2242222500255:550-51
Sudoku
080050060
789251364
460907108
462937158
005000029
135468729
970006500
978316542
000872031
546872931
300049000
321549876
004025003
894725613
010000480
217693485
603100007
653184297
3-SAT 5v
14.511.44-52-444112555555554444
1,2,3,44,5
522--4422-33112221::2333333331
4144551.22.31122551::34444333111
42235912254122211222214444444423
1.4.512.3.4244522244444552221211
3-SAT
2.2.112...77334.122333333...
1,2,3,4,,,6,-
6.122.5022.5.063344461244463666.
2-55223..71122332224411333212
722466122-122333611226--1115-55.
1266111266672222255556666666-66-
2.115-6..22-6-77225744445531111.
4435-913-445/22666122661411
91441124511155-011332223771141
3-SAT 9v
3.4-6111.52-77811:30223334444
1.2.3.4.5.66.77-88
1.911...992277901133525.....44
91233382237224446225556666333333
3.28:455.774466322-992999944
955.991.6677336-114444444444
12.611....5..229944799444444471
9/1,4,8
24


Reviews:
[{"review_id": "i9xj7UQC4j", "reviewer": "Reviewer_w8tZ", "summary": "This paper introduces Multi-granularity Diffusion Modeling (MDM), a method that prioritizes subgoals by difficulty to improve complex reasoning and planning in tasks where autoregressive models fall short. MDM significantly outperforms autoregressive approaches on challenging tasks like Countdown and Sudoku, achieving accuracy rates of 91.5% and 100%, respectively. The results demonstrate the promise of diffusion-based methods for advancing AI in complex language understanding and problem-solving.", "questions": "1. How does subgoal imbalance impact performance in real-world RL tasks? Could you provide concrete examples of this issue in practical applications?\n2. How does the performance of the proposed method compare to autoregressive models in terms of computational efficiency (i.e., training time and inference time)?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The authors address a novel and important problem in reinforcement learning (RL): subgoal imbalance in planning tasks. This issue is particularly relevant for the application of deep learning models in RL.\n- They propose a new method to tackle subgoal imbalance, which appears to be a reasonable approach to addressing this challenge.\n- The authors demonstrate that their method outperforms traditional autoregressive models on complex tasks, including Countdown, Sudoku, and Boolean Satisfiability Problems.", "weaknesses": "- The evaluation is limited to well-designed, controlled tasks, which may restrict the general applicability of the proposed method. It would be interesting to see if this approach is also practical in more realistic RL scenarios.\n- While they extend a diffusion model to handle subgoal imbalance, I wonder if a similar technique could be applied to autoregressive language models.  Extending this work to LLMs could make it more broadly applicable and impactful.\n- The authors provide a strong example of subgoal imbalance in Section 3. However, it would be helpful to understand how frequently this issue occurs in real-world RL settings. Including concrete, real-world examples would enhance the appeal of this work.", "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": "0AXYThGXaH", "reviewer": "Reviewer_SUNs", "summary": "This paper explores a few popular hard/challenging planning-based reasoning tasks such as Countdown and Sudoku, and whether a discrete diffusion model can solve them better than an autoregressive model. I am fairly unfamiliar with diffusion models and unable to judge whether the proposed method (multi-granularity diffusion modeling) is a significant contribution or a very minor one -- judged by Eq-8, it is a form of loss reweighting. The derivation for Eq (6) seems novel, but again, I am not fully sure if the derivation is a trivial reuse of previous work or a fully novel insight. \n\nI have a few questions about several of the paper's claims. If properly addressed, I think this would be a good paper for acceptance.", "questions": "1. I want to understand Proposition 1 better: for a sequence $q(x)$, let's take your toy planning task as an example; if that is my true distribution, then I can just do an MC sample from $q(x)$ (by conditioning on previous states). Why would it not follow an autoregressive pattern? Is Proposition 1 as straightforward as thinking because there is an action policy $\\pi$, therefore, $p_{\\pi_1}(x_n | x_{1:n-1})) \\not = p_{\\pi_2}(x_n | x_{1:n-1}))$? I fail to see why this is the case.\n2. \"The regretful compromise\" -- I fail to see what the compromise is. Can you tell me if this is you reinventing the concept of \"compounding error\" in [1]? If so, I am astonished that we seem to be recycling and relearning the same phenomenon every 5-7 years. \n3. Looking at Eq (8), which is a per-token loss reweighting -- it seems that we can apply per-token loss reweighting, too. Can the authors comment on whether per-token loss reweighting can help AR, too (address the issue discussed in Proposition 1)?\n\n[1] Lamb, Alex M., Anirudh Goyal ALIAS PARTH GOYAL, Ying Zhang, Saizheng Zhang, Aaron C. Courville, and Yoshua Bengio. \"Professor forcing: A new algorithm for training recurrent networks.\" Advances in neural information processing systems 29 (2016).", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "The paper is very well-motivated and well-written.\n\nI especially enjoyed Section 3, where the authors tried to use a simple metric (planning distance) to capture the difference between several methods (AR, Reverse-AR, and Diffusion). This alone is already a novel contribution.\n\nThe experimental result is mind-blowing. The proposed method almost completely solves Countdown (CD) and Game of 24, both are major challenges and are under active investigation. It would have been nice if, as reviewers, we had the time, energy, and resources to reproduce and replicate the results of the paper so that we could validate the claim. The authors have provided the code. A quick glance did not raise immediate red flags for me.", "weaknesses": "I do not consider these weaknesses, but any paper can always improve:\n\nThe datasets are mostly Countdown, Game of 24, and Sudoku. I think they are sufficient as of right now to show proof of concept. It would obviously be nicer if we saw other larger domains (such as Blocksworld [1]). However, the three datasets/domains investigated are the most popular ones and, therefore, sufficient.\n\n[1] Valmeekam, Karthik, Sarath Sreedharan, Matthew Marquez, Alberto Olmo, and Subbarao Kambhampati. \"On the planning abilities of large language models (a critical investigation with a proposed benchmark).\" arXiv preprint arXiv:2302.06706 (2023).", "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": "fiHtS2C9Td", "reviewer": "Reviewer_aoA3", "summary": "This work proposes to use a form of diffusion models as an alternative to auto-regressive models for reasoning and planning problems. The authors illustrate the weaknesses of auto-regressive models with planning by introducing the notion of subgoal imbalanced or the idea that the true data distribution may not follow an autoregressive pattern. Using the sub-goal learning as a motivation, the authors propose multi-granularity diffusion model which uses a re-weighting procedure for training efficiency. Through a variety of planning tasks including Countdown, Sudoku and the boolean satisfiability problem (SAT), the authors demonstrate both the efficiency of the multi-granularity diffusion model compared to auto-regressive models and much better performance on more difficult tasks.", "questions": "In addition to the questions in the weaknesses:\n- Could MDM exhibit the properties of in-context learning that AR models are known to have? Could it be possible once there is a sufficient amount of data?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- Novelty: This work is particularly insightful. Not only does the work show experimentally demonstrate the success of diffusion models, it also provides intuition through the sub-goal problem. As mentioned in the related work, diffusion models have been explored in other domains as alternatives to auto-regressive models but it has not been done so in the planning domain. Given how frequently auto-regressive models are used, using diffusion models successfully for planning is almost by definition novel. \n- Presentation and Clarity:  Related to the point above, by incorporating experiments through  early in the paper and providing intuition, the presentation made it much easier to understand.\n- Contribution: There are several strong contributions. First, is replacing auto-regressive models with diffusion models in general. The second is the introduction of the multi-granularity diffusion model. The third, is showing that diffusion models are more efficient (especially with regard to tokens and on cheaper hardware, V100) for planning tasks.", "weaknesses": "- The main weakness would be on the types of planning tasks that the model is evaluated on. Other works like Tree of Thought evaluate on other planning tasks such as Mini-Crossword (in addition to the game of 24). It would be interesting to see how MDM does on tasks that also require some ‘common sense.’ \n- Scaling of MDM: Part of the advantages of AR models is the ability to scale up due to the ease of generating data with next-token prediction. Could MDMs exhibit such a property? In the results of the CD task, we see that with CD 4 and CD 5, the 85 M does better than the 303 M.", "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": "D1kBsEk5sl", "reviewer": "Reviewer_gxqG", "summary": "The authors begin by pointing out that autoregressive (AR) models struggle with reasoning and long-term planning tasks. The authors propose _discrete diffusion (DD)_ as a better modeling approach for these classes of tasks. \n\nThey begin their contributions by providing a theoretical argument for their assertion from the viewpoint of what they deem \"subgoal imbalance\". Here, they argue that the subgoals faced when attempting reasoning and planning tasks autoregressively are not all necessarily of equal learning difficulty. The authors follow this initial theoretical setup with empirical evidence on a toy synthetic planning task demonstrating that autoregressive models struggle when the subgoals become more difficult, e.g. require longer planning horizons.\n\nThe authors pose the same toy task to a DD approach, which they demonstrate does not show similar signs of failure as the AR approach. Given this, they continue their contributions by theoretically outlining _why_ the DD approach does not struggle as much as the AR model on the same toy task. Here, they compare the losses of the AR approach and DD approach, and point out that the DD loss remains low compared to its AR counterpart. Moreover, they show that the unweighted loss of the DD approach reveals the following trend: as the number of timesteps increases, subgoals at smaller timesteps become easier to learn (face lower loss). The authors analyze this from a \"multi-view learning\" perspective, arguing that each x_t in DD can be interpreted as a different view of x_0. This, they argue, affords more effective learning of challenging subgoals.\n\nGiven these insights, the authors develop an extension to discrete diffusion tailored to planning and reasoning where the \"subgoal imbalance\" problem supposedly emerges. Their extension, dubbed multi-granularity diffusion modeling (MDM) consists in introducing a token-level reweighting mechanism in the Monte-Carlo sampled DD loss equation, designed to adaptively reduce the loss for easy tokens while emphasizing the loss for harder ones. \n\nThe authors then test their method on 3 \"real-world scenarios\", namely on Countdown, Sudoku and SAT. They compare MDM to existing diffusion models as well as a number of AR baselines. Throughout, the authors empirically demonstrate the superiority of DD approaches to these planning and reasoning tasks, as well as the performance edge gained with MDM over existing DD approaches.", "questions": "- I've already hinted this a little bit above, but one of the properties of these very popular AR models is their emergent generality: they can be successfully applied to a variety of downstream tasks while being trained on a simple autoregressive objective. In the paper you have shown that _when training specifically for planning_ discrete diffusion is a better approach than autoregression. My question is -- should we expect this finding to generalize to \"general\" diffusion solutions too? I.e. suppose one has some foundation *diffusion* LM -- do we expect such an LM to still be better at planning than an equivalent AR LM? Or is the assertion/result of the paper only for the narrow case of planning specific training? \n- If indeed this paper is about the narrow case, in which way, if any, would you envision incorporating this narrow case into a more general model capable of more than just planning?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The paper is very well structured. Ideas are presented in a organised, sequential manner with little to no jumping required. Evidence and explanations are provided soon after assertions. The positioning of every section makes sense.\n-  The paper is original in some ways: it contributes this novel concept of \"subgoal imbalance\" and later links \"multi-view learning\" as a possible explanation for diffusion's superiority over autoregressive methods when it comes to planning and reasoning tasks. \n- The main contribution of the paper: theoretical argumentation supported by empirical evidence for why diffusion works better than autoregression for planning and reasoning is a significant contribution considering the value of such capabilities, and may provide a further shift towards diffusion methods in the field.", "weaknesses": "- I found that the definition of \"subgoal imbalance\" was a bit awkward -- is this something that can only be talked about from an autoregressive perspective or is this a general aspect of multi-step tasks such as planning and reasoning? The more formal proposition at line 160 makes it sound like the former, but throughout the text my understanding leans more to the latter. I recommend trying to make this proposition more rigorous, more carefully outlining when it applies and when it does not. \n- I found the \"multi-view learning\" perspective a bit underdeveloped. I think it would help to define what \"view\" means in this context, rather than expecting the reader to refer to Xu et al. 2013. Especially given how well the rest of the paper flows. In general I think this part of the paper (L243-252) could have gone into more depth for readers to get a better understanding of what the authors were intending, because at the moment it reads as a bit vague and non-rigorous.\n- Perhaps because I did not fully grasp the \"multi-view learning\" perspective due to its underdevelopment, it is not 100% clear to me how what is observed in 3.2 led to the insights on which MDM is developed in 3.3. Perhaps it would have been better to drop the \"multi-view learning\" perspective and focus entirely on the subgoal imbalance idea, which I think transitions naturally to the extension presented in 3.3 without much of the argumentation in 3.2 needed. \n- I am confused by the inconsistency of the baselines in section 4. It would have been nice to always use all the baselines, but for some reason e.g. GPT-4 ToT is only used on Game of 24. \n- Similarly, since there are no claims of the generality of diffusion models which is on the other hand often touted as an emergent property of the autoregressive models, it would have been nice to have other non-general, planning-specific baselines to compare against. E.g. the question I am getting at is, many \"real-world\" scenarios could be solved analytically a Planning Domain Definition Language (PDDL), so why should one use MDM (or other diffusion methods) instead?", "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)