Single-Step Consistent Diffusion Samplers

Question 1

Could you explain the differences between CDDS and CD.

We appreciate the opportunity to clarify this point further. Once a sampler is available, our consistency distillation (CDDS) approach performs consistency training similar to the original method. However, the difference is how we generate intermediate noisy states for training:

• Generative Modeling Scenario: Intermediate state $X_{t_n}$ is generated directly by perturbing samples drawn from a given data distribution. Then, the adjacent state $X_{t_{n+1}}$ is produced by performing one Euler integration step from $X_{t_n}$ using the given pretrained sampler.
• CDDS (our scenario): Both states are generated from simulating the pre-trained sampler from $t=0$ to $n+1$. We do not first generate a dataset by simulating the sampler up to $t=T$, and then apply the generative model scenario.

Moreover, our Self-Consistent Diffusion Sampler (SCDS) further distinguishes itself from traditional consistency distillation by entirely removing the need for a pretrained sampler. Traditional consistency models rely on deterministic trajectories from data samples, providing explicit guidance towards high-density regions. In contrast, SCDS must also explore the space, which we achieve via a Brownian motion. Self-consistency is then enforced on deterministic sub-trajectories. This contrasts with the purely deterministic trajectories used in standard consistency models.

We have also clarified the distinctions in the form of a table in our response to Reviewer emRE, Question 2.

Question 2

Could you elaborate on whether other diffusion acceleration techniques, such as shortcuts or consistency trajectory models, are applicable to the generation of unnormalized distributions? If so, should these state-of-the-art diffusion acceleration methods be included in the comparison?

These techniques assume access to a dataset sampled from the target distribution. They lack an intrinsic mechanism to incorporate information from an unnormalized density. Our methods CDDS and SCDS precisely adapts the ideas of consistency distillation and shortcuts to the sampling problem, where only an unnormalized density is known. This difference in accessible information (also detailed in answers to Q1-3 of reviewer emRE) justifies the choice of baselines in our experiments. Please refer to [1] for an evaluation framework of diffusion samplers.

Question 3

The self-consistency loss in SCDS is designed to ensure that a large step yields the same result as multiple smaller steps. How sensitive is this self-consistency loss to the choice of step sizes, and are there optimal step sizes that lead to faster convergence or better sample quality?

We specifically chose step sizes that are powers of two because this minimizes the computational overhead of the self-consistency loss. The power-of-two schedule requires only 3 additional network function evaluations per training iteration. Choosing other arbitrary step sizes would increase computational cost, as intermediate points would need additional simulation.

However, beyond computational convenience, the core rationale behind the self-consistency loss is not overly sensitive to the exact choice of step size as long as:

• Larger steps are exactly composed of multiple smaller base steps, allowing consistent alignment of trajectories at different scales.
• Resolution: The smallest step size ($T/N$) is sufficiently small to ensure accurate approximation of the underlying probability flow ODE.

Conceptually, any choice of step sizes satisfying these criteria would ensure effective learning. Thus, while the power-of-two step size is computationally optimal and practically convenient, our method is robust and flexible enough to accommodate other step sizes without substantial differences in convergence speed or sample quality. A systematic exploration of other step-size schedules could be an interesting ablation for future work.

Question 4

Can the proposed SCDS be applied to existing image generation tasks, and what advantages does it have over existing methods?

SCDS is specifically designed for scenarios where the target distribution is provided through an unnormalized density. Instead, traditional image generation tasks, such as CIFAR-10 or MNIST, provide samples directly from the target distribution. We are not aware of large-scale image tasks formulated explicitly as unnormalized sampling problems [1].

[1] D. Blessing, X. Jia, J. Esslinger, F. Vargas, and G. Neumann. Beyond ELBOs: A Large-Scale. Evaluation of Variational Methods for Sampling. In International Conference on Machine Learning, 2024.

Dear Reviewer,

This is a gentle reminder that the Author–Reviewer Discussion phase ends within just three days (by August 6). Please take a moment to read the authors’ rebuttal thoroughly and engage in the discussion. Ideally, every reviewer will respond so the authors know their rebuttal has been seen and considered.

