Jump Your Steps: Optimizing Sampling Schedule of Discrete Diffusion Models

Authors: Yong-Hyun Park, Chieh-Hsin Lai, Satoshi Hayakawa, Yuhta Takida, Yuki Mitsufuji
Primary area: generative models
Keywords: Discrete diffusion models, Efficient sampling

Abstract:
Diffusion models have seen notable success in continuous domains, leading to the development of discrete diffusion models (DDMs) for discrete variables. Despite recent advances, DDMs face the challenge of slow sampling speeds. While parallel sampling methods like $\tau$-leaping accelerate this process, they introduce _Compounding Decoding Error_ (CDE), where discrepancies arise between the true distribution and the approximation from parallel token generation, leading to degraded sample quality. In this work, we present _Jump Your Steps_ (JYS), a novel approach that optimizes the allocation of discrete sampling timesteps by minimizing CDE without extra computational cost. More precisely, we derive a practical upper bound on CDE and propose an efficient algorithm for searching for the optimal sampling schedule. Extensive experiments across image, music, and text generation show that JYS significantly improves sampling quality, establishing it as a versatile framework for enhancing DDM performance for fast sampling.

Decision: accept (poster)
Reviewer scores: 5, 5, 8, 8

Reviews and discussion:
<thread>
creation date:
1730647453781

thread_type:
official_review

<thread_object>
creator: reviewer_vtwh

<content>
summary:
The paper describes a way to optimize sampling for discrete diffusion models. The idea is to select a good subset of optimal sample times based on an optimization criterion. Some theory and some experiments are given.

strengths:
The direction of this research is new and interesting.

weaknesses:
Some of the theory parts confounds heuristic and rigorous justification (by ignoring some long-time dependencies without actually emphasizing that this is a working simplification). Thus, some assumptions and/or limitations are not highlighted, which is making the paper look weaker. ----[update] this was addressed in the discussion, with reference to appendices

Also, I am not completely convinced with the interpretation of some of the experiments, but this may be just me. --- [update] the authors clarified and promised to amend some figure

questions:
1) line 189 saying "motivated by this observation" is not correct/precise I think, given that you list a bunch of factors at work and you just focus on one? the only motivating point is the dependence on sampling methods, and not the others.

2) line 202 "if we ignore the accumulated error from the previous steps": why should we ignore it? can you formulate this as an assumption somewhere?

I think that the paper's theory part might have very non-sharp estimates in case this error can't be ignored, can you convince me otherwise?

3) can you comment on when is the inequality in (7) sharp? I think this has to do with question 2.

4) line 216-217 "this theorem suggests" -- I think that it'd be accurate to say "based on the case that inequality (7) is sharp, we are tempted to hypothesize" and to mention that in case accumulated error from previous steps is negligible, then this heuristic becomes stronger.

Note that in applications in which the gap in (7) is large, then theorem 3.1 gives a lax and uninformative upper bound and this heuristic has no reason for holding.

5) about Theorem 3.2, when is the inequality in (8) sharp? Upper bounds like KLUB are useful to the extent to which they are close to sharp. So this limits the range of applicability of (8), and therefore it would be good to have a description of cases in which this is sharp.

6) about the hierarchical time breakdown (described in paragraph 3.4 and summarized in Figure 3) : this is one possible time breakdown strategy. Can you comment on alternative strategies, or on how this one is better than, or equivalent to, other strategies? 

7) line 298-299 what do you mean by "computational complexities"? This is vague and easy to misinterpret for a reader. I think you mean the thing mentioned in lines 350-360? If so, I suggest moving those lines up here, so that the reader has a concrete thing in mind since the beginning.

8) about the formula/equation at lines 308-309: what does the "approximately equal" sign actually mean? and why or under what hypotheses is this approximation valid? You mention something in the preceding paragraph, but this is not giving a clear cut answer to what approximation you have in mind, and on what are the underlying assumptions behind it.

9) about figure 9, I am puzzled/perplexed by the cases (a) and (c): 
In (a), why do you think that the JYS schedule has so dense subdivision in the first interval? I don't see anything in the "ground truth steps" graph that would indicate that this is optimal. 
In (c), why is the JYS deciding to just not include any timesteps in the first interval? In there, I think that the Ground Truth Steps part has a very dense time dependence, at least that's what I see in the figure.. so how do you justify this discrepancy?

