Wasserstein Convergence of Critically Damped Langevin Diffusions

We thank the reviewer for the positive assessment and for recognizing the novelty of our $\mathcal{W}_2$ analysis for CLD. Below, we address all concerns and describe the additions in the revision.

Explicit pseudocode

For clarity, we now include a concise pseudocode block for both training (Algorithm 1) and sampling (Algorithm 2). They are now integrated into the manuscript. We are happy to provide a more detailed version if the reviewers find it necessary.

Algorithm: CLD Training

Inputs: Dataset $\mathcal{D}$, batch size $B$, network $s_\theta(\cdot, t)$, a positive weight function $\lambda : [0,T] \to \mathbb{R}_+$, and noise level $\epsilon \geq 0$.

Precomputation: Compute $\tilde \Sigma_{0,t} = \Sigma_{0,t} + e^{t A} \, \mathrm{diag}(0 \cdot I_d, v^2 I\_d) \, (e^{t A})^\top$ (Note: The value of $\Sigma_{0,t}$ depends on $\epsilon$, see Lemma A.2 (Eq. 23).)

While not converged:

1. Sample {$x^{(i)}$}$_{i=1}^{B} \sim \mathcal{D}$
2. Sample {$t^{(i)}$}$_{i=1}^{B} \sim \mathcal{U}([0,T])$
3. Sample {$\varepsilon^{(i)}$}$\_{i=1}^{B} \sim \mathcal{N}(0, I_{2d})$
4. Compute $$$$

\overrightarrow{\mathbf{U}}_{t^{(i)}} = e^{t^{(i)} A} ( \overrightarrow{X}_0 \ 0_d )^\top + ( \tilde{\Sigma}_{0,t^{(i)}} )^{1/2} \varepsilon^{(i)}

\mathcal{L} \gets \frac{1}{B} \sum_{i=1}^{B} \lambda(t^{(i)}) || s_\theta(t^{(i)}, \overrightarrow{\mathbf{U}}_{t^{(i)}}) + (\tilde \Sigma_{0,t^{(i)}})^{-1/2} \varepsilon^{(i)}||^2

We thank the reviewer for their valuable feedback and have addressed each point below. The proposed revisions have been incorporated into the revised paper.

On the benefit of deriving $\mathcal{W}\_2$-bound.

Previous empirical work (Dockhorn et al., 2022) shows that CLD can outperform standard SGMs forwards on image benchmarks yet the only theoretical results were given in KL (Chen et al. (2023) and Conforti et al.(2025)) and not in $\mathcal{W}\_{2}$ -distance even though the latter enjoys nice geometric properties and is a true mathematical distance. The goal of this work is to close this gap and place CLD on the same footing as VP/VE SDEs in Wasserstein distance, while maintaining state‑of‑the‑art, minimal, regularity assumptions on the data distribution. In this way, as training in the coupled phase space has been shown by Dockhorn et al. (2022) to be empirically more stable (in particular the Hybrid Score Matching loss), our results certify that the associated Euler-Maruyama sampler remains $\mathcal{W}\_{2}$‑stable even when the score is learned only approximately (Assumption H3).

As also mentioned by the reviewer, our contribution concerns a different metric and setting as Chen et al. (2023) and Conforti et al.(2025). We give here the first $\mathcal{W}\_2$ guarantees for CLD, showing it attains the same convergence rates as SGMs. Note that the proof technique is completely different and rely on coupling techniques tailored to $\mathcal{W}\_2$ analysis. We also agree that KL divergence upper bounds can sometimes be translated to $\mathcal{W}\_2$ upper bounds; however, this is non-trivial and the price to pay may be high. In the case of compactly supported data distribution, the KL bounds constitute a valid upper bound to the $\mathcal{W}\_2$ metric. Nevertheless, this upper bound depends on the diameter of the space (Villani, 2008, Thm. 6.12, p. 82), which could therefore grow very large, especially in high dimensional settings, such as those encountered in image datasets. Regarding the log-concave case, we are not sure what was meant. If the reference was meant to Talagrand's inequality, this typically requires to verify regularity on the measure to the right of the KL (i.e. $\nu$ for $\mathrm{KL}(\mu||\nu)$) which in standard convergence results (Chen et al(2023), Conforti et al.(2025)) is usually the distribution generated by the algorithm and not the data distribution.