Thank you for your prompt participation!

Best regards,

AC

Dear Reviewer,

A gentle reminder that the extended Author–Reviewer Discussion phase ends on August 8 (AoE).

Please read the authors’ rebuttal and participate in the discussion ASAP. Regardless of whether your concerns have been addressed, kindly communicate:

• If your concerns have been resolved, please acknowledge this clearly.

• If not, please communicate what remains unaddressed.

The “Mandatory Acknowledgement” should only be submitted after:

• Reading the rebuttal,

• Engaging in author discussion,

• Completing the "Final Justification" (and updating your rating).

As per policy, I may flag any missing or unresponsive reviews and deactivate them once additional reviewer feedback has been posted.

Thank you for your timely and thoughtful contributions.

Best regards,

AC

Dear Reviewer,

Please notice that the “Mandatory Acknowledgement” should only be submitted after:

• Reading the rebuttal,

• Engaging in author discussion,

• Completing the "Final Justification" (and updating your rating).

Kindly review the authors’ response and reply as soon as possible. Regardless of whether your concerns have been fully addressed:

• If your concerns are resolved, kindly confirm this clearly.

• If not, please explain what issues remain.

Thank you for your cooperation.

Best,

AC

I acknowledge the authors' rebuttal. They have addressed some of my concerns, and while the paper presents incremental extensions to the existing content, I will keep my original score.

Question 1

Could the authors clarify which specific discrepancy measure is used in Equation (7)?

The discrepancy measure used in Equation (7) is a standard mean squared error. We have corrected (7) accordingly.

Question 2

How do the proposed methods compare to the baselines in terms of computational capacity, such as the number of parameters? This is important for fair comparisons

In our experiments, all baselines including PIS, DIS, DDS, and our proposed CDDS are of all the same size. The network is composed of a MLP with four layers and 64 units per layer along with a Fourier-based time embedding of the same dimensions. SCDS has an additional small 2 layers time-embedding to condition on the step-size. Therefore, all methods have a similar computational capacity and the comparisons are fair. We’ve included these details and discussions about the number of parameters in the revised manuscript.

Question 3

How strong/realistic is the assumption of a Lipschitz constraint in Theorem 1, especially in the SDE setting where involved quantities may not be differentiable? Additionally, how does the method account for or bound the error when the corresponding loss function exceeds zero?

We thank the reviewer for this insightful question. The Lipschitz continuity assumption in Theorem 1 applies to the learned consistency function defined as a deterministic mapping from states along the PF ODE. The Lipschitz continuity assumption in Theorem 1 is a standard condition used widely in the analysis of neural ODE (e.g., [1]). Typically, neural networks with standard activations and bounded parameters satisfy this constraint.

Theorem 1 proves the network can be arbitrarily accurate if the loss reaches zero, which has provided sufficient support of the theoretical correctness of our method. When the corresponding loss function exceeds zero gradually, we intuitively believe the sample quality will drop. Nonetheless, it could still perform well with some non-zero losses, as empirically verified by our experiments. For instance, single-step CDDS beats PIS and DDS with 128 steps in terms of Sinkhorn distance for the GMM experiment. CDDS also recovers all modes in one step as shown in Figure 4. Theoretically, deriving a precise and quantitative error bound as a function of the non-zero loss value is challenging, as it may be influenced by a complex interplay of factors like the non-convex nature of the loss landscape, the specific local minimum found by the optimization algorithm, and the generalization properties of the network, which cannot be captured by the single loss value alone. As the primary contribution of our paper is the proposal of single-step samplers, we believe this detailed analysis on error bound of non-zero loss scenarios is beyond the scope of this paper.

[1] Y. Song, P. Dhariwal, M. Chen, and I. Sutskever. Consistency Models. In International Conference on Machine Learning, 2023.

Dear Reviewer,

This is a gentle reminder that the Author–Reviewer Discussion phase ends within just three days (by August 6). Please take a moment to read the authors’ rebuttal thoroughly and engage in the discussion. Ideally, every reviewer will respond so the authors know their rebuttal has been seen and considered.

