Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning

## Abstract

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 (MDM), which prioritizes subgoals based on difficulty during learning. On complex tasks like Countdown, Sudoku, and Boolean Satisfiability Problems, MDM significantly outperforms autoregressive models without using search techniques. For instance, MDM 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.

## 1 Introduction

In recent years, autoregressive language models (LMs; Bengio et al. ) have dominated the landscape of natural language processing and artificial intelligence. Empowered by scaling laws (Kaplan et al., ), these models have demonstrated impressive performance across various applications (OpenAI, ; Achiam et al., ; Anthropic, ; Team et al., , *inter alia*). However, this apparent success masks significant limitations that are becoming increasingly evident. Autoregressive models inherently struggle with tasks requiring complex reasoning, long-term planning, and maintaining global coherence (Bubeck et al., ; Valmeekam et al., ; ; Dziri et al., ; Kambhampati et al., ). These shortcomings represent substantial challenges in developing AI systems capable of robust problem-solving and adaptable cognition (Wu et al., ; Zhao et al., ; Trinh et al., ; Yao et al., ; Shinn et al., , *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 to 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 et al., ; Besta et al., ) and the incorporation of backtracking supervision during training (Lehnert et al., ; Gandhi et al., ). 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 a fundamentally different modeling approach: discrete diffusion models. While most contemporary language models are autoregressive, diffusion-based models have become predominant in image (Dhariwal & Nichol, ; Rombach et al., ; Peebles & Xie, ) and video domains (Ho et al., ; Wu et al., ; Brooks et al., ). Diffusion models are also gaining traction in various other applications, such as protein desiging (Xu et al., ; Hoogeboom et al., ; Corso et al., ) and planning in reinforcement learning (Janner et al., ; Ajay et al., ; Chi et al., ). In this work, we 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 2) to illustrate why autoregressive models struggle with these types of problems, often achieving near-random performance. In contrast, we demonstrate how diffusion models effectively learn the subgoals 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., ). Each of these views is more manageable, resulting in an overall easier and more effective learning process.

Building upon these insights, we propose a natural extension to current discrete diffusion models, which we term multi-granularity diffusion modeling (MDM; §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., ) and Sudoku (Garns, ). These tasks demand extensive planning over a large number of combinations and pose challenges even for commercial Large Language Models (e.g., GPT-4 Achiam et al. ). Notably, without employing any search techniques, MDM achieves 91.5% and 100% accuracy on Countdown and Sudoku respectively, while its autoregressive counterpart only solves 45.8% and 20.7% of the problems. Additionally, we conduct experiments on the Boolean Satisfiability Problem (SAT), an NP-complete problem (Cook, ) that represents a wide range of constraint satisfaction problems. Our model exhibits superior performance in solving SAT problems with higher accuracy compared to the autoregressive alternative, particularly when dealing with an increased number of variables and constraints. Through this systematic exploration, we aim to demonstrate the potential advantages of diffusion-based approaches in addressing sophisticated language understanding and generation challenges. All associated code is available at [https://github.com/HKUNLP/diffusion-vs-ar](https://github.com/HKUNLP/diffusion-vs-ar).

## 2 Background

### 2.1 Auto-regressive Modeling

Let ${\mathbf{x}}\coloneq({\bm{x}}_{1},\dots,{\bm{x}}_{N})$ denote a sequence drawn from a data distribution $q({\mathbf{x}})$. For decades, it has been common to factorize the joint probabilities of a sequence of tokens as the product of conditional probabilities (Jelinek, ; Bengio et al., ):

$$ $p_{\bm{\theta}}({\mathbf{x}})=p_{\bm{\theta}}({\bm{x}}_{1})\prod_{n=2}^{N}p_{ \bm{\theta}}({\bm{x}}_{n}\mid{\bm{x}}_{1:n-1}),$ (1) $$

where ${\bm{\theta}}$ parameterizes the model distribution and ${\bm{x}}_{1:n-1}\coloneqq{\bm{x}}_{1},\dots,{\bm{x}}_{n-1}$. In order to optimize the generative model $p_{\bm{\theta}}({\mathbf{x}})$ to fit the data distribution $q({\mathbf{x}})$, we
optimize the negative log-likelihood:

$$ $L_{\text{AR}}=-\mathbb{E}_{q({\mathbf{x}})}\log p_{\bm{\theta}}({\mathbf{x}})= -\mathbb{E}_{q({\mathbf{x}})}\sum_{n=1}^{N}\log p_{\bm{\theta}}({\bm{x}}_{n} \mid{\bm{x}}_{1:n-1}).$ (2) $$

### 2.2 Discrete Diffusion Modeling

Discrete diffusion models (Sohl-Dickstein et al., ; Hoogeboom et al., ; Austin et al., ) are a class of latent variable models characterized by a forward noising process and a learned reverse denoising process. The forward process $q({\mathbf{x}}_{1:T}|{\mathbf{x}}_{0})=\prod_{t=1}^{T}q({\mathbf{x}}_{t}|{
\mathbf{x}}_{t-1})$ corrupts the original data ${\mathbf{x}}_{0}\coloneqq{\mathbf{x}}$ into a sequence of increasingly noisy latent variables ${\mathbf{x}}_{1:T}\coloneqq{\mathbf{x}}_{1},\dots,{\mathbf{x}}_{T}$. The backward process learns to gradually denoise the latent variables
to the data distribution given by:

$$ $p_{{\bm{\theta}}}({\mathbf{x}})=\sum_{{\mathbf{x}}_{1:T}\sim q}p({\mathbf{x}}_ {T})\prod_{t=1}^{T}p_{{\bm{\theta}}}({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t}).$ (3) $$

Due to the intractable marginalization, we typically
optimize a variational upper bound on the negative log-likelihood:

$$ $\displaystyle L_{\mathrm{DM}}=\mathbb{E}_{q({\mathbf{x}}_{0})}\bigg{[}$ $\displaystyle\underbrace{D_{\mathrm{KL}}[q({\mathbf{x}}_{T}|{\mathbf{x}}_{0})| |p({\mathbf{x}}_{T})]}_{L_{T}}+\sum_{t=2}^{T}\underbrace{\mathbb{E}_{q({ \mathbf{x}}_{t}|{\mathbf{x}}_{0})}\big{[}D_{\mathrm{KL}}[q({\mathbf{x}}_{t-1}| {\mathbf{x}}_{t},{\mathbf{x}}_{0})||p_{{\bm{\theta}}}({\mathbf{x}}_{t-1}|{ \mathbf{x}}_{t})]\big{]}}_{L_{t-1}}$ $\displaystyle\underbrace{-\mathbb{E}_{q({\mathbf{x}}_{1}|{\mathbf{x}}_{0})}[ \log p_{{\bm{\theta}}}({\mathbf{x}}_{0}|{\mathbf{x}}_{1})]}_{L_{0}}\bigg{]},$ (4) $$

where $L_{T}$ is a constant when one uses a fixed prior $p({\mathbf{x}}_{T})$.
By defining both the forward and backward distribution as categorical distribution, e.g., $q({\mathbf{x}}_{t}|{\mathbf{x}}_{t-1})=\mathrm{Cat}({\mathbf{x}}_{t};{\bm{p}}=
{\bm{Q}}_{t}^{\top}{\mathbf{x}}_{t-1})$ where ${\bm{Q}}_{t}$ is a pre-defined $K\times K$ transition matrix and $K$ is the size of categories, and $p_{{\bm{\theta}}}({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})=q({\mathbf{x}}_{t-1}|{
\mathbf{x}}_{t},f({\mathbf{x}}_{t};{\bm{\theta}}))$, the forward process posterior $q({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})$ and each KL term can be calculated analytically (Hoogeboom et al., ; Austin et al., ).

## 3 Subgoal Imbalance and Multi-granularity Diffusion Models

In this section, we employ a motivation example (§3.1) to elucidate the challenges faced by autoregressive models in specific scenarios. Through this analysis, we introduce the concept of *subgoal imbalance*—wherein some subgoals are inherently more difficult than others—which offers insights into these difficulties. We then extend our discussion in §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-granularity Diffusion Modeling (MDM; §3.3) as a 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

Figure: Figure 1: The planning task.
Refer to caption: x1.png

We designed a simple planning task to serve as our running example. Consider the example in Figure 2, where the input for the task consists of a set of shuffled edges from the graph shown below. At the end of the input sequence, the start and goal nodes are 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 $9$, the model must plan over a distance of $3$ nodes to determine that the correct next choice should be node $5$ rather than 0 0. We define this span as the Planning Distance (PD), a 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*.

###### Proposition 1

(Subgoal imbalance due to the unknown data distribution $q({\mathbf{x}})$)
The autoregressive modeling decomposes the model distribution $p_{\bm{\theta}}({\mathbf{x}})$ into multiple subgoals: $p_{\bm{\theta}}({\bm{x}}_{1}),\dots,p_{\bm{\theta}}({\bm{x}}_{N}\mid{\bm{x}}_{
1:N-1})$. Given that the true data distribution $q({\mathbf{x}})$ is unknown, the generation of individual tokens may not inherently follow an autoregressive pattern (i.e., ${\bm{x}}_{n}\not\sim q({\bm{x}}_{n}\mid{\bm{x}}_{1:n-1})$). Consequently, the difficulty of learning various subgoals can differ significantly. Given only the left context, some subgoals may require substantially more data to learn or may even be infeasible to learn.

###### Proposition 1

#### Setup.

We synthesize the data with only one distracting path. We randomize node numbers in $[0,10]$ and the intersection positions in $[0,5]$. We further designed this task to be symmetric, ensuring that simply training with reversed output, as suggested by Bachmann & Nagarajan (), cannot solve subgoals with all PDs. For comparison, we include Auto-regressive (AR), reverse AR (Bachmann & Nagarajan, ), and teacherless training (Monea et al., ; Bachmann & Nagarajan, ), 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 $6$M parameters. More details can be found in Appendix §C.

#### Discussion.

We examine the performance of all the models in two scenarios. In the first scenario, we generate a 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 0 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 simply 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 et al., ; Wei et al., ).
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 0 and 1 (or equivalently, 5 and 4 in the reverse setting) with only 10k data points. However, exponentially larger amounts of data are required to address increasingly challenging subgoals. Both teacherless training and diffusion models exhibit a similar U-shaped curve in their performance. This similarity can be attributed to the fact that teacherless training can be conceptualized as a special case of diffusion without an iterative noising and denoising process. In these models, solving edge PDs necessitates slightly more data. We hypothesize that this phenomenon occurs because the distance to other positions is shorter from the middle position (i.e., higher closeness centrality), thus providing the middle position with more nearby tokens to aid in 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 a fixed dataset size of 50,000 instances. Our findings reveal that scaling the original 6 million parameter model to 85 million, 303 million, and 1.5 billion parameters fails to resolve all PDs. Only upon fine-tuning a substantially larger model, specifically the LLaMA 7B model (Touvron et al., ), 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 between two complicated categoricals (Hoogeboom et al., ), we consider discrete diffusion with absorbing state and simplify it as the weighted cross-entropy losses (Austin et al., ; Zheng et al., ; Shi et al., ; Sahoo et al., ):

