We thank the reviewer for your valuable comments and positive assessment for our work. Below, we address the questions and concerns you have raised.

Q1. Details on the implementation of the rectification step (Eq. 6) and the training procedure with the new coupling are insufficiently described. Based on the experimental section, it appears that a distribution is computed by multiple sampling pairs, but this process is not clearly explained.

A1. We appreciate the reviewer’s comment. To obtain the new coupling by Eq.6, we first sample the initial point from the prior distribution by $X_0\sim p(X_0)$ and apply the iterative denoising using the pre-trained model to obtain a sample from the terminal distribution by $X_1\sim p_\theta(X_1|X_0)$. The obtained pair $(X_0, X_1)$ defines the new coupling i.e., $(X_0, X_1)\sim \pi_{k+1}(X_0, X_1)$, which derives the new conditional transition probability in Eq.2. The model $p_\theta$ is then trained with the updated transition probability with the same denoising objective, which corresponds to a rectification step.

We will add this detailed description to the preliminary section (Section 3.1) and link it to Eq. 6 to ensure clarity and completeness.

Q2. What is the computational cost of a single rectification step? A comparison of the overall training time with other baseline methods would help assess the efficiency of the proposed approach.

A2. We appreciate the comment. While our rectification process requires sequential training, we clarify that each rectification step is much more efficient than distillation for the following reasons. First, as discussed in Section 4.4. and Figure 4(b) in the main paper, the rectification process can be trained with only a small portion of the entire training data (50K images vs. 1M full training data). As a result, it greatly reduces the cost of forwarding pre-trained models and convergence speed of each rectification step. Second, unlike distillation approaches, our method requires only the student model during training, avoiding the cost of operating two models simultaneously.

Table A compares the training time of ReDi with the teacher model and other distillation approaches. It shows that the entire rectification process for ReDi$^3$-distill requires only ~3% of the time needed to train the initial teacher model, and faster training than distillation methods up to three rectification steps. We appreciate the reviewers’ comment and will incorporate this result in the paper.

Table A.

We thank the reviewer for your valuable comments and positive assessment for our work. Below, we address the questions and concerns you have raised.

Q1. In Section 3.1 you adopt a unified viewpoint … such miss leaves outside an interesting connection of rectified flows to IMF procedure that solves SB problem $Appendix A.2, 3$ with recent works on discrete space versions of IMF $4, 5$ .

A1. We thank the reviewer for these constructive comments and insightful suggestions. We agree that incorporating the suggested connections will significantly enrich the paper's context and clarify its position within the broader field.

In the revised manuscript we will clarify the unified viewpoint in discrete generative modeling and discuss the connection to Schrödinger Bridge Problem (SBP) and Iterative Markovian Fitting (IMF). Specifically, we will highlight the procedural similarities between our ReDi method and the IMF algorithm, noting that our iterative training and data generation steps are analogous to the Markovian and reciprocal projections of IMF. We will also briefly introduce recent works that apply the SBP/IMF framework to discrete domains as the reviewer suggested. We believe these additions will provide valuable context and strengthen the paper's theoretical foundation. We thank the reviewer again for this valuable feedback, which has helped us improve our manuscript.

Q2. In the text experiments, the authors do not assess diversity of samples.This is crucial, since there is a known pitfall of generative perplexity.

A2. We thank the reviewer for the insightful comment. Upon the reviewer’s suggestion, we measured the unigram entropy alongside the Generative Perplexity (PPL), following the methodology in $1$ . Aligned with the reviewer’s observation, we find that the entropy of the generated texts tends to decrease as we iterate more rectification steps. Upon closer inspection, we identified that this is inherited from the limited generative performance of the pre-trained model, perhaps attributed to its limited capacity (GPT-2-small backbone). Specifically, the generated distribution of the pre-trained model (PPL: 42.43 / entropy: 5.22) has a significant gap to the real data (PPL: 14.36, entropy: 5.44). When our rectification process is applied iteratively, this distributional gap can grow with each step, leading to low-diversity characteristics.

