Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planning

We sincerely thank Reviewer gxqG for your detailed review and are grateful for the time you spent on our submission. We are also glad to know you think our paper is very well structured and original. Below we would like to give detailed responses to each of your comments.

Weaknesses1: 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.

Thanks for your great suggestion. The subgoal imbalance is a general aspect of multi-step tasks like planning and reasoning, where different subgoals vary in difficulty, as shown in Figure 2. This imbalance exists regardless of whether the tasks are modeled autoregressively or using other approaches—it is inherent to the data and the nature of the tasks themselves. Since our work focuses on the weaknesses of the autoregressive (AR) paradigm, we explained subgoal imbalance from this perspective.

We intend to illustrate how subgoal imbalance particularly affects AR models due to their sequential processing nature. We acknowledge that the formal proposition at line 160 may give the impression that subgoal imbalance is specific to AR models. To address this, we have revised the proposition in the updated manuscript to involve only the data distribution, emphasizing that subgoal imbalance is a property of data distribution, not tied to any specific modeling approach.

Weaknesses2: 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.

Thank you for your suggestion. We have added detailed descriptions in the updated manuscript. Each $x_t$ can be interpreted as a distinct view of $x_0$, where each view provides different information about $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(x_0,x_t, n;\theta)$, our findings suggest that objectives challenging to learn in AR models become more feasible to learn in diffusion models.

Weaknesses3: 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.

Thank you for your suggestion. The multi-view learning perspective is an advantage introduced by diffusion models, as they naturally capture multiple representations (views) of the data. Our MDM method builds on this by applying weighting to the more challenging views to effectively address subgoal imbalance. We have clarified this connection in the revised manuscript to better explain how the observations in Section 3.2 led to the development of MDM in Section 3.3.

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

Thanks for your feedback. Our main objective was to analyze the performance of autoregressive (AR) models versus diffusion models under fair and comparable conditions, focusing on models of the same scale to ensure a balanced and affordable comparison. GPT-4 with Tree of Thought (ToT) was included only for the Game of 24 as a bonus reference to demonstrate that even larger AR models rely heavily on search strategies to perform well on complex tasks. This highlights that without such search mechanisms, large AR models may struggle, whereas our diffusion-based approach achieves strong performance without additional search.

We sincerely thank Reviewer aoA3 for your review and are grateful for the time you spent on our submission. We are also glad you think our methodology is novel and has strong contributions. Below we would like to give detailed responses to each of your comments.

Weakness 1: 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.’

Thanks for the great suggestion. The research on “common sense” often needs to scale up to larger models with more data, while in our current study, we began with small-scale, controlled tasks to thoroughly analyze and validate the effectiveness of our approach. We acknowledge the significance of scaling up to larger models and more complex and common sense-related tasks. As a next step, we are actively working on extending our model to handle more complex planning tasks with scaled models. We believe this will further showcase the capabilities of our approach and its potential impact on a broader range of applications.

Weakness 2 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.

Thanks for the great question. In our experiments, we compared models of different sizes (e.g., 85M and 303M parameters) using the same amount of training data for a strict comparison. The larger 303M model didn't outperform the 85M model on CD 4 and CD 5 tasks likely because the data was insufficient for the larger model to fully leverage its capacity, potentially leading to overfitting. This phenomenon also occurs in AR models under limited data conditions. We conducted additional experiments with increased data below, which shows that both increasing the model and data size leads to a consistent increase in performance on CD4 for both AR and MDM. We have included the experiment in the revised manuscript in Appendix D.3.

Question 1 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?

Thank you for the insightful question. We believe that with sufficient data and scaling, MDM could exhibit in-context learning properties similar to those of AR models, and we plan to investigate this in our future work.

We sincerely thank Reviewer SUNs for your review and are grateful for the time you spent on our submission. We are happy to know you think our paper is well-motivated and you enjoy Section 3 especially. Below we would like to give detailed responses to each of your comments.

Weakness1 “ 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.”

Thank you for your constructive feedback and for recognizing that the three datasets/domains we investigated are among the most popular and sufficient for demonstrating our approach. We agree that applying our method to larger domains such as Blocksworld [1] would further showcase its scalability and general applicability. Due to time constraints and computational resources, we were unable to include experiments on larger domains in this version of the paper. However, we will explore these domains in our future work.

Question 1: 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?