$$ $D_{\mathrm{KL}}[q({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t},{\mathbf{x}}_{0})||p_{{ \bm{\theta}}}({\mathbf{x}}_{t-1}|{\mathbf{x}}_{t})=-w(t)\sum_{n=1}^{N}\bm{1}_{ \mathrm{{\mathbf{x}}_{t,n}\neq{\mathbf{x}}_{0,n}}}{\mathbf{x}}_{0,n}^{\top} \log f({\mathbf{x}}_{t};{\bm{\theta}})_{n},$ (5) $$

where $w(t)=\frac{\alpha_{t-1}-\alpha_{t}}{1-\alpha_{t}}\in(0,1]$ is a time-dependent reweighting term which places higher weight when $t$ approaching 0.
We then rewrite Equation (4) as:

$$ $L_{\mathrm{DM}}=\mathbb{E}_{q({\mathbf{x}}_{0})}\sum_{n=1}^{N}\underbrace{\sum _{t=1}^{T}w(t)\mathbb{E}_{q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})}u({\mathbf{x}}_ {0},{\mathbf{x}}_{t},n;{\bm{\theta}})}_{-\log p_{\text{DM}}({\bm{x}}_{n}\mid{ \bm{x}}_{\neq n})},$ (6) $$

where $u({\mathbf{x}}_{0},{\mathbf{x}}_{t},n;{\bm{\theta}})\coloneqq-\bm{1}_{\mathrm{
{\mathbf{x}}_{t,n}\neq{\mathbf{x}}_{0,n}}}{\mathbf{x}}_{0,n}^{\top}\log f({
\mathbf{x}}_{t};{\bm{\theta}})_{n}$ is the cross entropy loss on token $n$.

Figure: Figure 3: Loss for a specific hard subgoal, i.e., PD$=$3, in Diffusion and AR modeling. We also show the unweighted loss $u({\mathbf{x}}_{0},{\mathbf{x}}_{t},n;{\bm{\theta}})$ at different timestep $t$ and context ${\mathbf{x}}_{t}$ in diffusion modeling.
Refer to caption: x3.png

We can now systematically compare the losses of autoregressive (AR) and diffusion models (DM), specifically $-\log p_{\text{AR}}({\bm{x}}_{n}\mid{\bm{x}}_{1:n-1})$ and $-\log p_{\text{DM}}({\bm{x}}_{n}\mid{\bm{x}}_{\neq n})$, as expressed 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 diffusion models are depicted using blue and red lines, respectively. The overall loss $-\log p_{\text{DM}}({\bm{x}}_{n}\mid{\bm{x}}_{\neq n})$ in the diffusion model remains relatively low compared to its autoregressive counterpart $-\log p_{\text{AR}}({\bm{x}}_{n}\mid{\bm{x}}_{1:n-1})$, corroborating the superior performance of the diffusion model on these challenging subgoals in our experiments.

Further analysis of the unweighted loss $u({\mathbf{x}}_{0},{\mathbf{x}}_{t},n;{\bm{\theta}})$ in the diffusion model, based on 1,000 samples of ${\mathbf{x}}_{t}\sim q({\mathbf{x}}_{t}|{\mathbf{x}}_{0})$, reveals a clear trend: as the number of timesteps increases, resulting in more noise in ${\mathbf{x}}_{t}$, objectives in smaller timesteps (i.e., recovery from less noisy data) become significantly easier to learn. From a multi-view learning perspective (Xu et al., ), each ${\mathbf{x}}_{t}$ can be interpreted as a distinct view of ${\mathbf{x}}_{0}$. In the diffusion process, by exploring the consistency and complementary properties of different views offered by a diverse range of interrelated objectives $u({\mathbf{x}}_{0},{\mathbf{x}}_{t},n;{\bm{\theta}})$, our findings suggest that objectives challenging to learn in AR models become more effective, promising, and exhibit better generalization in diffusion models.