To summarize, both points (a) and (c) show that JYS has high discrepancy in its "optimized" timestep choice, compared to "ground truth samples", so it would be good to understand what this means, how do we interpret this.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730661607165

thread_type:
official_review

<thread_object>
creator: reviewer_ibjs

<content>
summary:
This paper introduces Jump Your Steps (JYS), a method for discrete diffusion models (DDMs) that enhances generation quality by finding optimal sampling schedules when using efficient parallel sampling methods. Parallel sampling methods such as τ-leaping can speed up DDM inference, but they introduce a Compounding Decoding Error (CDE) that demotes sample quality. The authors address this by deriving a practical upper bound on CDE and developing an efficient schedule optimization algorithm that minimizes this bound without additional computational overhead. They perform experiments on image, music, and text generation tasks to show that the JYS approach improves sampling quality across different transition kernels using fast sampling methods like τ-leaping or k-Gillespie.

strengths:
The main idea of this paper is novel: the authors suggest reframing DDM sampling schedule optimization as a CDE minimization problem, introducing a Kullback-Leibler divergence upper bound (KLUB) metric to make the optimization tractable for improving sample quality in fast DDM sampling methods. The theoretical development seems sound, establishing connections between CDE minimization and path-space probability measures. This work is particularly interesting because it requires no architectural modifications or additional inference costs, unlike existing methods. As for the empirical validation, the paper includes an evaluation across diverse domains showing consistent improvements in the sample quality when using the same number of function evaluations (NFEs) as a baseline approach.

weaknesses:
1. The empirical analysis can be improved: 
 - While improvements in the generation quality in terms of FID or perplexity scores are shown, comparisons against other quality-improvement methods (e.g., predictor-corrector approaches mentioned in Section 1), and an explicit study of their computational costs are missing. The paper should either include such comparisons or justify why they aren't necessary.
 - The trade-off between schedule optimization time (or its memory requirements/end-to-end generation time) and sample quality gains isn't studied, making it difficult to assess the method's practical utility.
- No error bars or standard deviations are provided for the quantitative metrics. This is important since generative models can have high variance in their outputs and metrics. An elaboration of the significance of the reported results would be appreciated. 
-- The comparison is limited to a uniform scheduling baseline. Would it be possible to compare against an adaptive step size integrator method (Runge–Kutta method modifications if they are applicable to the discrete case) or other works like a recent one [5]? It’d be great to include this or explain why such comparisons aren't applicable or necessary.
2. The paper doesn't clarify how the diffusion model formulation (Eq. 1) relates to other DDM variants from [2, 3, 4]. Such discussion would be appreciated.
3. The single-state transition assumption, while mathematically convenient, limits the method's applicability: it can lead to suboptimal schedules for fast-changing processes. It’d be great to see ideas on possible extensions to handle multiple transitions while maintaining computational tractability.
4. Some technical details are difficult to follow or verify. For example, the authors introduce two techniques for KLUB computation in Section 3.5 but don't explain when each is preferable. It looks like both techniques are combined in Algorithm 1. No discussion of whether the hierarchical breakdown strategy converges to globally optimal schedules. The clarity of the paper could also be improved. There are suggestions on this in the Questions section.
5. No code is provided for reproducibility. 


[1] Hoogeboom, Emiel, et al. "Argmax flows and multinomial diffusion: Learning categorical distributions." Advances in Neural Information Processing Systems 34 (2021): 12454-12465.

[2] Santos, Javier E., et al. "Blackout diffusion: generative diffusion models in discrete-state spaces." International Conference on Machine Learning. PMLR, 2023.

[3] Austin, Jacob, et al. "Structured denoising diffusion models in discrete state-spaces." Advances in Neural Information Processing Systems 34 (2021): 17981-17993.

[4] Hoogeboom, Emiel, et al. "Autoregressive diffusion models." arXiv preprint arXiv:2110.02037 (2021).

[5] Chen, Yuzhu, et al. "Adaptive Time-Stepping Schedules for Diffusion Models." The 40th Conference on Uncertainty in Artificial Intelligence.