Thank you for your prompt participation!

Best regards,

AC

Dear Reviewer,

A gentle reminder that the extended Author–Reviewer Discussion phase ends on August 8 (AoE).

Please read the authors’ rebuttal and participate in the discussion ASAP. Regardless of whether your concerns have been addressed, kindly communicate:

• If your concerns have been resolved, please acknowledge this clearly.

• If not, please communicate what remains unaddressed.

The “Mandatory Acknowledgement” should only be submitted after:

• Reading the rebuttal,

• Engaging in author discussion,

• Completing the "Final Justification" (and updating your rating).

As per policy, I may flag any missing or unresponsive reviews and deactivate them once additional reviewer feedback has been posted.

Thank you for your timely and thoughtful contributions.

Best regards,

AC

Question 1:

In Eq. 2,3, what are $\pi$ and $\tau$? Are they supposed to be the diffusion prior $p_{\text{prior}}$ and target $p_{\text{target}}$?

Yes, $\pi$ and $\tau$ represent the initial probability distributions for the forward-time process and reverse-time process, respectively. As clarified explicitly in the manuscript (lines 93-94), by selecting $\pi = p_{\text{prior}}$ and $\tau = p_{\text{target}}$ in our setting, we generate samples from the target distribution $p_{\text{target}}$ using Eq. (2) with the optimal control $u^*$. We have clarified this in the revised version.

Question 2:

Which of Eq. 2 or 3 is the reverse (generative) process, and which is the forward process?

In our notation, Eq. (2) indeed corresponds to the generative (forward-time) process. Note that in the diffusion sampling literature [1], the generative process (from prior to target) is conventionally called the forward process because the noising process cannot be directly simulated from data. The two processes are time-reversals of each other.

Question 3:

In Eq. 6, should $P^{u, p_{\text{prior}}}$ and $P^{v, p_{\text{target}}}$ appear explicitly instead of $P^{u,\pi}$ and $\overline{P}^{v,\tau}$?

Thank you for pointing this out. We have updated Eq. (6) to remain general, explicitly using the general symbols $u, v, \pi, \tau$ for clarity and consistency, and we removed the parameter superscript to avoid confusion:

$\log\frac{\mathrm d {P}^{u,\pi} }{\mathrm d \overline{P}^{v,\tau}} = \int_0^T (u + v) \cdot \Big(u + \frac{v - u}{2} + \nabla \cdot (\sigma v - \mu)\Big)\mathrm d s + \int_0^T (u + v) \mathrm d W_s + \log \frac{\pi (X_0)}{\tau(X_T)}$

Question 4:

Is the consistency function $f(X_t, t)$ uniquely defined in Theorem 1, and what does it mean for it to be a ‘consistency function of a probability flow ODE’?

Yes, for any given initial condition (sample from the prior), the Probability Flow ODE is deterministic, thus uniquely defining the consistency function as the deterministic mapping from $X_t$ to $X_T$. The consistency function of the PF ODE, therefore, refers to the unique deterministic mapping induced by integrating the PF ODE forward in time.

Question 5:

Can the partial ODE integration used in CCDS be applied to standard consistency model training?

In principle, yes, although we have not empirically evaluated it. However, we anticipate greater integration errors compared to standard consistency training, which typically benefits from exact perturbation kernels that provide effectively ground-truth intermediate noisy data (assuming the drift is linear in X).

Question 6:

Should there be a minus sign between functions in Eq. 7?

Yes, thank you for pointing out this typo. We have corrected Eq. (7) accordingly

Question 7:

Where is the unnormalized density oracle used in the distillation process (CDDS)?

The unnormalized density is indirectly utilized during training of the teacher sampler, which guides the intermediate trajectories used for distillation. As you correctly suggest, an interesting future improvement could be to incorporate a sampling loss directly into the distillation step (CDDS) to potentially correct deficiencies of the pretrained sampler during distillation.

Question 8:

How is the sampling loss (Eq. 5) computed in practice?