We believe that employing higher-capacity pre-trained models can easily address this problem. To demonstrate this given a limited rebuttal timeframes, we conduct an experiment on a subset of OpenWebText with 1M tokens, ensuring that our teacher model has sufficient capacity to fit the distribution. As shown in Table A, B, in this setting, our method achieves significant PPL improvements in the few-step regime (2-8 steps) while simultaneously preserving high entropy. This demonstrates that when the base model is faithfully modeling data distribution, our method effectively improves few-step generation quality without sacrificing diversity.

We deeply appreciate the reviewer’s comment and will update the experiments in the paper with larger pre-trained backbones and full OpenWebText dataset.

Table A.

Table B.

$1$ Dieleman, Sander, et al. “Continuous diffusion for categorical data.”

Q3. Do not mention some important references that cover the problem of the factorization. $1, 2$ ; Why do you directly contrast the proposed method with distillation approaches?

A3. We thank the reviewer for suggesting these important references and for the insightful question regarding our choice of comparison.

The main reason that we choose the distillation approaches as baselines is because they share the same goal and settings with our method: improving the few-step generation performance of the pre-trained discrete diffusion model. While the suggested papers ( $1, 2$ ) are highly relevant for tackling the same factorization problem, they operate under different assumptions or with different goals, which makes a direct experimental comparison challenging:

• Discrete Copula Diffusion $1$ , like our work, aims to solve the factorization problem for few-step generation. However, it assumes the availability of an additional pretrained autoregressive model for a copula model. Our experimental setup does not include such a component i.e., we aim to directly improve the few-step performance of pre-trained discrete diffusion models without external models
• EDLM $2$ reduces the approximation error using an energy-based approach. While this can reduce the number of decoding steps, it is not optimized for sampling efficiency. As detailed in their algorithm, each decoding step requires multiple evaluations of the energy-based model. This fundamentally differs from our goal of a truly efficient one- or few-step generative process. Therefore, we believe that distillation-based models, which share our specific goal of improving a base model's sampling efficiency without these additional components or computational costs, represent a more direct and appropriate comparison group for our work.

Moreover, we agree that both papers are highly relevant as they address the same core problem of factorization error in discrete diffusion. We will add a detailed discussion of these works and a comparison of their methodologies to our Related Works section to better contextualize our contribution. We thank the reviewer again for these excellent suggestions.

$1$ Liu, Anji, et al. "Discrete Copula Diffusion."

$2$ Xu, Minkai, et al. "Energy-based diffusion language models for text generation."

Q4. In lines 273-274 it is quite hard to understand how Conditional TC is approximated; please provide some mathematical representation for better understanding.

A4. We thank the reviewer for this question and are happy to clarify our estimation procedure.

We approximated the Conditional TC via Monte Carlo sampling, as the exact transition probability $p(X_1|X_0=x_0)$ is intractable. Our process involved generating 10 samples of $x_1$ for each of 5,000 initial states $x_0$. From these samples, we estimated the empirical conditional distributions ($\hat{p}(X_1|X_0)$ and $\hat{p}(X^i_1|X_0)$) by counting occurrences, which then allowed us to compute the Conditional TC. We will add this detailed description to the appendix.

Q5. In the main theorem, I suppose there are missing assumptions, such as factorized conditional probability can’t be zero.

A5. We thank the reviewer for this precise question regarding the assumptions for our main theorem. The reviewer correctly points out that the KL-divergence $D_{KL}(p||q)$ requires that $q$ must be non-zero on the support of $p$. In the context of the Conditional Total Correlation, $p$ corresponds to the conditional joint distribution $p(X_1|X_0)$ and $q$ corresponds to the factorized product of conditional marginals, $\prod_i p(X^i_1|X_0)$.

However, we would like to clarify that no additional assumption is needed here. This is due to a fundamental property of probability distributions: for any given condition $X_0$, if a point $X_1$ has a non-zero probability under the conditional joint distribution (i.e., $p(X_1|X_0) > 0$), then each of its corresponding conditional marginal probabilities $p(X^i_1|X_0)$ must also be non-zero. Consequently, the factorized term $q = \prod_i p(X^i_1|X_0)$ is guaranteed to be non-zero on the support of $p(X_1|X_0)$. Therefore, the requirement for the KL-divergence is naturally satisfied without needing an explicit assumption in our theorem. We will add a footnote to the paper to make this point clear.

Q6. In checklist question 16, you mention using an LLM for a “proof sketch”. What do you mean by “sketch”? Specifically, is the proof in the Appendix considered a sketch proof?