questions:
1. As the empirical validation of the KLUB is a proxy for CDE, would it be possible to directly show that minimizing KLUB actually reduces this error?
2. Does the hierarchical breakdown strategy guarantee convergence to a globally optimal schedule?
3. The paper says “After K iterations, this hierarchical strategy yields a sampling schedule with 2^K NFEs, optimizing the schedule as the number of steps increases.” in lines 289-29. So, there is an exponential factor appearing in the number of function estimations. I wonder if it affects the overall complexity and method’s applicability, especially in resource-constrained settings.
4. I also wonder if the computational gain from this approach leads to higher memory costs. Is it this case from your observations?
5. Is there any intuition on how the method's error bounds scale with data size, or it doesn’t matter?
6. Is there any insight into the applicability of the proposed method to DDMs without fast sampling algorithms?
7. Do you think the hierarchical breakdown strategy converges to globally optimal schedules?

Small suggestions regarding the manuscript:
- There is no need to define DDMs (lines 12, 30, 94, for example) or CDE (line 156) a few times through the text if they were defined before.
- In Eq. 1 and this section, I’d suggest explicitly defining $x$ (data) and $y$ (corrupted data) variables.
- Line 133: Should it be “.” instead of “:”? 
- Maybe notations in line 148 can be combined with Section 2.1 or moved there? I’d suggest the authors go through their manuscript to make sure all notations are introduced and consistent.
NFE is not defined in line 290 or earlier. Also, it is useful to explain what it means (it is the number of calls to the score NN, right?).
- Maybe it’s worth moving Section 5 to the Intro section or a separate section at the beginning of the manuscript? Also, mentioning papers [1, 2] might be helpful.
- Adding equation numbers to Algorithm 1(line 7) would be helpful.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730715150893

thread_type:
official_review

<thread_object>
creator: reviewer_tf1f

<content>
summary:
This paper introduces Jump Your Steps (JYS), a novel approach that optimizes the allocation of discrete sampling timesteps. It achieves this by minimizing Compounding Decoding Error (CDE) without incurring any additional computational cost. The authors derive an upper bound on CDE and employ techniques to simplify the computation. Based on the upper bound, they propose an algorithm for searching for the optimal sampling schedule, thereby enhancing the sampling quality.

strengths:
The algorithm presented in this submission is straightforward and simple. Experimental results have been provided to demonstrate the feasibility of the proposed algorithm.

weaknesses:
1. The novelty of the proposed algorithm is rather limited. The authors merely transfer Align Your Steps (AYS), which is originally for image generation, to the framework of SEDD for text generation.

