Diffusion of Thought: Chain-of-Thought Reasoning in Diffusion Language Models

We sincerely thank Reviewer q6Ny 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.

Weakness 1: Confusion about the coupled sampling strategy

Thanks for pointing out this potential confusion. The main purpose of the coupled sampling mechanism is to enable the DoT$^{MP}$ model with the ability to correct potential errors in the previous thoughts. Without coupled sampling, the discrepancy arises between the use of correct previous thoughts during training and the possibly erroneous generated thoughts during testing, leading to error accumulation akin to that in autoregressive models. To alleviate this issue, we introduce the coupled sampling strategy for the model to learn to rectify past errors. This strategy involves injecting noise into previous thoughts during the training phase, enabling the model to possess the capability to see and correct errors in previous thoughts. We will add more details to clarify this confusion in the final version.

Weakness 2: Comparisons with the DiffuSeq model

Thank you for bringing up this point. We have this comparison in Table 2 and here is the summary of the comparison between DiffuSeq and DoT. We will clarify this confusion in the final version.

Weakness 3: Qualitative comparison of the reasoning paths between DoT and DoT$^{MP}$

Thank you for your suggestion. We observe that DoT$^{MP}$ outperforms DoT in correctness regarding the reasoning paths, while DoT slightly excels in diversity as depicted in Figure 4(b). Below we show some examples where DoT$^{MP}$ can predict the correct reasoning path while DoT fails. More content related to the reasoning path analysis will be incorporated in the paper accordingly.

Query: The Kennel house keeps 3 German Shepherds and 2 Bulldogs. If a German Shepherd consumes 5 kilograms of dog food and a bulldog consumes 3 kilograms of dog food per day. How many kilograms of dog food will they need in a week?

DoT: <<35=15>> <<73=21>> <<15+21=36>> #### 36

DoT$^{MP}$: <<35=15>> <<23=6>> <<15+6=21>> <<21*7=147>> #### 147

Query: Skyler has 100 hats on his hand with the colors red, blue, and white. Half of the hats are red, 3/5 of the remaining hats are blue, and the rest are white. How many white hats does Skyler have?

DoT: <<1/2100=50>> <<3/550=30>> <<100-30=70>> #### 70

DoT$^{MP}$: <<100/2=50>> <<100-50=50>> <<50*3/5=30>> <<50-30=20>> #### 20

Question 1: Confusion about L128-129

Thanks for pointing out this potential confusion. The first line of Table 2 is the tuned Plaid using the gradient-based token guidance to generate response and it achieves poor performance. In Plaid, the use of gradient-based guidance to inject conditions involves adjusting random source embeddings through gradients to match condition tokens. However, we observe discrepancies between recovered source tokens and condition tokens, which can adversely affect tasks requiring precise conditioning. Below we show an example on grade school math as a demonstration, where bold words in the query part are incorrectly recovered. We can see there are three recovered query tokens that exhibit minor differences due to soft gradient guidance, causing interference with the model's comprehension of the problem. That’s why we resort to hard control with gradient-free conditioning. We will add more details to clarify this confusion in the final version.

Groundtruth: Two trains leave San Rafael at the same time. They begin traveling westward, both traveling for 80 miles. The next day, they travel northwards, covering 150 miles. What's the distance covered by each train in the two days? <<280=160>> <<1502=300>> <<300+160=460>> <<460/2=230>> #### 230

Prediction: Three trains leave San Juan at the same time. They start traveling westward, both traveling for 80 miles. The next day, they travel southward, covering 150 miles. What's the distance covered by each train in the two days? <<3*80=180>> <<180+80+150=340>> <<340/ 30=12.5>> #### 12.5

Question 2: About probability $\epsilon$

Thank you for bringing up this excellent question. On one hand, we actually desire the presence of noise in $z_0$, as it forces the model to pose the self-correction ability. On the other hand, due to our small $\epsilon_\text{min}$, i.e., 0.95, we primarily rely on gold data, thereby preventing excessive noise in $z_0$ that could make training too challenging. We previously attempted a warmup process by integrating scheduled sampling after a certain number of steps, but we did not observe significant performance improvements. Hence, for the sake of simplicity, we opted for the current approach.

Question 3: How does the model decide on the number of rationales