Note that Gao et al. (2025) assume strong log-concavity of the data distribution (among other conditions) in $\mathbb{R}^d$, and explicitly state that their results are not implied by existing KL-based convergence bounds without further assumptions: "our results are not implied by the existing convergence results for SDE-based samplers." (Gao et al. (2025), p.5).

On the practical side, as noted in Strasman et al. (2025), Wasserstein upper bounds are often more tractable to approximate empirically, whereas KL divergences and Fisher information are difficult to estimate in high dimensions. It is also worth noting that previous results focus on specific parametrization choices ($a=1$ and $\sigma = 2$ so that the drift and diffusion coefficients are coupled in a specific manner), which limits their generality and expressiveness.

Response to Q1.

Note that our main contribution is Theorem 3.1, which establishes a quantitative bound for generic Langevin dynamics under weak Lipschtiz-like assumptions. Moreover, we propose in Section 4 an upper bound in the specific and widespread strongly log-concave setting. This framework allows to target more precisely the role of the parameters $a$, $\epsilon$ and $\sigma$. We agree that the initial version was a bit confusing. We now highlight that Section 4 is a particular case of Theorem 3.1. In this way, we turn the title of Section 4 into: "Specific analysis of $\epsilon$-CLD under strong log-concavity".

On the log-concavity at all times.

Regarding the question on Assumption H4, we agree that this is stronger than necessary. In fact, there is no need to require strong log‑concavity at all times. Similar to Gao et al. (2025) and Strasman (2025) (in the VPSDE case), it is enough to assume that the data probability density function is $c_0$-strongly log concave for the property to propagate at all times, i.e. for $\log p_t$ to be strongly concave for every $t \geq 0$. In fact, $c_t$ is explicitly available and given by,

$

c_t = \left( \frac{1}{ \min (c_0, v) \sigma^2_{\min}(e^{-tA})} + \lambda_{\max}(\Sigma_{0,t}) \right)^{-1} .

$

In particular, this expression also highlights the impact of the calibration of the initial variance $v$ of the velocity component, with respect to the data log-concavity constant and the implicit impact of $\epsilon$ through $\Sigma_{0,t}$. In addition, note that Section 4 provides quantitative upper bounds in the widespread setting of strongly log-concave distributions but Theorem 3.1 does not require this assumption.

Clarification on the assumptions.

To take into account this remark on the discussion of the assumptions and to better position our work in the literature, we have added a paragraph discussing our assumptions. Please refer to our answer to reviewer bmjs where we discuss in detail the assumptions of our work and relation to other results in the literature.

Regarding the benefit of CLDs and their $\epsilon$-variant.

The kinetic framework offers practical advantages that standard VP/VE SDEs do not offer:

-Forward tuning: the position–noise level $\varepsilon$, the velocity–noise level $\sigma$, and the initial velocity scale $v$ can be adjusted independently of the drift parameter $a$, whereas in VP/VE the drift and diffusion are rigidly coupled. This gives extra freedom to practitioners.

-Training objective: the extended phase space admits Hybrid Score Matching (HSM) in addition to classical DSM training, and Dockhorn et al. showed empirically that HSM lowers gradient variance during the training.

-Backward solver: because the reverse flow is Hamiltonian‑like, one can employ structure‑preserving integrators such as symplectic or leapfrog schemes; on our toy problems leapfrog consistently outperforms Euler–Maruyama, even in the $\varepsilon$‑regularized variant. We leave for future work the theoretical analysis of such schemes.

-$\epsilon$-regularized version: the introduction of the $\epsilon$ parameter was initially motivated theoretically by Lemma B.1, where we saw that larger $\epsilon$ loosens the stepsize condition on $h$—i.e., strengthens the effective contraction, although a full global trade‑off analysis remains open.We however, illustrated this effect empirically on toy datasets and saw that a small regularization yield improvement in the $\mathcal{W}_2$ metric on a Funnel dataset in dimension 50. We further illustrate this claim by extending our empirical analysis to dimension 100. Moreover, we have added results on two additional challenging synthetic datasets: MG-25 in dimension 100 (a 25-mode Gaussian mixture) and Diamonds (a 2D Gaussian mixture arranged in a diamond-shaped configuration). All these additional experiments (Tables 1, 2 and 3 in our response to reviewer N5Lx) further support our claims and exhibit trends consistent with our earlier findings.