A6. We apologize for the ambiguity. To be clear, the proof in the Appendix is a complete and formal proof, not a sketch. Our use of the term "proof sketch" in the checklist was intended to describe using the LLM for background research on mathematical tools (e.g., the Data Processing Inequality). The LLM did not write or generate the derivation, and the proof is our own work.

We thank the reviewer for the question that leads us to clarify the theorem. We would like to recall Theorem 1 and provide a justification how applying $\Phi$ to $Q_k$ yields $Q_{k+1}$.

Theorem 1. Let $\pi_k(X_0, X_1)$ be a coupling at iteration $k$, and let $\pi_{k+1}(X_0, X_1)=p(X_0)p_\theta(X_1|X_0)$ be the "rectified" coupling obtained via the ReDi procedure at iteration $k$. Then it satisfies the following:

$$TC_{\pi_{k+1}}(X_1|X_0) \leq TC_{\pi_k}(X_1|X_0).$$

In the proof we assumed that $p_\theta(X_1|X_0) = \prod_i p_\theta(X_1^i|X_0),$ rectifying the coupling with one-step decoding process. In this case, $\Phi(Q_k)=Q_{k+1}$ holds as:

$$\Phi(Q_k)=E_{X_0, X_1}[p(X_0) \prod_{i=1}^N p_{\pi_k, 1|0}(X_1^i|X_0) \delta_{X_0}(X_0')p_\theta(X_1'|X_0)]\\$$ $$=p(X_0')p_\theta(X_1'|X_0')\\$$ $$=p(X_0') \prod_{i=1}^N p_\theta(X_1^i|X_0)\\$$ $$=Q_{k+1}.$$

While the proof is simple and general to other $f$-divergences that support data processing inequality, it is limited to one-step decoding process.

---

Below, we provide an extended proof for more general M-step decoding processes, which is defined as $p_\theta(X_1|X_0)=\sum_{X_{t_1}, ... X_{t_{M-1}}}\prod_k\prod_ip(X_{t_{k+1}}^i|X_{t_k})$. The following properties are convenient:

Property 1. (Pythagorean Inequality for KL Divergence [1]) For a distribution $q$ in a log-convex set of distributions $Q$, if $q^* = \arg \min_{q\in Q} D_{KL}(p||q)$ and $r \in Q$, then

$$D_{KL}(p||r) \ge D_{KL}(p||q^* ) + D_{KL}(q^*||r) \quad \text{(Eq.A)}$$

Extended Proof of Theorem 1:

We make two assumptions:

• Assumption 1. Let $P$ be the family of M-step decoding processes. We assume that our model $p_\theta(X_1|X_0)$ lies within the log-convex hull of $P$.

• Assumption 2. At each rectification step, our model is trained to minimize $D_{KL}(p_{\pi_k, 1|0}(X_1|X_0)||p_{\theta}(X_1|X_0))$. We assume that the model has reached the minimizer.

Assumption 1 is plausible as $P$ becomes close to the hypothesis space of $p_\theta(X_1|X_0)$ for sufficiently large steps $M$. Assumption 2 is in line with step-wise optimality assumptions used in, e.g., GANs [2].

The proof that $TC_{\pi_{k}}(X_1|X_0) \ge TC_{\pi_{k+1}}(X_1|X_0)$ proceeds as follows:

$$TC_{\pi_{k}}(X_1|X_0) = E_{X_0}[D_{KL}(p_{\pi_k, 1|0}(X_1|X_0)||\prod_ip_{\pi_k, 1|0}(X_1^i|X_0))] \quad \text{(Definition of Conditional TC)}$$ $$\ge E_{X_0}[D_{KL}(p_{\pi_k, 1|0}(X_1|X_0)||p_{\theta}(X_1|X_0))] + E_{X_0}[D_{KL}(p_{\theta}(X_1|X_0)||\prod_ip_{\pi_k, 1|0}(X_1^i|X_0))] \quad \text{(by Eq.A)}$$ $$\ge E_{X_0}[D_{KL}(p_{\theta}(X_1|X_0)||\prod_ip_{\pi_k, 1|0}(X_1^i|X_0))] \quad (\text{since }D_{KL} \ge 0)\\$$ $$=E_{X_0}[D_{KL}(p_{\pi_{k+1},1|0}(X_1|X_0)||\prod_ip_{\pi_k, 1|0}(X_1^i|X_0))] \quad \text{(by Eq.6 in the paper)}\\$$ $$=E_{X_0}[D_{KL}(p_{\pi_{k+1},1|0}(X_1|X_0)||\prod_ip_{\pi_{k+1}, 1|0}(X_1^i|X_0)) + \sum_iD_{KL}(p_{\pi_{k+1}, 1|0}(X_1^i|X_0)||p_{\pi_{k}, 1|0}(X_1^i|X_0))] \\$$ $$\ge E_{X_0}[D_{KL}(p_{\pi_{k+1},1|0}(X_1|X_0)||\prod_ip_{\pi_{k+1}, 1|0}(X_1^i|X_0))] \quad (\text{since }D_{KL} \ge 0)$$ $$= TC_{\pi_{k+1}}(X_1|X_0). \quad \square$$

We sincerely thank the reviewer for their insightful feedback, which has helped us strengthen our proof. We will incorporate this extended version into the revised manuscript.

[1] Wolfer, Geoffrey and Watanabe, Shun. "Geometric Aspects of Data-Processing of Markov Chains."

[2] Goodfellow, Ian J., et al. "Generative adversarial nets."

We thank the reviewer for your valuable comments and positive assessment for our work. Below, we address the questions and concerns you have raised.

Q1. Are there any previous works that attempt to measure the gap between the conditional transition probability and the factorized product of marginals?

A1. We thank the reviewer for the question. We interpret the term "measure the gap" as having two potential meanings—(1) to empirically quantify it, and (2) to theoretically characterize it with a mathematical concept. We will address both aspects below.

Regarding the empirical measurement of this gap, to the best of our knowledge, our work is the first to propose a method to quantify it. We accomplish this by estimating the Conditional Total Correlation (CTC) via Monte Carlo sampling.

Regarding the theoretical characterization, previous works have indeed conceptualized this gap using various metrics. For instance, Di4C $1$ used Total Variation to represent it. We also acknowledge that Discrete Copula Diffusion $2$ uses a concept related to CTC to characterize the factorization error of model. We would like to note that our key contribution is distinct in that we focus on the CTC of the data coupling itself (rather than the model), and introduce a rectification mechanism to empirically minimize this gap by updating the coupling.

$1$ Hayakawa, Satoshi, et al. “Distillation of Discrete Diffusion through Dimensional Correlations.”

$2$ Liu, Anji, et al. "Discrete Copula Diffusion."

Q2. Are there other metrics that can be used to characterize/measure factorization error? Why KL would be more appropriate than other methods?

A2. We thank the reviewer for this insightful question. We would like to answer it from a theoretical perspective and an empirical measurement perspective. We also wish to clarify that KL-divergence was not used as a direct training objective in our work, but rather as a tool for our theoretical and empirical analysis.

From a theoretical perspective, the reviewer is correct that other metrics can be used. Any $f$-divergence is suitable for characterizing the factorization gap, as our Theorem 1 relies on the Data Processing Inequality (Appendix, L358-359), which holds generally for all $f$-divergences i.e., our rectification process reduces all $f$-divergences between the joint and factorized distributions, including Total Correlation.

From an empirical measurement perspective, however, KL-divergence offers a crucial advantage in computational tractability. Our method for measuring the gap relies on Monte Carlo estimation. To compute the KL-divergence $D_{KL}(p||q)$, we only need to evaluate the term $p \log(\frac{p}{q})$ for the points sampled from the true distribution $p$. In contrast, estimating other f-divergences (such as JS-divergence) would require evaluating terms over the entire support of the factorized model $q$. In a high-dimensional discrete space, this is computationally infeasible.

Therefore, while many metrics are valid in theory, KL-divergence was the practical and well-motivated choice for our empirical analysis. We will add this discussion to the appendix to clarify our reasoning.

Q3. Though ReDi's performance is outstanding for 1 step generation, for 4 step generation, it is mostly comparable or a little behind other methods. Is there something that TC fails to sufficiently capture?

A3. We thank the reviewer for this crucial question regarding the relationship between Conditional TC and generation quality. The full performance metrics are provided in Table A.