This phenomenon is particularly evident when examining scenarios where mask noise is applied to positions after the hard token, i.e., ${\mathbf{x}}_{t}={\bm{x}}_{1:n-1}$, where the diffusion model learns the hard subgoal similarly to AR models. We plot this loss as $-\log p_{\text{DM}}({\bm{x}}_{n}\mid{\bm{x}}_{1:n-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-granularity Diffusion Modeling

These observations provide valuable insights, suggesting that due to subgoal imbalance, we should prioritize different subgoals based on their difficulty during the learning process to achieve more effective learning outcomes and faster convergence. Building on this, we propose the multi-granularity diffusion model as a natural extension of the discrete diffusion model.

In practice, to optimize Equation (6), we typically employ Monte Carlo sampling, which results in:

$$ $L_{\mathrm{DM}}=\sum_{n=1}^{N}\sum_{t=1}^{T}w(t)u({\mathbf{x}}_{0},{\mathbf{x} }_{t},n;{\bm{\theta}}).$ (7) $$

For a sequence of length $N$, the probability of sampling the same ${\mathbf{x}}_{t}$ in AR is 1. However, in diffusion, this probability reduces to $1/C_{N-1}^{t(N-1)/T}$ due to the randomness in sampling ${\mathbf{x}}_{t}$, 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 ${\mathbf{x}}_{t}$. However, individual tokens within the sequence, given their imbalanced difficulties, are not properly reweighted.
To address this, we propose multi-granularity diffusion modeling (MDM), which introduces an additional token-level reweighting mechanism to enhance training efficiency:

$$ $L_{\mathrm{\textsc{MDM}}}=\sum_{n=1}^{N}\sum_{t=1}^{T}w(t)v({\mathbf{x}}_{t,n} )u({\mathbf{x}}_{0},{\mathbf{x}}_{t},n;{\bm{\theta}}),$ (8) $$

where $v({\mathbf{x}}_{t,n})=\alpha(1-\exp(-u(\cdot)))^{\beta}$ is the adaptive token-level reweighting term. Setting $\beta>0$ reduces the relative loss for easy tokens while emphasizing harder tokens, and $\alpha$ is used to control the relative reweighting magnitude. For inference, we employ an easy-first TopK decoding strategy, which has demonstrated superior performance compared to the random decoding method used by Austin et al. (). This finding aligns with similar observations documented in prior studies (Savinov et al., ; Zheng et al., ). We provide a detailed derivation and algorithm of the training and inference process in Appendix §A 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 to 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 planning. Detailed experimental setup can be found in Appendix §C.

### 4.1 Countdown

Countdown (Countdown, ) 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., ). The goal of Countdown is to use the given numbers and arithmetic operations ($+-*/$) to obtain a target number. For example, given 4 numbers “97,38,3,17” and a target number “14”, a step-by-step solution is “97-38=59,59-17=42,42/3=14”.

#### Setup.

**Table 1: Results on the Countdown (CD) task with increasing complexity.**
|  | Params | CD 3 | CD 4 | CD 5 |
| --- | --- | --- | --- | --- |
| Autoregressive |  |  |  |  |
| GPT-2 Scratch | 6M | 94.1 | 31.9 | 4.3 |
| 85M | 95.9 | 45.8 | 5.1 |  |
| 303M | 96.4 | 41.3 | 4.5 |  |
| Stream-of-Search | 250M | - | 54.2 | - |
| LLaMA | 7B | 95.7 | 41.1 | 6.7 |
| 13B | 96.5 | 51.1 | 7.4 |  |
| Diffusion |  |  |  |  |
| VDM | 85M | 99.1 | 73.4 | 16.3 |
| D3PM | 85M | 99.4 | 83.1 | 27.6 |
| RDM | 85M | 99.5 | 87.0 | 45.8 |
| MDM (Ours) | 6M | 98.1 | 52.0 | 27.0 |
| 85M | 99.5 | 91.5 | 46.6 |  |
| 303M | 99.9 | 88.3 | 39.0 |  |

We follow  Gandhi et al. () to generate 500k problems with target numbers ranging from 10 to 100 and randomly hold out 10% of the targets for ‘out-of-distribution’ evaluation. We 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., ) have recently been employed to address the limitations of AR, we also compare with such approaches by training on Countdown 4 and evaluating on the same game of 24 test set as Yao et al. ().

#### Baselines.

Our primary comparison involves autoregressive models trained from scratch, employing the GPT-2 architecture (Radford et al., ) 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., ) with sizes 7B and 13B. These models are fine-tuned using the same dataset. In addition, we compare with Stream-of-Search (Gandhi et al., ), which augments the dataset with search trajectory such that the AR model can be taught
to search.
Furthermore, we compare with several existing diffusion models, both continuous models VDM (Kingma et al., ) and discrete models D3PM Austin et al. () and RDM (Zheng et al., ). By default, we use the absorbing noise for discrete diffusion as it significantly outperforms the multinomial one (Austin et al., ; Zheng et al., ).
Finally, we consider in-context learning (Brown et al., ; Wu et al., ; Ye et al., ) based on GPT-4, including vanilla input-output (IO), chain-of-thought (CoT; Wei et al. ), CoT with Self-consistency (CoT-SC; Wang et al. ) and Tree-of-thought (ToT; Yao et al. ). We use 5 in-context examples following  Yao et al. .

#### Results on Countdown.

As shown in Table 1, diffusion-based 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 (Stream-of-search) does provide some benefits, its effectiveness is limited. Importantly, training the entire search trajectory as a 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.

#### Results on Game of 24.

**Table 2: Accuracy and token cost on game of 24.**
|  | Acc. | # Token |
| --- | --- | --- |
| Prompting |  |  |
| GPT-4 IO | 7.3 | x28 |
| GPT-4 CoT | 4.0 | x61 |
| GPT-4 CoT-SC | 9.0 | x241 |
| GPT-4 ToT | 74.0 | x186 |
| Supervised training |  |  |
| GPT-2 Scratch | 18.8 | x1 |
| MDM | 76.0 | x1 |

As shown in Table 4.1, the performance of the GPT-4 with IO, CoT, and CoT-SC prompting methods from  Yao et al. () is unsatisfactory for the given task, with only accuracy below 10%. The introduction of ToT, which incorporates a 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 4.1 with more details in Appendix D.1. ToT consumes 186 times more tokens than MDM, showcasing the ‘internal’ search capability 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 it is challenging for model scaling and decoding strategies to substitute the advantages of modeling paradigm.

### 4.2 Sudoku

Figure: Figure 4: (Left) Accuracy on Sudoku. (Right) Accuracy on boolean satisfiability problem with increasing difficulty.
Refer to caption: x4.png

Sudoku is a classical logic-based number placement puzzle that has gained popularity due to its rigorous intellectual demands.
The goal of Sudoku is to meticulously fill a 9 × 9 grid with numerical digits, ensuring that every column, row, and 3 × 3 subgrid contains all the numbers from 1 to 9.

#### Setup.

We collect one million solved games from  Park () 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 × 9 grid into a sequence of 81 digits, which serves as the model input. To illustrate, an example input appears as “080050060…603100007” (omitted for brevity), while the corresponding output is represented as “789251364…653184297”. During tokenization, we treat each digit as a separate token.