We hope these comments meet the expectations and would be happy to clarify any remaining questions.

References:

• Villani, C. (2008). Optimal Transport: Old and Young. Springer.

We thank the reviewer for the positive and thoughtful review, as well as the opportunity to clarify our contributions. We address each point below and have already incorporated the corresponding changes into the revised manuscript. All references are cited consistently with those in the original version.

[W1] On the Theoretical Contribution and Practical Relevance of CLD.

Previous empirical work (Dockhorn et al., 2022) shows that CLD can outperform standard SGMs forwards on image benchmarks, yet no $\mathcal{W}\_{2}$ analysis existed; existing theory for CLD is restricted to KL (Chen et al., 2023; Conforti et al., 2025) and to a specific parametrization of the Langevin dynamics. Our goal is to close this gap (Theorems 3.1 and 4.1) and place CLD on the same footing as VP/VE SDEs in Wasserstein distance, while maintaining the state‑of‑the‑art minimal regularity assumptions on the data distribution.

By providing the first $\mathcal{W}\_{2}$ guarantees for CLD, we show that the method is not only empirically competitive but also theoretically sound and tunable, offering an alternative to classical SGMs. In this way, as training in the coupled phase space has been shown by Dockhorn et al. (2022) to be empirically more stable (in particular the Hybrid Score Matching loss), our results certify that the associated sampler remains $\mathcal{W}\_{2}$‑stable even when the score is learned only approximately (Assumption H3).

Moreover, as the reverse SDE is Hamiltonian‑like, it is possible to design structure‑preserving integrators such as symplectic or leapfrog schemes, which are inaccessible to standard SGMs. On our toy example leapfrog integrator consistently outperforms Euler–Maruyama in generation quality, even in the $\varepsilon$‑regularized setting. This perspective paves the way for improved discretization error rates. While a full analysis of these structure‑preserving schemes is beyond the scope of this work, our results lay the groundwork for future improvements in discretization error rates.

[W2] Convergence rate for $\mathcal{W}\_2 ( \pi\_{\mathrm{data}}, \mathcal{L}(\bar X\_T^{\theta}) )$.

Since both training and sampling operate in the extended phase space, we established a convergence result for the joint distribution, even though the main quantity of interest is the marginal data distribution. The error on the data distribution is directly controlled by the joint bound via contraction under projection of the $\mathcal{W}_2$ distance, $\mathcal{W}\_2 ( \pi\_{\mathrm{data}}, \mathcal{L}(\bar X\_T^{\theta}) ) \leq \mathcal{W}\_2 \left( \pi\_{\mathrm data}\otimes \pi\_v ,\mathcal{L}( \mathbf{\bar U}^\theta_T) \right) .$ We have added this clarification to the revised version of the paper.

[W3] Extended experiments and impact of $\varepsilon$.

We understand the concern regarding the use of real, high-dimensional datasets. However, training such models is very computationally intensive, and as our work is primarily focused on theoretical guarantees, we restricted our empirical validation to toy datasets that serve to illustrate our mathematical claims. In this way, the introduction of the $\varepsilon$ parameter was initially motivated theoretically by Lemma C.1, where we showed that larger $\varepsilon > 0$ loosens the stepsize condition on $h$—i.e., strengthens the effective contraction. However, we acknowledge that a complete global analysis of this trade-off remains an open question.

To address this concerns within the limited rebuttal window, we extended our numerical experiments to a higher-dimensional toy setting. Specifically, we provide an analysis of the sensitivity of the $\varepsilon$ parameter on the 100-dimensional Funnel dataset across different choices of drift $a$. We report the mean and standard deviation over 5 runs in Table 1. In this extended empirical analysis it appears that the benefit of a small $\varepsilon$ regularization is confirmed even for different choices of $a$. Moreover, note that in our ablation study, we set $\sigma$ to $\sqrt{2}$ so that the case where $a=2$ and $\epsilon=0$ exactly matches the setting recommended by Dockhorn et al. (2022). We also adopted the same time renormalization as them to make the comparison as fair as possible (see Appendix C.1).