Table A.

As shown in the table, we observe different trends in performance. In the multi-step regime, the 4-step FID begins to increase after $K=2$. In contrast, the 1-step FID consistently improves as TC decreases. This phenomenon arises because TC captures the factorization error but does not capture the discrepancy between our model's learned distribution, $p_\theta(X_1)$, and the true data distribution, $p(X_1)$. In an ideal scenario where the base model is perfect ($p_\theta = p$), each rectification step would improve generation quality. However, this discrepancy may exist in practice, meaning each rectification step can introduce a small distributional drift even as it reduces the factorization error.

The net effect of this trade-off is highly dependent on the number of sampling steps. In the multi-step regime, the factorization error is already relatively small. Here, the marginal benefits of further TC reduction can be overshadowed by the accumulated distributional drift. In contrast, in fewer-step regimes (e.g., 1-step generation) where the initial factorization error is much larger, the significant performance gain from reducing TC far outweighs the minor impact of the drift. This explains the outstanding 1-step performance.

We are grateful for this question and will add this detailed discussion and the corresponding data to the main paper.

Q4. Does each row in Table 1 correspond to a different CFG value?

A4. The reviewer is correct that we report the results with the best CFG value for each method in Table 1. To ensure reproducibility, we will add the below table in our appendix.

Table B.

Q5. Does ReDi with different inference step count work better with different CFG values?

A5. The reviewer is correct that the optimal CFG value varies with the number of inference steps. In response to the reviewer's request, we provide the full ablation study on this relationship below. This analysis will be added to the appendix as a curve, following the reviewer's helpful suggestion.

Table C. ReDi$^1$

Table D. ReDi$^2$

Q6. Could the authors kindly define delta function in Eqn(1)?

A6. We thank the reviewer for pointing out this omission. The '$\delta_x(y)$' in Equation (1) refers to the Kronecker delta function, defined as:

$$\delta_x(y) = 1 \text{ if } x=y\text{, else } 0.$$

We will add this definition to the manuscript.

Q7. How do you determine the stopping criteria for rectification?

A7. The number of rectification iterations ($K$) is determined empirically by evaluating the generation quality after each step. Specifically, we monitor both FID and IS for 4-step generation to find the best trade-off, as shown in Table E.

Table E.

Based on the trade-off between FID and IS, we identified $K=2$ as a "sweet spot" that provides a large IS gain without a significant FID penalty on 4-step generation. Accordingly, we reported the scores for ReDi$^1$ and ReDi$^2$ as they represent the most compelling performance points.

Q8. How do the time/compute requirements of ReDi compare with existing approaches?

A8. We thank the reviewer for this important question regarding computational cost. While our rectification process requires sequential training, we clarify that each rectification step is much more efficient than distillation for the following reasons. First, as discussed in Section 4.4. and Figure 4(b) in the main paper, the rectification process can be trained with only a small portion of the entire training data (50K images vs. 1M full training data). As a result, it greatly reduces the cost of forwarding pre-trained models and convergence speed of each rectification step. Second, unlike distillation approaches, our method requires only the student model during training, avoiding the cost of operating two models simultaneously.

Table F compares the training time of ReDi with the teacher model and other distillation approaches. It shows that the entire rectification process for ReDi$^3$-distill requires only ~3% of the time needed to train the initial teacher model, and faster training than distillation methods up to three rectification steps. We appreciate the reviewers’ comment and will incorporate this result in the paper.

Table F.

Q9. Why are different CFG levels used for training ReDi$^1$+ReDi$^2$ v.s. ReDi$^3$-distill?

A9. The CFG levels to define rectified coupling was determined empirically, and our ablation studies show it differs for multi-step rectification versus one-step distillation. Specifically, lower guidance (CFG 1-2) was optimal for direct multi-step sampling (Table G), whereas strong guidance (CFG 8) was required to generate the best coupling for the one-step distillation task (Table H). This analysis will be added to the appendix for clarity.

Table G.

Table H.

We thank the reviewer for your valuable comments and positive assessment for our work. Below, we address the questions and concerns you have raised.

Q1. What is the actual training cost versus knowledge distillation?