Thank you for your insightful question. In the case of Figure 1, when the given goal is 9, the value 5 is not simply generated through random sampling; rather, it must be inferred by reasoning backward from the goal of 9. This backward inference is challenging for AR as it only utilizes the preceding context for its predictions. Broadly speaking, we would like to point out that there are some subgoal imbalance issues at the data level, especially in complex planning tasks where a certain action/token is determinative to the subsequent action/tokens. We show such a subgoal imbalance issue can be very challenging for AR. We have improved the description for Proposition 1 in the updated version.

Question2 "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.

Thank you for your question. "Regretful compromise" here highlights a phenomenon: despite correctly calculating earlier steps, the model's incorrect planning forces it to make an incorrect calculation or decision at the final step in an attempt to achieve the overall goal. The model has the computational ability to find the correct solution, but because it cannot backtrack, it ends up making a mistake in trying to reach the goal. For example, in our paper, given the 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. The first two equations are calculated correctly, but the last one is incorrectly forced to equal 87. As shown in Figure 6(b), 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 in its realization of being unable to achieve the goal only at the end. By that point, it cannot correct the previous errors, leading it “regretfully” to rely on incorrect equations to reach the goal (e.g., the number 87 in the above case).

"Compounding error" refers to the accumulation of small prediction errors over sequential steps, leading to a divergence from the expected trajectory, so it is more like a general cause, whereas "regretful compromise" here is the specific phenomenon we observe specifically in the goal-conditioned planning tasks when the model, unable to backtrack, makes a compromise at the final step due to the incorrect intermediate planning.

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

We would like to make our main point again first. Our starting point is to explore why current LLMs in the text field struggle with complex reasoning and long-term planning tasks. To this end, we made four contributions:

• The insights on why planning is hard, typically for AR, through the lens of subgoal imbalance: We first point out the phenomenon of subgoal imbalance often present in complex reasoning and planning data, which makes it very difficult for the autoregressive (AR) structure of LLMs to learn certain subgoals (i.e., tokens), leading to poor prediction performance.
• The design of Minimal planning task: We designed a minimal planning task to describe the concept of planning and experimentally validated the shortcomings of the AR structure.
• The explanation of how diffusion learns subgoals that challenge the AR model both theoretically and empirically: We theoretically explain why diffusion methods can better address the subgoal imbalance issue and validate this on the aforementioned planning task.
• The propose of MDM: We propose MDM to enhance the efficiency of existing diffusion models in optimization, which prioritizes subgoals based on difficulty during learning. We have validated that MDM significantly outperforms autoregressive models on more complex real-world tasks such as Countdown, Sudoku, and Boolean Satisfiability Problems.

Most of the current generation of LLMs are AR-based. Our work highlights the potential of diffusion-based approaches in advancing AI capabilities for sophisticated language understanding and problem-solving tasks, and we hope to promote the development of the next generation of diffusion-based LLMs.

Weakness 1: 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.

Thanks for the feedback. Since our focus is on exploring the issues with LLMs, we primarily concentrate on the text field. Extending the ideas and methods presented in this paper to an RL setting is interesting and feasible, but it falls outside the scope of our work.

Weakness 2: 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.

Thank you for your comment. The “subgoal” studied in this work particularly refers to “token” in language. Our work specifically addresses challenges in autoregressive (AR) language models when dealing with subgoal imbalance, and we find that diffusion models naturally offer a solution. The token reweighting technique can be applied to AR models as well, but it doesn’t overcome the inherent limitations of AR—namely, their inability to plan for future subgoals or revise past outputs based on future context. In our experiments, applying token reweighting to AR models still cannot solve subgoals with PD>1 (accuracy around 50%), similar to that in the original AR.

Weakness 3: 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.

Thank you for your feedback. Our work primarily focuses on analyzing the shortcomings of AR language models in handling subgoal imbalance, which is a practical issue observed in existing AR LLMs, especially when dealing with complex reasoning and planning tasks. While subgoal imbalance may also impact the RL field, our main interest is in addressing this problem within the context of language modeling tasks where AR models face challenges in planning and reasoning due to their sequential generation process. Therefore, in this paper, we concentrate on scenarios relevant to AR LLMs rather than RL, and highlight the practical implications of subgoal imbalance in current language modeling applications.

Dear Reviewer,

We would like to first express our sincere gratitude for your time and effort in reviewing our paper, and we truly appreciate your constructive feedback regarding our work.

As the author-reviewer discussion period is coming to a close, we wonder if you could kindly share some of your thoughts so we can keep the discussion rolling to address your concerns if there are any. We are eager to address any additional questions or issues you may have.

We would greatly appreciate it if you would consider adjusting the score based on our responses and the other review comments.

Thank you once again for your constructive feedback!

Sincerely,

The Authors