2. The experiments conducted on the CIFAR-10 and Countdown datasets lack practical significance. In the realm of image generation, efficient sampling methods have been well explored. For CIFAR-10, even with 256 NFEs, the proposed approach can only achieve an FID higher than 10. In contrast, some popular sampling methods for image generation, such as DPM-Solver++ (https://arxiv.org/pdf/2211.01095) and DMD (https://arxiv.org/pdf/2311.18828), can attain an FID less than 4 using only 10 or even just 1 sampling step. Additionally, the Countdown dataset is a synthetic one. I suggest that the authors place the results for CIFAR-10 and Countdown in the appendix. Moreover, as I am not familiar with music generation, it appears to me that although JYS can lead to a significant drop in Hellinger distance, the Hellinger distance for uniform sampling timesteps is already small (less than 0.4).

3. It seems to me that the practically meaningful experimental results are mainly showcased in Figure 8 for text generation. Nevertheless, upon examining Figure 8, I cannot observe a significant improvement over the use of uniform sampling timesteps.

4. Some problems in the writing:
- The right arrows used on Line 52 seem strange. Additionally, such right arrows occur multiple times in the main content. It is suggested to use normal notations instead (and I don't think there is a need to write down the whole sequence for timesteps repeatedly).
- The term "KLUB" is first seen in Figure 2 on page 4; however, its explanation is only provided in Theorem 3.1 on page 5.
- There is no necessity to write down the Sampler in Figures 6 and 7.
- In the section on "$k$-Gillespie's Algorithm" (Line 125), '$k$' apparently means updating $k$ tokens in parallel. However, in Appendix B.1 KLUB COMPUTATION, the use of '$p(t|k)$' is presented and it is explained that this is determined by the pretrained noise schedule. As a result, I am not certain whether '$t$' signifies time. If it does, then it appears inconsistent with '$k$'. Please provide a more detailed explanation about $k$-Gillespie's Algorithm.
- The condition $\mathbb{P}_{t_0} = \mathbb{Q}_{t_0}$ in the last sentence of the proof of Theorem 3.1 should be presented as an assumption in the statement of Theorem 3.1.


--------

After rebuttals:

I would like to thank the authors for their detailed responses. However, the authors have only provided Figure 13 (which has been shifted to illustrate the performance gains of JYS compared to those achieved by increasing model size) to address my major concern in W3. Currently, although I have slightly raised my score to 5, my concern about the practical significance and relevance of the proposed approach still persists. As a result, I am still inclined towards rejection.

questions:
N/A
</content>
</thread_object>

</thread>
<thread>
creation date:
1730718964452

thread_type:
official_review

<thread_object>
creator: reviewer_tvar

<content>
summary:
This paper proposes a method, Jump Your Steps, to optimize the discretization schedule of discrete diffusion models, by minimizing the compounding decoding error (CDE). The authors provide a KL-divergence upper bound (KLUB) for different continuous-time Markov chains of discrete diffusion models and propose techniques to efficiently approximate the timesteps minimizing KLUB.

strengths:
This paper proposes an approach to improve the sampling quality and efficiency by optimizing the sampling schedule. This paper is well-written and easy to follow. The sampling and acceleration of diffusion models is an interesting topic for the community.

weaknesses:
- Theorem 3.1 seems to be standard entropy bounds, and Theorem 3.2 seems to directly follow from Equations 3.2 and 3.4 of Ding & Ning (2021) and Theorem 3.2 of Sabour et al. (2024). The authors should illustrate their technical contributions and the novelty of their theoretical results with a comparison to the previous literature. 
- The authors only compare the JYS schedule with the uniform schedule. Additional experiments on other schedules, e.g. EDM, Linear LogSNR, and Cosine LogSNR, should be presented.

questions:
- The proposed scheduling strategy only considers $2^K$ numbers of function evaluations (NFE). How can the schedule extend to other values of NFE? 
- In Figure 5, is there any intuition to explain the error increase of the $k$-Gillespie sampler with the JYS schedule when NFE=16? 
- In Figure 10, the comparison between uniform and JYS is not significant for NFE=32, 128, or 256. Please consider providing additional experiments on other image datasets, e.g., toy 2D datasets in Sabour et al. (2024).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732279306342

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank each reviewer for their careful reading of our work and their thoughtful feedback. 

We appreciate the positive remarks, including that our paper is well-written and easy to follow (tVar), straightforward and simple (TF1f), and novel (ibjS, VTwh) with sound theoretical developments (ibjS). We are also grateful for the recognition of our approach, which introduces performance improvment without no architectural changes or added inference costs (ibjS).

We have updated the manuscript following the suggestions in the reviews and marked the additions in blue to facilitate finding these passages. All figures mentioned in the global and individual reviewer rebuttals refer to those in the revised PDF. 

---

### 1. Difference from align your steps

> **Reviewer tVar and TF1f asked us to clarify the distinct contribution of our work compared to Align Your Steps (AYS) [1].**

We are pleased to provide clarification and emphasize that our contributions extend beyond a simple adaptation of AYS.

First, our paper presents an original theoretical contribution. In Theorem 3.2, we derive an upper bound for the sampling error of discrete diffusion models (DDMs), referred to as KLUB, using Girsanov's theorem for continuous-time Markov chains (CTMC) from Ding & Ning (2021) [2]. To the best of our knowledge, this is the first study to characterize the sampling error of DDMs using KL divergence. Notably, our result offers a key insight: the sampling error of DDMs can be elegantly expressed as **a ratio of the rate matrix**, introducing a novel framework for error analysis.

We respectfully point out that results from continuous diffusion models (CDMs) cannot be directly transferred to DDMs because they are based on fundamentally different mathematical frameworks. Specifically, the KLUB for stochastic differential equations (SDEs) does not directly apply to continuous-time Markov chains (CTMCs). In SDEs, the KLUB derived by Chen et al. (2022) [3] depends on the difference between the drift terms of two SDEs (i.e., $||f_1 - f_2||_2^2$). In contrast, for CTMCs, the KLUB is defined by the ratio of the rate matrices between two CTMCs (i.e., $\log (R_1 / R_2)$), which cannot be obtained through simple adaptation. While AYS [2] demonstrates how KLUB can be utilized for optimizing sampling schedules, our work shares only the general idea of using KLUB for sampling schedule optimization; the theoretical development and application are entirely different.

Second, we provide a distinct technical contribution that sets our work apart from AYS. In CDMs, KLUB can be directly computed numerically, whereas directly calculating KLUB for DDMs is infeasible with practical computational resources. To address this challenge, we developed hierarchical breakdown strategy (Section 3.4) with Techniques 1 and 2 (Section 3.5), leading to highly efficient optimization algorithms (Algorithms 1 and 2). This advancement not only makes sampling schedule optimization practical for DDMs but also significantly reduces computational costs. For example, for CIFAR-10, while AYS requires approximately 6 GPU hours to optimize the sampling schedule for NFE=50, **our approach completes the optimization in under 45 seconds for NFE=64**. We discuss our algorithm's efficiency in greater detail in Global Rebuttal #2.

Lastly, we would like to emphasize our focus on addressing the unique challenge of accelerating DDMs. While CDMs benefit from established numerical solvers like DPM and DPM++, which greatly improve sampling efficiency and image quality, such solutions are not available for discrete data. In discrete settings, there are no equivalent particle-wise differential equation solvers suitable for DDM sampling. Therefore, we analyzed the specific problem of Compounding Decoding Error (CDE) inherent in DDM sampling and proposed a principled, mathematically grounded solution. Consequently, we developed an alternative approach: optimizing the sampling schedule to mitigate CDE. We have tested our method across various modalities and transition kernels, demonstrating its broad applicability and effectiveness.

[1]: Sabour et al., Align Your Steps: Optimizing Sampling Schedules in Diffusion Models, ICML24
[2]: Ding & Ning, Markov chain approximation and measure change for time- inhomogeneous stochastic processes, Applied Mathematics and Computation, 2021
[3]: Chen et al., Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions, ICLR23
</content>
</thread_object>

</thread>
<thread>
creation date:
1732523626841

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We sincerely thank the reviewers and area chair for their efforts and would like to request any missing responses to our initial rebuttal.

Given the extensive experiments in our rebuttal, we understand that reviewing the revised manuscript may be challenging. To assist, we are providing a brief summary of the newly added appendix sections:

- **Appendix A.1**: Limitations of JYS (NFE = 2^K, single-state assumption)
- **Appendix A.2**: Discussion on JYS design choices (Ignoring Accumulated Error, Golden Section Search, Hierarchical Breakdown)
- **Appendix B.1 & B.2**: Equality conditions of Theorems 3.1 and 3.2, Tightness comparison with ELBO
- **Appendix B.4**: Derivation for Technique 1 approximation
- **Appendix C.1**: JYS optimization for k-Gillespie
- **Appendix C.3**: Comprehensive explanation and algorithm for k-Gillespie
- **Appendix E.1**: Analysis of JYS computation cost
- **Appendix E.2**: Error bars and standard deviation for evaluation metrics
- **Appendix E.3**: Comparison of JYS performance improvement with model size scaling
- **Appendix E.4**: Comparison with other heuristic sampling schedules
- **Appendix E.5**: Qualitative results on Toy 2D dataset

We are grateful for your continued support and are happy to provide any further clarifications if needed.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This work addresses the challenge of slow sampling speeds in discrete diffusion models (DDMs) by introducing Jump Your Steps (JYS), a novel method to optimize discrete sampling timesteps. JYS minimizes Compounding Decoding Error (CDE), a key issue in parallel sampling methods like τ-leaping, which degrade sample quality. The authors derive a practical upper bound on CDE and propose an efficient algorithm to determine an optimal sampling schedule, all without additional computational cost. Experiments across image, music, and text generation demonstrate that JYS significantly improves sampling quality. Following the author response and author-reviewer discussions, the reviews are still mixed. Considering the substantial theoretical contributions, which outweigh the relatively weaker empirical experiments, I recommend to accept.

In the final version, I encourage the authors to include a discussion on DNDM, proposed by [Chen et al., 2024]. This closely related work shares a similar concept of utilizing predetermined transition times (i.e., allocating discrete sampling timesteps) to accelerate the sampling process in discrete diffusion models.

[Chen et al. 2024] Fast Sampling via Discrete Non-Markov Diffusion Models with Predetermined Transition Time