#### 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 a performance of only 32.9 with 13B parameters. In contrast, our model, which has only 6M parameters, is 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, is a foundational problem in computer science that has been rigorously proven to be NP-complete (Cook, ). This challenging combinatorial problem is attractive as a broad range of search problems from
domains such as software verification, test pattern generation, planning, scheduling, and combinatorics can all routinely be solved by reducing to an appropriate SAT problem (Gomes et al., ).
The goal of SAT is to determine whether a 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 example formula with three variables can be $(x_{1}\vee\neg x_{2})\wedge(\neg x_{1}\vee x_{2}\vee x_{3})\wedge\neg x_{1}$ and an corresponding assignment is $x_{1}=0,x_{2}=0,x_{3}=1$.

#### Setup.

Given the number of variables $n$ and clauses $m$, we iteratively sample $k=3$ variables (and their polarities) uniformly at random until $m$ clauses are obtained. To ensure that we get relatively hard instances of SAT, we take advantage of the well-studied family of random $k$-SAT problem (Ding et al., ) and set the $m$ to be close to $m=4.258n+58.26n^{-2/3}$ given $n$, as it has been observed that SAT solvers are slow to determine the satisfiability of a formula when $m$ is near the threshold (Crawford & Auton, ). We consider increasing numbers of variables from {5,7,9} and generate 50k training data for $n=5,7$ and 100k for $n=9$, as well as additional 1k testing data for each $n$.

#### Result.

As shown in the right figure of Figure 4, MDM 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 performance gap between the two models widens as the difficulty of the task increases. This indicates that our diffusion model exhibits a more pronounced advantage in handling more challenging tasks than the AR counterpart.

### 4.4 Analysis

Figure: Table 3: Ablation on training reweighting strategies and inference decoding methods.
Refer to caption: x5.png

#### On the Effect of Training and Decoding Strategies.

As listed in Table 3, we find that changing the sequence-reweighting strategies has only led to a slight improvement in performance. However, when a suitable parameter is selected for token-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., ; Zheng et al., ). We compare the evaluation accuracy along the training process in Figure 5. By aligning the AR training steps with the diffusion process, we can see AR converges rapidly, with the performance tends to plateau afterward. The utilization of our multi-granularity loss, which incorporates sequence and token reweighting, demonstrates superior performance, particularly during the middle stages of training. This implies that the inclusion of such a loss function contributes to enhanced convergence during the training process.

#### On Decoding Speed.

Figure: 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).(b) Ratio of planning (left y-axis) and calculation (right y-axis) error at each reasoning step on Countdown 4.
Refer to caption: x6.png

We assess the trade-off between accuracy and decoding speed by comparing the performance of the AR model (GPT-2 Scratch 85M) and MDM (85M). The speed metric is determined by the number of samples processed using a batch size of 1 on the NVIDIA GeForce V100 GPU.
As shown in Figure 6(a), MDM can flexibly control the trade-off between accuracy and decoding speed by varying the diffusion timesteps.
Notably, by employing just one diffusion step, MDM demonstrates a remarkable 10x improvement in speed compared to AR, while maintaining superior accuracy with 75% and 12.7% compared to 45.8% and 5.1% of AR on Countdown 4 and 5, respectively.
We observed that the slope of countdown 4 is smaller compared to countdown 5 in the trade-off. This suggests that for tasks with lower complexity, diffusion demonstrates a more noticeable speed advantage by setting a smaller diffusion step. In addition, it also indicates that sacrificing some efficiency for performance improvements becomes particularly evident when dealing with more intricate tasks.

#### Error Analysis: The Regretful Compromise.

To gain a deeper understanding of error patterns in AR and MDM, 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 “97-38=59,59-17=42,42/3=14”. We divide the solution into four parts: equation 1, equation 2, equation 3, and answer checking.
First, from a calculation perspective, we assess the error ratio in each equation by comparing the left-hand side and the right-hand 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*4=64,11+64=87” while the correct solution is “16/4=4,40+51=91,91-4=87”.
We can observe that in the last equation, the model realizes that it needs to output 87 as the final result. Therefore, even though the model most likely knows that 11 + 64 actually equals 75 given the low calculation error in the former equations, it reluctantly outputs 87 due to it being the last equation. 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 in Figure 6(b). It is evident that there is a 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.

## 5 Related Work

### 5.1 Autoregressive Modeling

Starting from Bengio et al. () and later Sutskever et al. (), the autoregressive modeling paradigm, where the prediction of a token only depends on the preceding context, is widely adopted in modeling language, until recently (OpenAI, ; Achiam et al., ; Anthropic, ; Team et al., ; Touvron et al., ; Jiang et al., ; Bai et al., , *inter alia*). Theoretically, the autoregressive Transformers have limited expressive power, but their capabilities can be expanded given sufficient chain-of-thought intermediate steps (Wei et al., ; Merrill & Sabharwal, ; Malach, ). However,  Lin et al. () demonstrates that the expressing of some next-tokens requires super-polynomial computational resources and is NP-hard to approximate.
Numerous advancements have been made upon the AR paradigm to compensate for modeling deficiencies, such as reverse training (Lee et al., ; Golovneva et al., ),fill-in-the-middle training (Bavarian et al., ), future-token prediction (Qi et al., ; Gloeckle et al., ), lookahead attention (Du et al., ) during the training stage, as well as search-augmented decoding (Lu et al., ; Xie et al., ; Yao et al., , *inter alia*) during inference.
In practice, autoregressive next-token predictors are shown to be ill-suited for planning tasks (Bubeck et al., ; Valmeekam et al., ; ; Dziri et al., ; Kambhampati et al., ).
Besides,  Bachmann & Nagarajan (); Lin et al. () 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.

### 5.2 Non-autoregressive Modeling

The non-autoregressive (NAR) generation method, which produces all target tokens simultaneously given the source context, is first proposed by  (Gu et al., ) in the text field for machine translation, primarily due to the efficiency consideration. While a series of advancements have been made afterward (Lee et al., ; Gu et al., ; Ghazvininejad et al., ; Qian et al., ; Huang et al., , *inter alia*), traditional NAR models still tend to underperform AR models in terms of generation quality (Xiao et al., ).
Diffusion models (Sohl-Dickstein et al., ; Ho et al., ), a powerful class of generative models known for their impressive image-generation capabilities (Dhariwal & Nichol, ), have recently been applied to the field of text generation (Hoogeboom et al., ; Austin et al., ; Li et al., ; Campbell et al., ; Dieleman et al., ; Chen et al., ; Ye et al., ; Lovelace et al., ), reinforcement learning (Janner et al., ; Chi et al., ) and protein design (Xu et al., ; Hoogeboom et al., ; Corso et al., ). 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, ), connecting them to score matching (Hyvärinen & Dayan, ) and energy-based models (LeCun et al., ).
For text, the diffusion model can be seen as an extension of the traditional iterative NAR models (Gong et al., ) and has been shown to approach or outperform AR models on text perplexity (Han et al., ; Lou et al., ; Gulrajani & Hashimoto, ), diversity (Gong et al., ; ; Zhang et al., ) as well as various seq-to-seq tasks (Wu et al., ; Zheng et al., ; Ye et al., ).
In contrast, we compare diffusion with AR from a perspective of subgoal imbalance and demonstrates the effectiveness of diffusion in tasks requiring complex reasoning and planning.

## 6 Conclusion