The sampling loss (KL or LV divergence) is computed explicitly using discretization of the Radon–Nikodym derivative. Algorithm 5 in the appendix provides detailed steps for this computation. We will explicitly mention that the detailed steps are provided in Algorithm 5 in the revision.

Question 9:

Does estimating log normalization constant or ESS require multiple diffusion steps, thus increasing NFE beyond 1?

When performing single-step sampling (CDDS/SCDS), the NFE is still one. Specifically, the interior integral term collapses leaving a single evaluation of the likelihood ratio at the boundary points only.

Question 10:

The discretized Radon–Nikodym derivative should be explained in more detail.

In practice, we approximate the integrals in the Radon–Nikodym derivative by replacing them with finite sums over a discrete time mesh. The implementation of the discretized Radon–Nikodym derivative is provided in Appendix Section A.1, under the paragraph “Discretizing the Radon–Nikodym derivative.”

Question 11:

In Appendix line 813, why is $v=0$ assumed?

This assumption is made explicitly to ensure uniqueness of the solution to Eq. (4). Without fixing the noising process (the backward process defined by $v$), multiple solutions may exist [1]. We have explicitly stated our choice of a Variance-Preserving SDE as the noising process in the revised manuscript.

Question 12:

In lines 248–249, is $N$ the number of sampling steps, and what is the running-cost term?

Yes, $N$ is the number of sampling steps; we have clarified this explicitly in the revised manuscript. The running-cost term corresponds to the integral over time in Eq. (6), and, when discretized, corresponds to the sum computed in line 12 of Algorithm 5. We have clarified this terminology when first introducing the RND (Eq. 6).

Question 13:

Have the authors considered computing posterior predictive log likelihood on a held-out test dataset for Table 2?

For Table 2, we compute only the ELBO since we do not have direct access to real samples from these targets. Posterior predictive log likelihood would require samples or predictions on held-out data, which is not available in these tasks as formulated.

Question 14:

Could the authors provide a clearer summary of the empirical performance of CCDS and SCDS, especially considering statistical significance?

CDDS consistently provides strong single-step sampling performance. SCDS, while slightly underperforming in some high-dimensional scenarios, still achieves strong performance, demonstrating its capability to learn effective single-step sampling without pretrained samplers. We have included standard deviations explicitly in Appendix Tables 4 and 5 to strengthen the empirical claims further.

[1] Lorenz Richter and Julius Berner. Improved sampling via learned diffusions. In International Conference on Learning Representations, 2024.

Dear Reviewer,

This is a gentle reminder that the Author–Reviewer Discussion phase ends within just three days (by August 6). Please take a moment to read the authors’ rebuttal thoroughly and engage in the discussion. Ideally, every reviewer will respond so the authors know their rebuttal has been seen and considered.

Thank you for your prompt participation!

Best regards,

AC

Question 1

Why should the self-consistency objective work? Why does this trick on ‘synthetic’ samples from the control function work? If there would be a theoretical/intuitive/toy problem justification for this, the paper would be a lot stronger, as SCDS is the main contribution of the paper.

The self-consistency objective intuitively works because it explicitly enforces that taking a single large step from an intermediate noisy state should yield the same outcome as multiple smaller steps. Specifically, as detailed in the second paragraph of Section 5.2, the model is trained such that predictions of larger steps (step size 2d) match the results from sequences of smaller incremental steps (two steps of size d).

Initially, the network learns a reliable “base” transition for the smallest step size $d = T/N$, guided by the stochastic diffusion process from the unnormalized density $\rho$. This base transition helps discover high-density regions (Section 5.2, first paragraph). For larger step sizes, the model’s predicted state is matched against states obtained by recursively applying previously learned smaller-step transitions. Hence, the model uses its own “synthetic” states (with gradients stopped) as targets to enforce consistency at progressively larger scales.

To better highlight this intuition, we have refined the Overview Section 5.1:

Even though CDDS demonstrates that single-step sampling is feasible, it requires a pretrained diffusion sampler for distillation. In this section, we further introduce self-consistent diffusion sampler (SCDS) that achieves single-step sampling without requiring a pre-trained diffusion sampler. The intuition behind SCDS is to reconcile the conflicting objectives of diffusion-based sampling and consistency functions. Diffusion-based samplers typically learn a time-dependent control function designed to steer an SDE from the prior toward the target through multiple small steps. In contrast, consistency models aim to directly map an intermediate noisy state from an ODE trajectory to the final target state.