Thank you for bringing up this point. During training, we append a special token <EOS> to the last thought, so when the model generates a thought followed by <EOS>, it stops generating further. We will add this detail to our manuscript.

We sincerely thank Reviewer Y1KX for your review and are grateful for the time you spent on our submission. Below we would like to give detailed responses to each of your comments.

Weakness 1: Direct Comparison Baseline

Thank you for your suggestion. We conduct the answer-only setting to further validate the effectiveness of DoT. The result table reveals that fine-tuning diffusion models solely with answer data leads to inferior performance compared to DoT, mirroring the degradation of AR models in the absence of CoT.

"The paper does not adequately separate between the specific contributions of the DoT methodology from the inherent advantages of using diffusion language models."

Thank you for bringing up this point. Firstly, one of our contributions is exactly employing diffusion language models for multi-step text reasoning. To the best of our knowledge, we are the first to bring diffusion into the realm of complex text reasoning such as mathematical reasoning. Additionally, we show fine-tuning a pretrained diffusion model is a non-trivial work. As demonstrated in the ablation study (Table 2), directly following the pretraining approach leads to subpar results. Therefore, we resorted to an infilling approach and further proposed a series of sampling strategies and multipass variants to enhance the model's performance. The comparison between applying CoT and DoT to diffusion is presented below.

Weakness 2: Missing iCoT context

Thank you for your suggestion about more description about iCoT. In Section 3.1, we list 3 parallel approaches of modeling CoT: AR, iCoT, DoT, theoretically. Below, we discuss some similarities and differences between iCoT and DoT. DoT shares similarities with iCoT in the following 3 high-level aspects: i) Both DoT and iCoT try to tackle the time cost of auto-regressively generating the chain-of-thought rationales; ii) Both DoT and iCoT process “thoughts” “vertically’’ in hidden dimension. But DoT presents the hidden information across different diffusion timesteps (in temporal dimension), while iCoT presents the hidden information across the model’s different layers (in spatial dimension); iii) Both DoT and iCoT evaluate the model’s CoT ability, so we refer to some experiment settings of iCoT, including datasets and baselines.

However, in terms of methodology, iCoT and DoT are completely different. iCoT still relies on next-token prediction for autoregressive generation, while DoT utilizes diffusion, which offers additional advantages such as a flexible performance-efficiency trade-off and self-correction capability beyond efficiency. We will add more clarification about DoT and iCoT in the paper.

Thanks for pointing out the potential confusion about the names of ‘single pass’ and ‘multi-pass’, we will clarify them in the paper.

Weakness 3: Task-Specific Fine-Tuning Requirement

Thank you for your constructive comments. The main reason is that the current pre-trained diffusion language models are relatively small, resulting in the underutilization of their in-context learning capabilities, as we discussed in the limitations. Exploring in-context learning for text diffusion is another interesting topic.

Weakness 4: Throughput Comparison

Thank you for your suggestion. In Appendix L713-L714, we provide a detailed description on T for in Table 1. Specifically, we utilize T = 1 for digit multiplication, T = 2 for the boolean logic dataset, and T = 64 for grade school math. It is worth noting that adjusting the parameter T itself is an advantage of DoT, as it allows us to allocate computational resources more efficiently to challenging tasks, while using a smaller T for simpler tasks. We show how T affects performance on grade school math in Figure 3, and below we also show how T affects throughput for Plaid DoT$^{MP}$. The relationship between throughput and T appears to be nearly linear.

Q1: comparisons with answer-only and CoT for diffusion language models

Please see the response for weakness 1.

Limitation: The need for more reasoning steps as tasks become more complex could be further elaborated upon.

Thank you for your suggestion. The main point of the CoT paper is to improve the reasoning ability by involving more intermediate steps. In DoT, we also observed that more diffusion timesteps (computing FLOPs) yield better results. [1] also presents a similar idea, and this topic is worthy of further investigation.

[1] Pondernet: Learning to ponder. ICML 2021.

We sincerely thank Reviewer 5GKr for your review and are grateful for the time you spent on our submission. We are also glad you think our paper is novel and significant. Below we would like to give detailed responses to each of your comments.

Weakness 1: The presentation regarding color and figure rendering can be enhanced

Thank you for bringing to our attention the potential confusion regarding color and figure rendering. We will address and clarify this issue in the final version of the paper.

Weakness 2: Comparison to larger open language models (e.g., Llama)