On the one hand, we observe that this parametrization is highly effective on toy datasets. On the other hand, several $\varepsilon$-regularized variants perform even better. To further support this observation, we are extending our evaluation to additional challenging toy datasets: MG-25 (a 25-mode 100D Gaussian mixture) and Diamonds (a 2D Gaussian mixture with diamond geometry). Preliminary results suggest that the same conclusions hold (see Tables below). If of interest, we are happy to include standard deviations in the revised paper and to reproduce plots analogous to Figure 1 for these new datasets.

Table 1: Comparison of mean Wasserstein distance for different noise levels $\varepsilon$ on the Funnel‑100D (Mean ± standard deviation across 5 runs; lower is better)

Table 2: Comparison of mean Wasserstein distance for different noise levels $\varepsilon$ on the MG25‑100D

Table 3: Comparison of mean Wasserstein distance for different noise levels $\varepsilon$ on the Diamond‑2D

We thank the reviewer once again for their insightful comments. The revisions directly address each point raised and, we believe, further strengthen the manuscript. We hope the changes align with the reviewer’s expectations and remain available to clarify any remaining questions.

We thank the reviewer for their valuable feedback and have addressed each point below. The proposed revisions have been incorporated into the revised paper. Based on these comments mentioning that the paper is difficult to follow and lacks of intuitive remarks and comments or examples on the assumptions, we simplified and clarified the presentation of the main results in the revised version. Without loss of generality and without modifying our results we propose the two following revisions.

• First, we provide a unified result that holds for all $\epsilon \geq 0$, rather than treating the cases $\epsilon=0$ and $\epsilon>0$ separately. This reformulation aims to make the theoretical contributions more accessible to a broader audience and clarifies the discussions (see below for our comments on the assumptions).

• To clarify our assumptions and better compare our work with the existing literature (see, e.g., [1]), we have added a paragraph discussing our assumptions, in particular the one-sided Lipschitz assumption. Such assumptions on the score function provide a bound on the largest eigenvalue of its Jacobian and are the weakest assumptions in the score-based generative models literature, see [1,2]. Our work aligns with the most recent results, and our assumptions cover in particular data distributions which are not covered by the strong logconcave assumptions used in the literature. In order to simplify the paper and to compare our framework with the existing literature we propose to replace H2 and H4 in the revised paper by the following assumption.

New Assumption H2'. The data distribution is of the form $\pi_{\rm data} (x) \propto \exp{(-u(x) + a(x))}$ and satisfies:

• There exist $\alpha>0$ and $L_0>0$ such that $\alpha I_{d} \preceq \nabla^{2}u(x) \preceq L_0 I_{d}$ for all $x\in\mathbb{R}^{d}$.
• There exists $K>0$ such that $|a(x)-a(y)|\le K||x-y||$ for all $x,y\in\mathbb{R}^{d}$ with $||x-y||\le1$.

This minor adaptation of our Lipschitz-type assumptions is consistent with the weakest assumptions found in the literature on classical score-based generative models (i.e. without the Langevin forward dynamics) and has no impact on our main results. This assumption models the data distribution as a strongly log-concave component $u(x)$ perturbed by a term $a(x)$. Intuitively, this assumption allows the distribution to deviate from strict log-concavity via the perturbation $a(x)$, while still maintaining sufficient regularity for analysis. When $a(x)=0$, the distribution reduces to a strongly log-concave case.

Discussion on Assumptions.

The novel assumptions proposed above are satisfied by standard distributions such as Gaussian and mixtures of Gaussian distributions. They are strictly weaker than classical assumptions like strong log-concavity (H4), which holds only for non-degenerate Gaussian distributions and excludes many practically relevant settings although it is still common in the literature. Regarding Assumption H3, as discussed in Section 3, it quantifies the score approximation error and is standard in recent works that establish non-asymptotic Wasserstein convergence guarantees for diffusion-based models (e.g., [2, 3]).