We propose a unified network architecture that conditions the control function $u_\theta(X_t, t, d)$ simultaneously on the current time $t$ and the desired step size $d$. By varying the step size parameter $d$, the model smoothly transitions between small, incremental steps (as in diffusion-based samplers) and large leaps (as in consistency models). For smaller steps $d = T/N$, the model learns stochastic incremental diffusion-based sampling guided by the density $\rho$, allowing exploration of the target distribution. For larger step sizes, the model is trained such that its outcomes match those obtained from sequences of smaller steps, thereby enforcing consistency across different scales. This step-conditioned control architecture amortizes the learning of both small transitions and large jumps within a single network. During inference, setting $d = T - t$ enables single-step sampling directly from any intermediate noisy state, whereas setting $d = T/N$ recovers traditional multi-step diffusion sampling.

Question 2:

Why would the SCDS method not degrade the sample quality? Perhaps there are some theoretical results you could draw on regarding samplers not perfectly recovering the reverse trajectory (prior to target) but still having good performance.

Previous work [1] has formally derived the Radon–Nikodym derivative (RND) of Equation (6). If the RND is driven to 1, the forward and backward path measures become identical. In practice, perfect recovery of these reverse trajectories is not necessary to obtain high-quality samples. Like ours, multiple previous works such as [2] have obtained good results minimizing discretized versions of (6).

Intuitively, the self-consistency loss doesn't conflict with Radon–Nikodym derivative because our network models a step-size-conditioned control (scaled score) instead of a consistency function. By conditioning the on the step size, our architecture accommodates both short-step diffusion updates (required for accurate RND estimation) and larger shortcut steps across the Euler integration. Consequently, this unified modeling approach does not compromise but rather supports accurate estimation of the RND, and facilitates recursive learning of single-step sampling.

Question 3:

Can CDDS or SCDS be used for image generation tasks? If not, clarify limitations. Showing results on CIFAR-10 or MNIST would broaden impact.

We appreciate this insightful suggestion. However, CDDS and SCDS as presented are specifically designed for sampling from unnormalized target densities, where no direct samples are available. Commonly used image datasets such as CIFAR-10 or MNIST inherently provide samples rather than unnormalized densities. We are not aware of standardized unnormalized-density large image generation tasks. We refer to [2] for an evaluation framework and standard sampling benchmark.

Question 4:

Negative societal impacts were not addressed. How could these samplers negatively affect society?”

While we do not anticipate major negative societal impacts from our sampling methods themselves, though we acknowledge potential misuse of efficient sampling algorithms in sensitive applications, such as synthetic data generation or privacy. We’ve updated the Limitations section in the Appendix.

[1] J. Berner, L. Richter, and K. Ullrich. An Optimal Control Perspective on Diffusion-based Generative Modeling. Transactions on Machine Learning Research, 2024.

[2] D. Blessing, X. Jia, J. Esslinger, F. Vargas, and G. Neumann. Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. In International Conference on Machine Learning, 2024.

Dear Reviewer,

This is a gentle reminder that the Author–Reviewer Discussion phase ends within just three days (by August 6). Please take a moment to read the authors’ rebuttal thoroughly and engage in the discussion. Ideally, every reviewer will respond so the authors know their rebuttal has been seen and considered.

Thank you for your prompt participation!

Best regards,

AC

The reviewer thanks the author for their response. The intuition behind the self-consistency objective is appreciated.

Furthermore, while the reviewer does not have the background to adjudicate on the significance of the experiments, the responses to other reviewers has been helpful.

This reviewer will bump the score up by one. However, for a stronger score, some kind of theoretical result, perhaps on a simplified problem, is necessary with respect to SCDS.

Question 1

Can authors explain in which way the proposed method is new? and not “just” distill a sampler, as in the original consistency sampler?