Thank you for your suggestion. We add the results for (LoRA) fine-tuning LLMs on the same dataset, listed in the following table. Please note that the current diffusion pretrained model is much smaller than Llama 7B, so this comparison is not fair and we just list them for reference. We have validated that our DoT is better than the same scale autoregressive model GPT-2 (Table 1), which shares the similar architecture with Llama. We believe that further exploration of diffusion language models will lead to larger models that can compete with current LLMs, allowing DoT to achieve results more comparable to Llama.

Q1: About the confusion in line 155

Thank you for bringing up this question. For the first “former”: in the inference stage of diffusion, the timestep t starts from T and progressively decreases to 1, so we refer the larger time $u$ to “former” steps. Regarding two “forward” words in this line, this term carries different meanings: the first “forward” refers to the forward process in diffusion, which involves adding noise to data, while the last “forward” denotes the forward pass of the model, in contrast to the backward gradient-backpropagation pass. We will avoid using the same word with different meanings to prevent misunderstandings.

Q2: A comparison to the paradigm in [1] could be informative.

Thank you for sharing this paper ‘Visual CoT of diffusion models’. This paper borrows the idea of CoT in LLMs which involves intermediate steps to improve the performance. Their model acts in two steps: first generate the CLIP embedding and then generate the final image. This paper and our DoT both mentioned the CoT in diffusion models, but there is a big difference: the mentioned paper only borrows the idea of CoT but cannot perform the CoT reasoning, while DoT focuses on the reasoning ability of text models, as the alternative to autoregressive CoT in LLMs. We will add this comparison in related work.

We sincerely thank Reviewer E7FJ for your review and are grateful for the time you spent on our submission. We're pleased you find our method effective. Below, we provide a point-by-point rebuttal to clarify your concerns.

Weakness 1: Discussion of recent work

Thank you for sharing this paper Can mamba learn how to learn?. We noticed that there are several differences between this paper and DoT. First, mamba is a new model architecture as an alternative to traditional Transformers with full attention. For our diffusion models, we currently use traditional Transformers architecture, and it is orthogonal to the design of mamba. It would be interesting to see DoT’s performance with mamba as the base model. Second, the mamba paper mainly discusses the ability of in-context-learning, while our experiment setting focuses on the chain-of-thought reasoning. In all experiments except for few-shot ChatGPT baselines, we didn’t use in-context demonstrations. Exploring ICL ability of diffusion models is another interesting topic.

Weakness 2: Standard deviation

Thank you for bringing up this point. All experimental results were obtained by averaging the results of 3 separate trials, with a confidence interval of p < 0.01. Experimental results also reveal significant disparities in accuracy among different models.

Question 1: Rationale example in Figure 2

Thank you for bringing up this question. In the problem-solving tasks chart, we have two rationales and one final answer, 3 CoT steps in total <<2/2=1>><<2+1=3>>####3. For AR models, it will generate each token one by one. For single-pass DoT, it will generate the whole CoT steps in parallel: <<2/2=1>><<2+1=3>>####3 (Table 3). For multi-pass DoT, it will generate <<2/2=1>> first in parallel, then <<2+1=3>>, and then ####3.

Question 2: Comparison with no-DoT finetune

Thank you for your suggestion. We conduct the answer-only setting to further validate the effectiveness of DoT. The result table reveals that fine-tuning diffusion models solely with answer data leads to inferior performance compared to DoT, mirroring the degradation of AR models in the absence of CoT.

Limitation: Ablation study and general performance

Thank you for your suggestion. We have listed the comparison with the no-DoT fine-tune and will add it to the paper. The current ablation experiments further validate DoT’s effectiveness.

In this work, we mainly focus on the reasoning ability of models including both logical reasoning and mathematical reasoning. For more general testing like performing a variety of complex tasks such as for ChatGPT, further advancements are still required to enhance the scalability of diffusion language models for general ability, as described in our limitation section.

We sincerely thank Reviewer abfr 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.

Weakness 1: Simple datasets

The reasoning ability contains arithmetic reasoning and logical reasoning. From the experimental results, we validated that DoT performs as well as GPT models after fine-tuning on the boolean logic dataset. Yet, arithmetic reasoning such as grade school math is a more challenging task for all models. Besides, by introducing the relatively simple datasets, we aim to demonstrate that DoT not only performs well on relatively simple tasks but also exhibits a higher efficiency compared to GPT models.