In the particular case where the data distribution is Gaussian, the score function admits a closed form expression and does not require approximation via a neural network. Consequently, the score approximation error vanishes, and the constant $M$ in our bound can be as chosen as $M=0$.

Clarification on the step size $h$ and its impact.

We acknowledge that the definition of $h$ was not clearly specified earlier. The constant $h$ in Theorems 3.1 and 4.1 corresponds to the step size in the Euler--Maruyama discretization of the backward SDE. This is now clearly stated in the revised version of the paper.

In Theorem 3.1, the result holds for all $h > 0$; the only condition we implicitly use in the proof is the elementary inequality $h \leq (1 - \mathrm{e}^{-h}) \mathrm{e}^h$, which is equivalent to $h + 1 \leq \mathrm{e}^{h}$, a bound satisfied for all $h > 0$ (as mentioned a few lines earlier in the proof). For Theorem 4.1, we already stated that the result holds for $h < h_\epsilon$, where $h_\epsilon$ is defined in Appendix B. We have now added the explicit form of $h_\epsilon$ to the main text for clarity.

Regarding the dependence on $h$ in $c_2$, the discretization error introduces a term of order $\mathcal{O}(\sqrt{h})$. In fact, the discretization error takes the form $c_2 \sqrt{h} + o(h)$. Since $h$ is typically chosen to be sufficiently small in practice (see also [3, 4]), the higher-order term becomes negligible, and the $\sqrt{h}$ term dominates. In practice, discretization is performed using a uniform grid, with step size $h=T/N$, where $T$ is the total time horizon and $N$ is the number of discretization steps.

Comparison with KL divergence bounds and known results.

In the strongly log‑concave setting, our Wasserstein‑2 results cannot be recovered from existing KL‑divergence bounds. Indeed, if the data distribution $\pi_{\mathrm{data}}$ is $c$-strongly log-concave, Talagrand’s inequality yields

By contrast, for distributions with compact support one may also bound the squared Wasserstein‑2 distance by the KL divergence ([6], Thm. 6.12, p. 82), but the corresponding constant depends on the diameter of the support, which can grow prohibitively large in high dimensions (e.g., image data).

Since KL divergence is not symmetric, existing KL-based results only control $\mathrm{KL}\left(\pi_{\mathrm{data}}\otimes \pi_v, \mathcal{L} (\bar{\mathbf{U}}^\theta_T)\right)$, so our convergence bound cannot be deduced as a corollary. Indeed, [3] also assumes strong log-concavity of the data distribution (among other conditions) in $\mathbb{R}^d$ and explicitly emphasizes that “our results are not implied by the existing convergence results for SDE-based samplers” ([3], p. 5). Finally, Theorem 3.1 in our work holds under far weaker assumptions (requiring neither strong log-concavity nor compact support). Moreover, existing KL-based analyses (e.g., [4, 5]) focus on a specific case of Langevin dynamics with $a = 1$ and $\sigma = 2$, which yields an isotropic Gaussian stationary distribution—this significantly simplifies the analysis. In contrast, our analysis covers the general CLD setting and extends beyond it to include an additional noise parameter $\varepsilon$.

We hope that these additions fully address the reviewer’s suggestions and contribute to strengthening the paper. We are grateful for the valuable feedback and trust that the revised version meets the expectations raised.

References

[1] Stéphanovitch, A. (2025). Regularity of the score function in generative models. arXiv:2506.19559.

[2] Gentiloni-Silveri, M. and Ocello, A. (2025). Beyond Log-Concavity and Score Regularity: Improved Convergence Bounds for Score-Based Generative Models in W2-distance. https://arxiv.org/abs/2501.02298.

[3] Gao, X., Nguyen, H. M., and Zhu, L. (2025). Wasserstein convergence guarantees for a general class of score-based generative models. Journal of Machine Learning Research.

[4] Chen, S., Chewi, S., Li, J., Li, Y., Salim, A., and Zhang, A. (2023). Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. In International Conference362on Learning Representations.

[5] Conforti, G., Durmus, A., and Silveri, M. G. (2025). Kl convergence guarantees for score diffusion models under minimal data assumptions. SIAM Journal on Mathematics of Data Science.

[6] Villani, C. (2008). Optimal Transport: Old and Young. Springer.