This paper presents an 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 MDM, that performs significantly better than AR and previous diffusion models on various sophisticated planning tasks.
Our findings underscore the necessity to reevaluate the sequence modeling paradigm for modern large language models, especially in tackling challenging problem-solving tasks.

## Appendix A More Background and Derivation of Discrete Diffusion Models

### A.1 Background

Discrete diffusion probabilistic models are first introduced in Sohl-Dickstein et al. () for binary data, and later extended to categorical data in (Hoogeboom et al., ). Austin et al. () provides a general form of discrete diffusion and introduces multiple transition matrices, including an absorbing variant that draws close connections to masked language models (Devlin et al., ).
Several subsequent works push this line of research further from various aspects, such as incorporating editing-based operations (Johnson et al., ; Reid et al., ), casting permuted language models (Yang et al., ) as diffusion models (Hoogeboom et al., ), developing a continuous-time framework (Campbell et al., ), parameterizing the routing mechanism (Zheng et al., ), and investigating score functions for learning the reverse process Sun et al. (); Lou et al. ().

### A.2 Derivation Setup

We now provide a detailed derivation of the loss in Equation (6).
For a clear illustration, we initiate derivation with a single random variable ${\bm{x}}_{0}$ and ultimately link it with the multi-variable sequence ${\mathbf{x}}_{0}$.
Suppose ${\bm{x}}_{0}\sim q({\bm{x}}_{0})$ is a discrete random variable with $K$ possible categories and represented as a one-hot vector. The forward process $q({\bm{x}}_{1:T}|{\bm{x}}_{0})=\prod_{t=1}^{T}q({\bm{x}}_{t}|{\bm{x}}_{t-1})$ corrupts the original data ${\bm{x}}_{0}$ into a sequence of increasingly noisy latent variables ${\bm{x}}_{1:T}\coloneqq{\bm{x}}_{1},\dots,{\bm{x}}_{T}$. The learned backward process $p_{{\bm{\theta}}}({\bm{x}}_{0:T})=p({\bm{x}}_{T})\prod_{t=1}^{T}p_{{\bm{\theta
}}}({\bm{x}}_{t-1}|{\bm{x}}_{t})$ gradually denoises the latent variables
to the data distribution.
In discrete diffusion, both the forward and backward distribution are defined as categorical distribution, e.g., $q({\bm{x}}_{t}|{\bm{x}}_{t-1})=\mathrm{Cat}({\bm{x}}_{t};{\bm{p}}={\bm{Q}}_{t}
^{\top}{\bm{x}}_{t-1})$ and $p_{{\bm{\theta}}}({\bm{x}}_{t-1}|{\bm{x}}_{t})=q({\bm{x}}_{t-1}|{\bm{x}}_{t},f
({\bm{x}}_{t};{\bm{\theta}}))$, where ${\bm{Q}}_{t}$ is a pre-defined transition matrix of size $K\times K$ (Hoogeboom et al., ; Austin et al., ).

### A.3 The marginal and posterior

Starting from ${\bm{x}}_{0}$, we obtain the following $t$-step marginal and posterior at time $t-1$:

$$ $\displaystyle q({\bm{x}}_{t}|{\bm{x}}_{0})=\mathrm{Cat}\left({\bm{x}}_{t};{\bm {p}}=\overline{{\bm{Q}}}_{t}^{\top}{\bm{x}}_{0}\right),\quad\text{with}\quad \overline{{\bm{Q}}}_{t}={\bm{Q}}_{1}{\bm{Q}}_{2}\ldots{\bm{Q}}_{t}$ $\displaystyle q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})=\frac{q({\bm{x}}_{t} |{\bm{x}}_{t-1},{\bm{x}}_{0})q({\bm{x}}_{t-1}|{\bm{x}}_{0})}{q({\bm{x}}_{t}|{ \bm{x}}_{0})}=\mathrm{Cat}\left({\bm{x}}_{t-1};{\bm{p}}=\frac{{\bm{Q}}_{t}{\bm {x}}_{t}\odot\overline{{\bm{Q}}}_{t-1}^{\top}{\bm{x}}_{0}}{{\bm{x}}_{t}^{\top} \overline{{\bm{Q}}}_{t}^{\top}{\bm{x}}_{0}}\right),$ (9) $$

where $q({\bm{x}}_{t}|{\bm{x}}_{t-1},{\bm{x}}_{0})=q({\bm{x}}_{t}|{\bm{x}}_{t-1})$ due to the Markov property of the forward process. The KL divergence between $q$ and $p_{\bm{\theta}}$ can be computed by simply summing over all possible values of each random variable.
The cumulative products $\overline{{\bm{Q}}}_{t}$, which can be computed in closed form or precomputed for all $t$ depending on the choice ${\bm{Q}}_{t}$, may be prohibitive for large $T$ and number of categories. Therefore, two commonly used forms of ${\bm{Q}}$ are introduced by  Hoogeboom et al. () and  Austin et al. (), which ensures $\overline{{\bm{Q}}}_{t}$ can still be computed efficiently, allowing the framework to scale to a larger number of categories.

### A.4 Transition Matrix

Austin et al. () introduced multiple types of the transition matrix ${\bm{Q}}_{t}$, such as uniform (Hoogeboom et al., ), absorbing, discretized Gaussian and token embedding distance.
The absorbing noise for discrete diffusion has been demonstrated to outperform the others (Austin et al., ), where the transition matrix is given by :

$$ $\displaystyle\left[{\bm{Q}}_{t}\right]_{ij}=\begin{cases}1\quad&\text{if}\quad i =j=m\\ 1-\beta_{t}\quad&\text{if}\quad i=j\neq m\\ \beta_{t}\quad&\text{if}\quad j=m,i\neq m\end{cases}.$ $$

The transition matrix can also be written as $(1-\beta_{t})I+\beta_{t}\bm{1}e^{\top}_{m}$, where $e_{m}$ is a vector with a one on the absorbing state $m$ and zeros elsewhere.
Since $m$ is an absorbing state, the corruption process converges not to a uniform distribution but to the point-mass distribution on $m$.
The transition matrices $\overline{{\bm{Q}}}={\bm{Q}}_{1}{\bm{Q}}_{2}\ldots{\bm{Q}}_{t}$ can be computed in closed form.
Specifically, we transition to another token with probability $\beta_{t}$ and stay the same with probability $1-\beta_{t}$ in each step. After $t$ steps, the only operative quantity is the probability of not yet having transitioned to another token, given by ${\alpha_{t}}=\prod_{i=0}^{t}(1-\beta_{i})$. Therefore, we have $\overline{{\bm{Q}}}_{t}={\alpha_{t}}I+(1-{\alpha_{t}})\bm{1}e_{m}^{\top}$.

### A.5 Derivation of ELBO

In order to optimize the generative model $p_{{\bm{\theta}}}({\bm{x}}_{0})$ to fit the data distribution $q({\bm{x}}_{0})$, we typically
minimize a variational upper bound on the negative log-likelihood, defined below:

$$ $\displaystyle-\log p_{{\bm{\theta}}}({\bm{x}}_{0})$ $\displaystyle=-\log\int p_{{\bm{\theta}}}({\bm{x}}_{0},{\bm{x}}_{1},\dots,{\bm {x}}_{T})d{\bm{x}}_{1}\cdots d{\bm{x}}_{T}$ $\displaystyle=-\log\int\frac{p_{{\bm{\theta}}}({\bm{x}}_{0},{\bm{x}}_{1},\dots ,{\bm{x}}_{T})}{q({\bm{x}}_{1},\dots,{\bm{x}}_{T}|{\bm{x}}_{0})}q({\bm{x}}_{1} ,\dots,{\bm{x}}_{T}|{\bm{x}}_{0})d{\bm{x}}_{1}\cdots d{\bm{x}}_{T}$ $\displaystyle=-\log\mathbb{E}_{q({\bm{x}}_{1},\dots,{\bm{x}}_{T}|{\bm{x}}_{0}) }\Big{[}{\frac{p_{{\bm{\theta}}}({\bm{x}}_{0},{\bm{x}}_{1},\dots,{\bm{x}}_{T}) }{q({\bm{x}}_{1},\dots,{\bm{x}}_{T}|{\bm{x}}_{0})}}\Big{]}$ $\displaystyle\leq-\mathbb{E}_{q({\bm{x}}_{1},\dots,{\bm{x}}_{T}|{\bm{x}}_{0})} \Big{[}{\log\frac{p_{{\bm{\theta}}}({\bm{x}}_{0},{\bm{x}}_{1},\dots,{\bm{x}}_{ T})}{q({\bm{x}}_{1},\dots,{\bm{x}}_{T}|{\bm{x}}_{0})}}\Big{]}$ $\displaystyle=-\mathbb{E}_{q({\bm{x}}_{1},\dots,{\bm{x}}_{T}|{\bm{x}}_{0})} \Big{[}{\log\frac{p_{{\bm{\theta}}}({\bm{x}}_{0}|{\bm{x}}_{1})p_{{\bm{\theta}} }({\bm{x}}_{T})\prod_{t=2}^{T}p_{{\bm{\theta}}}({\bm{x}}_{t-1}|{\bm{x}}_{t})}{ q({\bm{x}}_{T}|{\bm{x}}_{0})\prod_{t=2}^{T}q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{ x}}_{0})}}\Big{]}$ $\displaystyle=-\mathbb{E}_{q({\bm{x}}_{1},\dots,{\bm{x}}_{T}|{\bm{x}}_{0})} \Big{[}{\log p_{{\bm{\theta}}}({\bm{x}}_{0}|{\bm{x}}_{1})-\sum_{t=2}^{T}\log \frac{q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})}{p_{{\bm{\theta}}}({\bm{x}}_ {t-1}|{\bm{x}}_{t})}-\log\frac{q({\bm{x}}_{T}|{\bm{x}}_{0})}{p_{{\bm{\theta}}} ({\bm{x}}_{T})}}\Big{]}$ $\displaystyle=-\mathbb{E}_{q}\Big{[}{\log p_{{\bm{\theta}}}({\bm{x}}_{0}|{\bm{ x}}_{1})-\sum_{t=2}^{T}D_{\mathrm{KL}}[{q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}} _{0})}||{p_{{\bm{\theta}}}({\bm{x}}_{t-1}|{\bm{x}}_{t})}]-\underbrace{D_{ \mathrm{KL}}[q({\bm{x}}_{T}|{\bm{x}}_{0})||p({\bm{x}}_{T})]}_{L_{T}(\text{ const})}}\Big{]}$ $\displaystyle=\underbrace{-\mathbb{E}_{q({\bm{x}}_{1}|{\bm{x}}_{0})}{\log p_{{ \bm{\theta}}}({\bm{x}}_{0}|{\bm{x}}_{1})}}_{{L}_{0}}+\sum_{t=2}^{T}\underbrace {\mathbb{E}_{q({\bm{x}}_{t}|{\bm{x}}_{0})}\Big{[}{D_{\mathrm{KL}}[{q({\bm{x}}_ {t-1}|{\bm{x}}_{t},{\bm{x}}_{0})}||{p_{{\bm{\theta}}}({\bm{x}}_{t-1}|{\bm{x}}_ {t})}]}\Big{]}}_{{L}_{t-1}}+L_{T}(\text{const}).$ $$

### A.6 Derivation for Equation ( 6 )

The categorical distribution $q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})$ based on Equation (9) is given as:

$$ $\displaystyle q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})$ $\displaystyle=\frac{{\bm{Q}}_{t}{\bm{x}}_{t}\odot\overline{{\bm{Q}}}_{t-1}^{ \top}{\bm{x}}_{0}}{{\bm{x}}_{t}^{\top}\overline{{\bm{Q}}}_{t}^{\top}{\bm{x}}_{ 0}}$ $\displaystyle=\!\frac{\left[(1-\beta_{t}){\bm{x}}_{t}+\beta_{t}\sigma_{{\bm{x} }_{t}}\bm{1}\right]\odot\left[\alpha_{t-1}{\bm{x}}_{0}+(1-\alpha_{t-1})e_{m} \right]}{\alpha_{t}{\bm{x}}_{t}^{\top}{\bm{x}}_{0}+(1-\alpha_{t}){\bm{x}}_{t}^ {\top}e_{m}}$ $\displaystyle=\!\frac{(1-\beta_{t})\alpha_{t-1}{{\bm{x}}_{t}\!\odot\!{\bm{x}}_ {0}}\!+\!(1-\beta_{t})(1\!-\!\alpha_{t-1}){{\bm{x}}_{t}\!\odot\!e_{m}}\!+\! \beta_{t}\alpha_{t-1}\sigma_{{\bm{x}}_{t}}\!{\bm{1}\!\odot\!{\bm{x}}_{0}}\!+\! \beta_{t}(1\!-\!\alpha_{t-1})\sigma_{{\bm{x}}_{t}}\!{\bm{1}\!\odot\!e_{m}}}{ \alpha_{t}{\bm{x}}_{t}^{\top}{\bm{x}}_{0}+(1-\alpha_{t}){{\bm{x}}_{t}^{\top}e_ {m}}}$ $\displaystyle=\!\frac{(1-\beta_{t})\alpha_{t-1}{{\bm{x}}_{t}\!\odot\!{\bm{x}}_ {0}}+(1-\beta_{t})(1\!-\!\alpha_{t-1}){\sigma_{{\bm{x}}_{t}}{\bm{x}}_{t}}+ \beta_{t}\alpha_{t-1}\sigma_{{\bm{x}}_{t}}{{\bm{x}}_{0}}+\beta_{t}(1\!-\! \alpha_{t-1})\sigma_{{\bm{x}}_{t}}{e_{m}}}{\alpha_{t}{\bm{x}}_{t}^{\top}{\bm{x }}_{0}+(1-\alpha_{t}){\sigma_{{\bm{x}}_{t}}}},$ $$

where $\sigma_{{\bm{x}}_{t}}\coloneqq e_{m}({\bm{u}}={\bm{x}}_{t})$ represents the probability of noise drawn from $e_{m}$ being equal to ${\bm{x}}_{t}$. Note ${\bm{x}}_{t}\odot{\bm{x}}_{0}=0$ if ${\bm{x}}_{t}\neq{\bm{x}}_{0}$ otherwise 1. Thus the computation of $q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})$ breaks down into two cases:

$$ $\displaystyle q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})=\begin{cases}\eta_{t }{\bm{x}}_{t}+\left(1-\eta_{t}\right)e_{m},&\text{ if }{\bm{x}}_{t}={\bm{x}}_{ 0}\\ \lambda_{t}{\bm{x}}_{0}+\left(1-\lambda_{t}\right)e_{m}({\bm{x}}_{t}),&\text{ if }{\bm{x}}_{t}\neq{\bm{x}}_{0},\end{cases}$ $$