Our approach indeed builds upon the principles of consistency models introduced by [1], but differs in terms of the information available. Original consistency models were developed for generative modeling, assuming access to a dataset sampled from a target distribution. This enables easy generation of intermediate states via known perturbation kernels from these real samples.

In our problem formulation, we have no direct access to samples from the target distribution, only an unnormalized density oracle. Consequently, we cannot generate noisy intermediate states from target samples. To overcome this, we introduce two distinct methods tailored to the sampling setting:

• Consistency-Distilled Diffusion Sampler (CDDS): This method distills a pretrained sampler for unnormalized distributions. Critically, we leverage noisy states obtained directly from intermediate partial integrations of the sampler from the prior. It entirely removes the need for data from the target or datasets of fully diffused samples.
• Self-Consistent Diffusion Sampler (SCDS): Our second approach is even more novel. Since there is no data to reveal high-density regions, SCDS simultaneously explores high-density regions via Brownian motion and learns to skip intermediate Euler steps. To our knowledge, SCDS is the first attempt to combine consistency with stochastic control-based diffusion samplers to achieve single-step sampling without external pretrained models or data.

Our methods extend the scope of consistency models to sampling scenarios that were previously unexplored.

Question 2

Could you provide a clear step-by-step comparison of consistency models for generative models and for unnormalized distribution sampling?

Yes, we provide a comparison table which clearly compares the differences between existing consistency distillation and consistency training methods, with our proposed CDDS and SCDS:

We’ve added this concise comparison table to the revised manuscript.

Question 3

Regarding the sentence "While our distillation approach builds upon the core principles…" it seems the proposed method only requires a sampler, and once you have the sampler, the distillation is the same. Is this correct? If not, please clarify.

We appreciate the opportunity to clarify this point further. Once a sampler is available, our consistency distillation (CDDS) approach performs consistency training similar to the original method. However, the difference is how we generate intermediate noisy states for training:

• Generative Modeling Scenario: Intermediate state $X_{t_n}$ is generated directly by perturbing samples drawn from a given data distribution. Then, the adjacent state $X_{t_{n+1}}$ is produced by performing one Euler integration step from $X_{t_n}$ using the given pretrained sampler.
• CDDS (our scenario): Both states are generated from simulating the pre-trained sampler from $t=0$ to $n+1$. We do not first generate a dataset by simulating the sampler up to $t=T$, and then apply the generative model scenario.

Moreover, our Self-Consistent Diffusion Sampler (SCDS) further distinguishes itself from traditional consistency distillation by entirely removing the need for a pretrained sampler. Traditional consistency models rely on deterministic trajectories from data samples, providing explicit guidance towards high-density regions. In contrast, SCDS must also explore the space, which we achieve via a Brownian motion. Self-consistency is then enforced on deterministic sub-trajectories. This contrasts with the purely deterministic trajectories used in standard consistency models.

Question 4

Can you comment on the significance/challenge of the experimental benchmarks used?

We use a range of well-established and challenging benchmarks standard in the Bayesian inference and sampling literature, such as the Gaussian mixture model, funnel distribution, many-well, log-Gaussian Cox process (1600 dimensions), and Bayesian logistic regression on the ionosphere dataset representing realistic inference problems of considerable complexity and scale. Our results show that our methods significantly accelerate sampling via learned diffusion (1 compared to 128 sampling steps). We refer to [2] for an evaluation framework of diffusion-based samplers.

[1] Y. Song, P. Dhariwal, M. Chen, and I. Sutskever. Consistency Models. In International Conference on Machine Learning, 2023.

[2] D. Blessing, X. Jia, J. Esslinger, F. Vargas, and G. Neumann. Beyond ELBOs: A Large-Scale Evaluation of Variational Methods for Sampling. In International Conference on Machine Learning, 2024.

Dear Reviewer,

This is a gentle reminder that the Author–Reviewer Discussion phase ends within just three days (by August 6). Please take a moment to read the authors’ rebuttal thoroughly and engage in the discussion. Ideally, every reviewer will respond so the authors know their rebuttal has been seen and considered.

Thank you for your prompt participation!

Best regards,

AC