Weakness 2: Throughput

Thank you for bringing up this point. When we show the throughput of DoT is superior to AR-CoT, what we want to emphasize is not purely throughput but its flexibility. As a hyperparameter, diffusion timesteps are determined through a held-out validation set. The number of time steps is influenced by the complexity of the task at hand. For instance, in simpler tasks such as digit multiplication, only a small number of timesteps is sufficient to obtain desirable performance. However, in more challenging reasoning tasks like GSM, we can increase the number of timesteps to enhance the performance. In this case, the final throughput is not superior to AR-CoT but we can achieve better results. In other words, we can spare more time on “thinking” on complex tasks (this interesting idea is introduced in Section 4.4 and Figure 4a). We argue that this flexibility is exactly the advantage of DoT over auto-regressive CoT models.

Moreover, AR models can utilize key-value caching to enhance throughput during generation, but they still decode token-by-token for longer outputs, whereas diffusion models operate with fixed timesteps for longer context. Also, by removing kv-caching memory, the inference of diffusion models can be conducted with larger batch size, as mentioned in the SEDD paper [1]. In our paper, we highlight the potential efficiency advantage of DoT in our tested cases, and we believe that the efficiency of diffusion models in the wild is another interesting topic and worthy of further investigation.

[1] Discrete Diffusion Modeling by Estimating the Ratios of the Data Distribution, ICML 2024.

Weakness 3: novelty of the methodology

In terms of various approaches that can be used in the fine-tuning stage of diffusion models, DoT is non-trivial. Other alternatives include directly fine-tuning Plaid 1B and using back-gradients to control the generation, and directly fine-tuning DoT models initialized with GPT2 parameters as shown in the Appendix. Compared with them, the current DoT model stands out as the most effective. These empirical findings have not yet been addressed in any other existing literature.

Moreover, we propose the multi-pass variant of DoT and two sampling strategies to further enhance the performance, which we believe are significant in the diffusion realm.

Finally, our paper also extensively discussed the potential advantages of DoT over autoregressive models particularly for reasoning tasks, such as the reasonability-efficiency trade-off, and the self-correction, which have not been explored to the best of our knowledge.

Q1: Connection between the first sentence of Section 3.2 and Table 2

Thanks for pointing out this potential confusion. The Plaid model is based on the continuous embedding space, and the model relies on gradient-based guidance to control the token generation. If we briefly follow this model to do the continue-training on the fine-tuning datasets (first row in Table 2), the performance is poor and we believe it is because the gradient-based guidance fails to do accurate conditioning. This motivates us to use DiffuSeq-style training. Please refer to the comments for detailed examples.

Q2: Sampling temperature

In our experiment, we found that enabling softmax temperature would slightly decrease the quality compared to greedy decoding (accuracy gap within 1%). But the diversity can be enhanced. As a result, we find a noticeable performance boost (4% on Plaid and 12% on SEDD) after performing self-consistency marginalization. We set the softmax temperature to be 0.5 according to the validation set in experiments.

Q3: Variable lengths

Yes, we use [PAD] to control the output length. The model will generate 128 tokens in total but all the latter [PAD] will be removed.

Q4: Connection between schedule sampling and self-conditioning

Thank you for bringing up this great question. The similarity between our schedule sampling and self-conditioning is that they both consider the condition of the model predicted sequence. However, they serve different purposes and are complementary in nature. The goal of self-conditioning is to use the previously estimated $\tilde x_0$ as additional feature besides the original $x_t$, so the network models $f(x_t,\tilde x_0,t)$, where $x_t$ is always corrupted from the oracle data. While the goal of our scheduled sampling is to add inference-time noise to $x_t$ to be consistent with the inference stage. So the network learns to model $f(\tilde x_t,t)$. They are based on different purposes, and can be used together. For example, we can potentially model $f(\tilde x_t,\tilde x_0, t)$ by providing $\tilde x_0$ as an additional feature.

Q5: Details for throughput results

For GPT models, we use KV-caching when decoding. For all models, considering both the small model size and context size, we didn’t use flash-attention.

Q6: Details for GSM-Aug dataset

Following the dataset setting in paper implicit-CoT, we keep the natural language in problem description but remove the natural language in the CoT response and only keep symbolic expressions in <<>>.