A1. We appreciate the comment. While our rectification process requires sequential training, we clarify that each rectification step is much more efficient than distillation for the following reasons. First, as discussed in Section 4.4. and Figure 4(b) in the main paper, the rectification process can be trained with only a small portion of the entire training data (50K rectification data vs. 1M full training data). As a result, it greatly reduces the cost of forwarding pre-trained models and convergence speed of each rectification step. Second, unlike distillation approaches, our method requires only the student model during training, avoiding the cost of operating two models simultaneously.

Table A compares the training time of ReDi with the teacher model and other distillation approaches. It shows that the entire rectification process for ReDi$^3$-distill requires only ~3% of the time needed to train the initial teacher model, and faster training than distillation methods up to three rectification steps. We appreciate the reviewers’ comment and will incorporate this result in the paper.

Table A.

Q2. Moreover, improvements in generation quality appear modest.

A2. We would like to kindly remind the reviewer that our primary goal is to improve the few-step generation performance by reducing the factorization errors in discrete diffusion models. As shown in Table 1 and discussed in Section 4.2, our method achieves significant improvements in the one-step setting compared to both the teacher model and other distillation baselines. While it shows comparable 4-step performance to competing methods, our method still exhibits much lighter memory for training compared to other distillation baselines that require additional teacher networks, and faster convergence using only a small portion of data for rectification (Fig 4(b) in the paper. please also see A1).

Q3. The method has not yet demonstrated scalability to larger datasets.

A3. We would like to clarify that the ImageNet and OpenWebText are widely used datasets to evaluate the image and discrete diffusion models. Although time and computational constraints prevented us from testing on larger datasets during the rebuttal phase, we believe that our method can easily scale to larger models or datasets for following reasons. First, our rectification process and its theoretical foundation are model-agnostic and are applicable on top of pre-trained discrete diffusion/flow models agnostic to their architecture and scale. Second, unlike distillation approaches that require both teacher and student models during training, our method does not increase the memory footprint beyond that of the original pre-training. We plan to open source our code to facilitate the future community efforts in scaling our approach.

Q4. Can this progressive coupling mechanism be extended to continuous flows? Why was the discrete setting prioritized?

A4. We would like to clarify the connection of our work to continuous flows and the motivation for prioritizing the discrete setting.

Although similar progressive refinement mechanisms have been explored in the continuous domain, as we note in our Related Works (Sec. 2.3), our work was prioritized for the discrete setting because the core challenges are fundamentally different. For continuous flows, the primary obstacle to few-step generation is the non-straightness of probability paths. In discrete flows, however, the concept of path straightness is not well-defined, as there are no intermediate states between distinct categories. The critical bottleneck we address is instead the factorization error inherent to these models.

Our key contribution is thus to characterize this specific factorization error and resolve it with our proposed rectification process. Interestingly, this process, despite being derived specifically to address the unique problem of discrete factorization, closely resembles the rectification methods used for the different challenges in continuous flows. We will add a discussion of this point to our method section to further clarify this distinction and connection.

Q5. How does generation quality (e.g., FID for images) improve compared to the baseline factorized model in the multi-step regime? Finally, how sensitive is performance to the number of coupling iterations $K$? Could $K$ be adaptively chosen?

A5. We would like to kindly clarify our model's performance in the multi-step regime and its sensitivity to the number of coupling iterations, $K$.

Regarding the generation quality versus the baseline in the multi-step regime, our method exhibits a trade-off between FID and IS, as shown in Table B. While the Inception Score improves compared to the baseline, the FID moderately degrades as $K$ increases. It is important to note that our primary contribution lies in the few-step regime, where the model shows significant gains from the rectification as demonstrated in Table 1.

Table B.

Concerning the sensitivity to $K$, this behavior can be understood as a trade-off between two competing factors. On one hand, each rectification process reduces the Conditional Total Correlation (TC), which significantly improves performance with fewer-step regimes where the initial factorization error is large. On the other hand, each rectification can introduce a small amount of distributional drift due to imperfections in the base model. In the multi-step regime, this accumulated drift can offset the benefits of TC reduction, leading to the observed performance changes.

Finally, regarding the adaptive selection of $K$, we agree with the reviewer that this is a viable approach. The general principle would be to monitor a chosen performance metric, such as FID or IS, after each rectification iteration and select the $K$ that provides the most desirable outcome based on that evaluation.