where $\eta_{t}\coloneqq 1-\frac{\beta_{t}(1-\alpha_{t-1})e_{m}({\bm{u}}={\bm{x}}_{t}
)}{\alpha_{t}+(1-\alpha_{t})e_{m}({\bm{u}}={\bm{x}}_{t})}$, $\lambda_{t}\coloneqq\frac{\alpha_{t-1}-\alpha_{t}}{1-\alpha_{t}}$, and $e_{m}({\bm{x}}_{t})=(1-\beta_{t}){\bm{x}}_{t}+\beta_{t}e_{m}$ denotes a noise distribution that interpolates between ${\bm{x}}_{t}$ and $e_{m}$.

Recall the distribution $p_{{\bm{\theta}}}({\bm{x}}_{t-1}|{\bm{x}}_{t})$ is parameterized by $q({\bm{x}}_{t-1}|{\bm{x}}_{t},f({\bm{x}}_{t};{\bm{\theta}}))$, the KL divergence between $q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})$ and $p_{{\bm{\theta}}}({\bm{x}}_{t-1}|{\bm{x}}_{t})$ becomes 0 when ${\bm{x}}_{t}={\bm{x}}_{0}$.
In the case of absorbing diffusion, ${\bm{x}}_{t}=e_{m}$ if ${\bm{x}}_{t}\neq{\bm{x}}_{0}$ and $e_{m}({\bm{x}}_{t})=e_{m}$. $q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})$ has probability $\lambda_{t}$ on index $x_{0}$ and $1-\lambda_{t}$ on the absorbing state. The model $f({\bm{x}}_{t};{\bm{\theta}})$ has zero-probability on the absorbing state as it never predicts the mask token. Therefore, $p_{{\bm{\theta}}}({\bm{x}}_{t-1}|{\bm{x}}_{t})$ also has $1-\lambda_{t}$ probability on the absorbing state. Putting them together, we derive the KL divergence as:

$$ $\displaystyle D_{\mathrm{KL}}[q({\bm{x}}_{t-1}|{\bm{x}}_{t},{\bm{x}}_{0})||p_{ {\bm{\theta}}}({\bm{x}}_{t-1}|{\bm{x}}_{t})]$ $\displaystyle=1_{x_{t}\neq x_{0}}[\lambda_{t}\log\frac{\lambda_{t}}{f({\bm{x}} _{t};{\bm{\theta}})_{x_{0}}}+(1-\lambda_{t})\log\frac{1-\lambda_{t}}{1-\lambda _{t}}]$ $\displaystyle=-\lambda_{t}1_{x_{t}\neq x_{0}}{\bm{x}}_{0}^{\top}\log f({\bm{x} }_{t};{\bm{\theta}})+C,$ $$

where $1_{x_{t}\neq x_{0}}$ is 1 if $x_{t}\neq x_{0}$ otherwise 0, and $C$ is a constant. Moreover, given $\alpha_{0}=1$ by definition and $\lambda_{0}=1$, we have:

$$ $\displaystyle L({\bm{x}}_{0})=-\mathbb{E}_{q({\bm{x}}_{0})}\sum_{t=1}^{T} \lambda_{t}\mathbb{E}_{q({\bm{x}}_{t}|{\bm{x}}_{0})}1_{{\bm{x}}_{t}\neq{\bm{x} }_{0}}{\bm{x}}_{0}^{\top}\log f({\bm{x}}_{t};{\bm{\theta}})$ $$

for a single random variable, and

$$ $\displaystyle L({\mathbf{x}}_{0})=-\sum_{n=1}^{N}\mathbb{E}_{q({\mathbf{x}}_{0 ,n})}\sum_{t=1}^{T}\lambda_{t}\mathbb{E}_{q({\mathbf{x}}_{t,n}|{\mathbf{x}}_{0 ,n})}1_{{\mathbf{x}}_{t,n}\neq{\mathbf{x}}_{0,n}}{\mathbf{x}}_{0,n}^{\top}\log f ({\mathbf{x}}_{t,n};{\bm{\theta}})$ $$

for ${\mathbf{x}}_{0}$ that represents a sequence of random variables ${\mathbf{x}}_{0}=({\bm{x}}_{0,1},\dots,{\bm{x}}_{0,N})$, where the $\lambda_{t}$ also represents the reweighting term $w(t)$ in Equation (6).

## Appendix B Algorithms for Training and Inference

The detailed algorithms for training and inference are illustrated in Algorithm 1 and 2, respectively. For conditional training and inference, we split ${\mathbf{x}}$ into $[{\mathbf{x}}^{\text{src}};{\mathbf{x}}^{\text{tgt}}]$ and freeze the condition part ${\mathbf{x}}^{\text{src}}$ during training and inference.

Figure: Algorithm 1 Training MDM

Figure: Algorithm 2 Sampling from MDM

## Appendix C Additional Experimental Details

### C.1 Task Details

**Table 4: Dataset statistics. Minimal and CD are short for the minimal planning task and Countdown, respectively.**
|  | Minimal | CD3 | CD4 | CD5 | Sudoku | 3-SAT 5v | 3-SAT 7v | 3-SAT 9v |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| Train Instance | 50k | 500k | 500k | 500k | 100k | 50k | 50k | 100k |
| Test Instance | 1k | 1k | 1k | 1k | 1k | 1k | 1k | 1k |
| Avg Input Token | 47 | 11 | 13 | 16 | 81 | 245 | 269 | 305 |
| Avg Output Token | 21 | 16 | 25 | 35 | 81 | 9 | 13 | 17 |
| 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 4 and Table 8, respectively.

**Table 5: Model parameters with varying model size.**
|  | Tiny | Base | Medium |
| --- | --- | --- | --- |
| Parameters | 6M | 85M | 303M |
| Num of Layer | 3 | 12 | 24 |
| Num of Head | 12 | 12 | 16 |
| Hidden Dim | 384 | 768 | 1024 |

### C.2 MDM 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 MDM and AR.
We set the learning rate to 1e-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 MDM for 1200 epochs on the minimal planning task, 300 epochs on Sudoku, and 600 epochs on other datasets.
By default, we set the diffusion sampling steps to $T=20$ for tasks with average output tokens larger than 20, otherwise $T=10$. We use a decoding temperature $\tau=0.5$ for all tasks.
For all the experiments, we have verified the statistical significance by running them multiple times.

### 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 MDM.

For LLaMA (Touvron et al., ), we use LoRA fine-tuning (Hu et al., ) with lora rank setting to 16. We use a learning rate of 1e-4, a batch size of 256, and train for a maximum of 20 epochs to ensure the model has converged.
For GPT-4, we borrow the numbers from Yao et al. ().

For all the diffusion baselines, we use the same transformer architecture as GPT-2 to control the variables. We set the training parameters the same as MDM, 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 a clear performance improvement when scaling to 1024.

## Appendix D Additional Experiments

### D.1 Token Consumption on Game of 24

We show the detailed accuracy and token consumption on the game of 24 in Table 6.

**Table 6: Detailed accuracy and token consumption on game of 24.**
|  | Accuracy | Prompt Tokens | Generate Tokens |
| --- | --- | --- | --- |
| GPT-4 IO | 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-Scratch | 18.8 | 11 | 26 |
| MDM | 76.0 | 11 | 26 |

### D.2 Case Study

**Table 7: Example predictions on Countdown 4. For each sub-equation, we mark the planning error in red and the calculation error in bold. AR exhibits more calculation errors in the last equation due to incorrect planning in the previous steps.**
| Numbers | Goal | Groundtruth | AR Prediction | MDM Prediction |
| --- | --- | --- | --- | --- |
| 64,36,52,42 | 14 | 64-52=12,36/12=3,42/3=14 | 64/36=2,52/2=26,42-26=14 | 64-52=12,36/12=3,42/3=14 |
| 9,73,99,75 | 81 | 75-73=2,9*2=18,99-18=81 | 99+75=174,174/9=16,73+16=81 | 75-73=2,9*2=18,99-18=81 |
| 2,52,20,73 | 57 | 52-20=32,32/2=16,73-16=57 | 2*20=40,73-52=21,40+21=57 | 52-20=32,32/2=16,73-16=57 |
| 9,80,4,5 | 89 | 9+80=89,5-4=1,89*1=89 | 9-5=4,4/4=1,80+1=81 | 9+80=89,5-4=1,89*1=89 |
| 65,2,61,22 | 96 | 65-61=4,2+22=24,4*24=96 | 65-61=4,22*4=88,2+88=90 | 65-61=4,2+22=24,4*24=96 |
| 42,47,9,14 | 81 | 47-42=5,14-5=9,9*9=81 | 47-42=5,14*5=70,9+70=89 | 47-42=5,14*5=70,9+70=81 |
| 41,4,48,20 | 96 | 41*4=164,48+20=68,164-68=96 | 48-41=7,20-7=13,4*13=92 | 4*20=80,41-40=2,48*2=96 |
| 21,36,3,42 | 39 | 42-36=6,3*6=18,21+18=39 | 36-21=15,15/3=5,42-5=37 | 42-21=21,36/3=12,21-12=39 |

We show more prediction cases of the autoregressive model and our model in Table 7.

**Table 8: Task details by showing example input and output for each dataset.**
| Task | Input Example | Output Example |
| --- | --- | --- |
| Minimal Planning | 2,10/10,4/11,5/2,0/8,2/0,11/6,2/1,9/5,3/4,1-8,3 | 8,2/2,0/0,11/11,5/5,3 |
| Countdown 3 | 15,44,79,50 | 44-15=29,79-29=50 |
| Countdown 4 | 86,28,13,31,96 | 86+28=114,31-13=18,114-18=96 |
| Countdown 5 | 50,36,82,44,31,51 | 44-36=8,82*31=2542,<br>8+2542=2550,2550/50=51 |
| Sudoku | 080050060<br>460907108<br>005000029<br>970006500<br>000872031<br>300049000<br>004025003<br>010000480<br>603100007 | 789251364<br>462937158<br>135468729<br>978316542<br>546872931<br>321549876<br>894725613<br>217693485<br>653184297 |
| 3-SAT 5v | 1,4,5/1,-4,-5/2,-4,5/-1,-2,5/3,4,5/-2,-4,-5/2,3,-4/-2,-3,5/1,2,4/1,-2,3/-1,3,5/1,-2,-4/1,4,-5/1,-2,-5/1,2,-5/-1,-3,-4/-1,3,-5/-1,3,4/2,-4,-5/-1,-4,5/1,-3,-5/1,3,-5/1,-3,-4/-2,3,5/1,2,5/-1,2,-4/1,-2,4/1,-4,5/3,4,-5/-1,2,-3/1,-3,5/-2,4,5/1,-2,5/-1,2,5/1,3,-4/-1,-4,-5/-2,-3,-4/2,4,5/-2,3,-4/-3,4,5/2,-3,5 | 1,2,3,-4,5 |
| 3-SAT 7v | -2,-3,-7/2,-4,-7/-3,4,-5/1,2,-3/1,5,-7/-5,-6,-7/2,-5,6/2,-5,-6/-3,-4,6/-1,2,-4/-3,6,7/-2,-5,6/2,3,-7/-1,2,3/-2,3,-4/-1,3,7/1,-2,-7/2,4,6/1,2,-7/2,-3,-6/1,-2,6/-1,5,7/3,-6,-7/2,6,7/-2,-6,-7/-2,3,-5/3,5,-6/-2,6,-7/-1,-2,-7/5,-6,-7/2,-6,-7/-2,5,7/-3,-4,5/2,3,-4/-3,5,-7/3,-4,5/-2,3,-6/1,2,-6/1,4,-7/1,4,7/2,4,5/1,5,-6/1,3,4/2,3,7/1,-2,4 | 1,2,3,4,5,6,-7 |
| 3-SAT 9v | 3,-4,-6/1,3,5/2,-7,8/1,-3,6/2,-3,-8/-4,-5,-7/1,-6,-9/1,8,-9/2,3,-9/3,-5,9/-3,7,9/-2,-3,9/-1,-5,-9/-2,-7,-9/-1,3,5/2,-5,-9/4,-7,-9/-2,3,-8/2,3,7/2,-4,6/-2,3,5/-2,-6,-8/-3,-4,-8/-2,6,7/-3,4,6/-3,-6,9/2,7,-9/2,4,-5/-3,-5,8/-4,5,-7/-4,-6,-8/2,-6,9/2,-5,9/1,4,-9/5,8,9/1,-6,7/-3,6,-9/1,4,-5/4,-6,9/-1,2,6/1,-2,-5/1,-2,-9/-4,7,9/-1,-4,-7/-3,5,-8/-1,-3,6/-2,-3,6/-3,6,9/-1,-5,8/1,-5,-9/1,4,8 | 1,2,3,4,-5,6,-7,-8,9 |

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. How does subgoal imbalance impact performance in real-world RL tasks? Could you provide concrete examples of this issue in practical applications?
2. How does the performance of the proposed method compare to autoregressive models in terms of computational efficiency (i.e., training time and inference time)?
soundness: 3
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.
- They propose a new method to tackle subgoal imbalance, which appears to be a reasonable approach to addressing this challenge.
- The authors demonstrate that their method outperforms traditional autoregressive models on complex tasks, including Countdown, Sudoku, and Boolean Satisfiability Problems.
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.
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.
- 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.
- 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.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 4
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.
2. "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. 
3. 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)?

[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).
soundness: 2
strengths: The paper is very well-motivated and well-written.

I 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.

The 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.
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. 

I have a few questions about several of the paper's claims. If properly addressed, I think this would be a good paper for acceptance.
weaknesses: I do not consider these weaknesses, but any paper can always improve:

The 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.

[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).

---

confidence: 3
contribution: 4
initial_rating: 8
presentation: 4
questions: In addition to the questions in the weaknesses:
- 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?
soundness: 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. 
- 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.
- 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.
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.
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.’ 
- 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.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
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? 
- 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?
soundness: 2
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.
-  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. 
- 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.
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. 

They 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.

The 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.

Given 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. 

The 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.
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. 
- 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.
- 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. 
- 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. 
- 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?

Meta-review: This paper proposes a particular discrete diffusion model to address reasoning and planning. All reviewers appreciate the merit and novelty of its technical contribution. Furthermore, it establishes a well-presented and clear connection between the bottlenecks in planning and why traditional autoregressive models would struggle with them. In particular, the problem is subgoal imbalance. Still, the definition of “subgoal imbalance” could leverage more illustrative examples. The theory developed can be more elaborate. The experiment setting is solid with hard problems. More real-world-based baselines definitely are welcome. Overall, it is solid writing with clear technical contributions.

Decision: Accept (Poster)