Beyond Scores: Proximal Diffusion Models
Authors: Zhenghan Fang
Mateo Diaz Diaz
Sam Buchanan
Jeremias Sulam
Venue: NeurIPS 2025 poster ()
Keywords: Generative models
Diffusion models
Proximal operators
Backward discretization
Primary area: deep_learning
TL;DR: This paper introduces Proximal Diffusion Models (ProxDM), derived from backward discretization of an SDE and based on learned proximal operators, achieving provable faster sampling complexity and empirically much faster convergence.

Abstract: Diffusion models have quickly become some of the most popular and powerful generative models for high-dimensional data. The key insight that enabled their development was the realization that access to the score---the gradient of the log-density at different noise levels---allows for sampling from data distributions by solving a reverse-time stochastic differential equation (SDE) via forward discretization, and that popular denoisers allow for unbiased estimators of this score. In this paper, we demonstrate that an alternative, backward discretization of these SDEs, using proximal maps in place of the score, leads to theoretical and practical benefits. We leverage recent results in _proximal matching_ to learn proximal operators of the log-density and, with them, develop Proximal Diffusion Models (`ProxDM`). Theoretically, we prove that $\widetilde{\mathcal O}(d/\sqrt{\varepsilon})$ steps suffice for the resulting discretization to generate an $\varepsilon$-accurate distribution w.r.t. the KL divergence.
Empirically, we show that two variants of `ProxDM` achieve significantly faster convergence within just a few sampling steps compared to conventional score-matching methods.

Beyond Scores: Proximal Diffusion Models
Zhenghan Fang Mathematical Institute for Data Science Johns Hopkins University zfang230jhn edu edu
Mateo Díaz
Mathematical Institute for Data Science Johns Hopkins University mateodd@jhu edu
Sam Buchanan Toyota Technological Institute Chicago sdbuchanann@eerkkleeyy edu
Jeremias Sulam
Mathematical Institute for Data Science Johns Hopkins University jul am1 @u juu.
Abstract
Diffusion models have quickly become some of the most popular and powerful generative models for high-dimensional data. The key insight that enabled their development was the realization that access the score density at different noise levels allows for sampling from data distributions by solving a reverse-time stochastic differential equation (SDE) via forward discretiza- tion, and that popular denoisers allow for unbiased estimators of this score. In this paper, we demonstrate that an alternative, backward discretization these SDEs, using proximal maps in place of the score, leads to theoretical and practical benefits. We leverage recent results in proximal matching to learn proximal operators the log-density and, with them, develop Proximal Diffusion Models (ProxDM). Theoretically, prove that Old/ Ve) steps suffice for the resulting discretization to generate an e-accurate distribution w.r.t. the KI divergence. Empirically, we show that variants of ProxDM achieve significantly faster convergence within just a few sampling steps compared to conventional score- wore-matching methods.
gradient of the log-
Introduction
Within a remarkably brief period, diffusion models [72, 27] have overtaken generative artificial intelligence, emerging as the de facto solution across diverse domains including medical imaging [13], video [36] and audio [28]. drug discovery [15], and protein docking prediction [12]. Diffusion models combine forward diffusion process that converts the data distribution we wish to sample into noise and a reverse-time process that reconstructs data from noise. The forward process is easy to simulate by progressively adding noise to data. The reverse process, which allows for sampling data from noise, involves the so-called score function-the gradient of the log-density of the noisy data at any time point. Though unknown in general, such score can be estimated by common denoising algorithms that approximate minimum mean- rrean-squardd-erro solutions. Various diffusion algorithms have been recently proposed [75, 53, 81], all relying on forward discretizations of the reverse stochastic differential equation, and allowing for impressive generative modeling performance.
Despite these capabilities, current state-of-ther offthheart diffusion-bran based methods face drawbacks: high- quality samples require a large number of model evaluations, which tends to be computationally expensive, and they are sensitive to hyperparameter choices and the lack regularity of the data distribution [67, 35]. Naturally, computational cost, and sensitivity to hyperparameters and data regularity are not unique to diffusion models. Many numerical analysis and optimization algorithms grounded forward discretizations of continuous-time processes share similar limitations. For
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

instance, when optimizing smooth convex function f. the go-to algorithmic solution, gradient descent (GD), which updates 2 f z=z) represents forward Euler discretization of gradient flow å - fxx) [1]. GD is known to converge when the learning rate h is on the order of I where I. is the gradient Lipschítz constant. Consequently, when f approaches nonsmoothness (i.e., is large), GD must proceed with very small steps, thus converging slowly.
Backward discretization schemes for continuous-time processes offer potential remedy these issues. For example, the backwart ward Euler discretization for gradient flow updates iterates via the implicit equation x- = x hV hVf (z=). This algorithm known as the Proximal Point Method, and, remarkably, it converges for any positive h [62]. The tradeoff is that 24 is defined implicitly, requiring solving an auxiliary sub- problem that is potentially as expensive as the original one. However, in the context of generative models, where f represents the log-density of data distribution, one could learn such an update directly from data via neural networks (instead of learning the score), making the per-step complexity of a backward discretization comparable to that of a forward discretization. This motivates the central question of this work.
Can backward discretization improve diffusion models?
In this paper, we give a partial affirmative answer to this question. We develop a diffusion sampler based on an implicit, backward discretization scheme of the reverse-time stochastic differential equation (SDE). Our approach leverages recent results that allow us to train proximal operators from data via proximal matching [23]. We study two variants: one relying exclusively on backward discretization, and another hybrid approach that combines backward and forward terms. Our methods are competitive with score-based diffusion
inceudiinn those based on probability flow when using& large number of discretization steps, while significantly outperforming them with fewer steps. Additionally, we provide supporting convergence theory that yields state-of-the-ar rates under the idealized assumption of having access to the exact implicit update or proximal operator. Employing backward discretization of stochastic processes has been studied in both earlier [60, 5, 20, 79, 66, 42, 68, 29] as well as more recent [11, 48, 22, 49, 85, 26, 39, 50, 58, 80] works, mostly in the context of proximal version of Langevin dynamics. To our knowledge, no backward discretization has yet been explored for diffusion models. We defer more thorough discussion of related works to after the introduction of the necessary background. Code for reproducing the experiments in this work is available https:// /tttp:://gthubb com/ whwnnganaanggg//eeeeee.aaaaaaaee
www.gghaanaaan/PoDMM..... ProxDDM.
2 Background
Diffusion and score-based generative models. Diffusion models perform generative modeling by posing reversed stochastic differential equation (SDE) of process that transforms sample from the data distribution, Xo Po. Xo E Rd, into Gaussian noise [72, 27]. A widely adopted SDE for the forward process is given by
dX, (t) /B())W, EE [0.
(1)
where We is a standard Wiener process in Rd (i.c. standard Brownian motion) [75]. At inference time, one can produce new sample from po by reversing the SDE above, resulting in the stochastic process
dx, = [- - - ((t)Vlnpp dt + V /(()dW,
(2)
where W, is reverse-time Brownian motion (with time flowing backwards from T to 0). This process can be simulated as long as the score function, st(x) V (x)), can be computed (or approximated for unknown real-world distributions).
The continuous-time time process in (2) must be discretized to be implemented. By far, the most common strategy is applying the forward Huler-Maruyama discretization, leading to the score-based sampling algorithm:
X X2 + 7k X + V Iipp, (Xx)] + T
FEuler-aauyyama
where T/k = B((e)(1e - 4-11 and 25 (0, T). Many variations of these processes have been proposed. For instance, instead of the prreavcc- formulation of the process in (1), one

---

could instead consider a variance-exp exxllodiii alternative (by removing the drift) [72, 75]. Likewise, instead of considering stochastic processes, one can find deterministic ordinary differential equations (ODEs) that reverse the forward process and match the marginals of the reversed SDE, resulting in probability-flon ODE [75. 70].
In order to apply these sampling algorithms for generative modeling on real-world data, corresponding score functions must be obtained for every distribution pt- The popular score matching approach [32, 72, 74, 59] learns parametric regression functions s0 t) (with parameters 0) by minimizing the objective
-8 v1 Ct 22)
X(t)
so (X,,t)
(3)
No.E
where Xo Po. 8 (O a = exp (- ((s)ss) x, VaiXo + vI a,e, and X(t) is function that weighs different terms (for different 1) appropriately. It can be shown [75] that the solution to such an optimization problem results the desired score, s* (ar, t) = V In pt (x). There are two important remarks about this approach: (i) such function amounts to denoiser-thal is, it seeks to recover a signal corrupted by (Gaussian) noise' -nnd (ii) such a denoiser is a minimum mean-squared error (MMSE) estimate, as it minimizes the reconstruction error with respect to squared Euclidean norm. This can also be seen alternatively by Tweedie's ' formula [21, 61], which relates the score to the conditional expectation of the unknown, clean sample var namely Xl X2] X, + V In Po (X;), where = vI - ar.. Such conditional mean is the minimizer of the MSE regression problem [57]. showing that MMSE denoisers provide direct access to the score of the distribution pt- Fortunately, the design and development of denoising algorithms is very mature, and strong algorithms exist [57]. As an example, it is standard to employ convolutional networks with U-net architectures [65] for score-matching in image generation [27, 75].
Proximal operators. Proximal operators are ubiquitous in optimization, signal and image process- ing, and machine learning. For functional f: E + I. the proximal operator of f is mapping from IRd to R2 defined by
f arg min f(u) + g soin-xpp x112.
Intuitively, the proximal operator finds a point that is close to the anchor, x. and has low function value in f. It is convenient to introduce an explicit parameter À E R++ to control the relative strength of the functional f, and consider prox (x) instead. In this way, p + x when À whereas prox y (x) - arg min, f(u) when x - co. Further, when f L-smooth and S 1 first-order optimality conditions yield that x+ = prox2)(2) if, and only if, x+ = x AV which illustrates its relation to backward discretizations.
Importantly, proximal operators can be equivalently described as critmuum----osteriinn (MAP) denoising algorithms under Gaussian corruption model. Indeed, it is easy to verify² that when y = x + E, for x ~ Po. Gaussian noise & ~ (N, I²),, and f In po. then prox (v) provides the mode of the conditional distribution of the unknown x given the measurements y; i.e.,
prox x2 (v) arg max (Xl))-
(5)
x
This view of proximal maps in terms of MAP denoising algorithms should be contrasted with that of the score, which relies on MMSE denoising algorithms: unlike the latter, the former does not require the function f = In Po to be differentiable nor continuous." 3 This has been extensively exploited in nonsmooth optimization [3, 14, 18, 19], leading to speedups in convergence. The advantages of employing proximal operators (instead of gradients) for sampling have also been recognized in earlier works [60, 5, 20, 79, 26, 39], leading to faster convergence or less stringent conditions on the parameters of the problem, like step sizes. More generally, MAP estimates are known to provide samples that are closer to the support of the data distribution than MMSE denoisers [16]. In the following section, we will harness proximal maps to develop a new discretization of denoising diffusion models, leveraging MAP denoisers instead of MMSE ones.
Moor precisely, (3) seeks to equivalently estimate the noise that has been added to the ground truth X, XX||X(/222" /(2o2)) exx(-f(X)) and In is monotonic.
2(5) follows since po (Xlp)) x P(yoX)pp(X)) x exp(-1ly-Xp x||2
It simply requires f be weakly-conver or prox-regule [63].
3

---

Algorithm I Proximal Diffusion Model (ProxDM)
Input: Map prox_, pn (.), function 5(-), grid {0 = to < ti <... < IN =T}. boolean hybrid.
Output: Sample Xo-
Generate X (o,
for k N to
Generate 21e
10) and compute Yts f- S/ B(s)ds
if hybrid then
Update Xx-1 f prox- In pe [(1 + \176)) X + Tk2k
else
Update X f prox 2nk Impre-1
2-76 (Xx + v
return Xo
3
Proximal-based generative diffusion modeling
Recall from Section 2 that all score-based generative models apply forward discretization to the reverse-time SDE, resulting in updates of the general form in (Guler-Maruyane with the score playing central role. We now propose different avenue turn this SDE into a sampling algorithm.
3.1 From backward discretization to proximal-based sampling
Consider now the backward discretization of the reversed-time SDE from (2), that is:
X = X 7k [ X + In (Xx-1)
VTk
(6)
Unlike the forward (Culer-Maruyanie discretization, this is an implicit equation of X is easy to verify that this update corresponds to the first-order optimality conditions of the problem min, g()),
where g(u)= := 2%k In Pt+ (u) +
2-70 (X + Th=a) precisely proximal step. Such
2-76
an update can be written succinctly as
2
X prox
(Xx VTEEE)
(PDA)
la Pty
7/k
We refer to this step the proximal diffusion algorithm (PDA). It remains to define the discretization of the step Ts: recall that for forward discretizations, one can take 7 = ((xx) Here, we instead make the choice of setting _, (ss)ds. Note that both are equivalent when At,- - 0.
It is worth remarking that the discretization in (6) is fully backward discretization: one where both terms depending on the time-continuous (t) are discretized as - Yet, one can also consider hybrid approach by considering the updates given by X X2 7 [ Xx + V In Pta-, (Xı-1)) +
VTk2k, zk. where both X₃ and X are employed in the discretization of the drift term. In this case, the resulting update can be concisely written as
X = prox - In In [(1 + 1 Xx + V
(PDA-hybrid)
The resulting hybrid approach is akin to a is complementary to the update in (PDA), and will have advantages and disadvantages, which we will expand shortly. For now, we have everything we need to formally define our sampling algorithm for Proximal Diffusion Models (ProxDM), presented in Algorithm 1.
method in nonsmooth optimization [3].
Remarks, First, we note that this technique for deriving proximal versions of diffusion models via backward discretization is quite general and widely applicable to other variants of diffusion sampling schemes. We have shown this idea for variance preserving (VP) SDE and the Euler-Maruyame solver, but the same can be done for the reverse diffusion sampler and the ancestral sampler (see [75, Appx. EJ), as well as variance exploding (VE) SDE and probability flow ODE [75]. In each case, one would obtain a new proximal-based sampling algorithm that mirrors the diversity of score-based sampling algorithms (see further discussions in Appendix F).

---

Compare now the update in (PDA) to the score ( update in (Euler- Maallr-Maruyymaa,, First, both algorithms call denoiser of the data distribution each step; the score-based method uses an MMSE denoiser based on gradients, whereas PDA uses a MAP denoiser via proximal operators. Second, both updates consist of a denoising step and a noise addition step, but in switched order: the forward version applies denoising first, followed by noise addition. As result, the final samples from score-based models do contain small amounts of noise, requiring an additional, ad-hoc denoising step at the end to mitigate this and improve sampling performance (see e.g. [75, Appx. GJ). In contrast, PDA performs denoising after noise injection. This avoids the need for extra modifications and yields more principled and general scheme.
Given the benefits of backward discretization. why should one bother with hybrid approach (PDA-hybrid)? As we will see in the next section, (PDA) enjoys faster convergence rates. However, this comes at the cost of limitation in the minimal number of sampling steps: one can verify from (PDA) that we must ensure N < 2 (as otherwise the regularization parameter, 20u,, would 2- be negative), resulting in an upper- bound on the step size and thus a lower- bouud to the number of sampling steps. The hybrid approach in (PDA-hybrid) does not suffer from this hard constraint, resulting in practical benefits, albeit at the expense of slower theoretical rate.
The reader should note that the connection between backward discretization and proximal- algorithms is not new and has been explored in different contexts across optimization and sampling. For example, while forward discretization of the gradient flow leads to the gradient descent algorithm, backward discretization corresponds the proximal point method [62]. In the context of sampling, forward discretization of the Langevin dynamics yields the (unadjusted) Langevin algorithm (ULA), whereas backward discretization leads to the proximal Langevin algorithm (PLA) [60, 5, 20, 79]. is interesting to note that, in this context, proximal algorithms that leverage backward discretizations often provide benefits that the simple forward ones do not have: they are often more stable to the choice of [60, 79]. A less related but interesting connection between forward-backwar of measures) and sampling has also been studied in recent works 177, 42, 68, 11, 80), and can provide unbiased asymptotic samples (which ULA does not provide) [80]. These advances, however, have all been provided in the context of Langevin dynamics and, to the best of our knowledge, our PDA approach in Algorithm is the first instantiation of these ideas in the context of diffusion processes.
and can do away with requiring differentiability of the potential function
methods (in the space
Lastly, score- based methods need the gradient of the log- density at time tx. V In Pta, which can be easily estimated via score- matchin.. On the other hand, both proximal updates in Algorithm I require prox- A, in - estimate such proximal operators.
i.c. the proximal operator for the log-density at time f Next, we show how to
3.2 Proximal matching for proximal operator estimation
Learning data-dependent proximal operators has only recently begun to receive attention [30, 56, 47, 23]. We will make use of recent work by Fang et al. [23], who presented general approach to obtaining learned proximal networks (LPNs). They present a that ensures the recovery of the proximal of the log-density" of continuous distributions at arbitrary regularization strengths from i.i.d. samples from that distribution. In a nutshell, mroximal-matchinn augments the random samples X Po with Gaussian noise, Y = X + ac, for ~ (O, Ia). Then, it performs a regression using a special loss function and schedule to learn a MAP denoiser for the data distribution, with the noise level dependent the desired regularization strength À in the proximal. To learn prox_ one minimizes the problem Ex ExExxx (FoM()o(Y), x: with the moximal-matching loss
loss function
(p, yi () = exp (-
[12 112-112 dr2
Fang et al. [23, Theorem 3.2] showed that, in the limit of < + O, the problem above recovers the proximal operator of the log density of the distribution with regularization strength o²: i.c.
"The work in [23] also proposes a specific parametrization of neural networks based on gradients of convex functions (such as those arising from ICNNs (2]) that guarantees exact proximal operators for any obtained network.
'We remark that the form of 4pM used here differs slightly from the definition in [23], but it does not change the optimization problem nor its guarantees.
5

---

fa. = -prox- o22 In po Intuitively, the proximal matching loss can be interpreted as smoothed version of the in pseudo-norm. In one dimension, minimizing the -io-diiferentiable) &g loss recovers the mode of a distribution. The proximal matching loss generalizes this idea to higher dimensions and differentiable functions, making it amenable first-order optimization methods. Fig. 6 visualizes the proximal matching loss along with the mean-squared error and & losses to illustrate their difference.
Here, we extend proximal-matching to a setting that will enable us learn a set of proximals for varying densities and regularization strengths, as needed for ProxDM: the set {prox_ la for range of t and X's. Our generalization of the proximal matching objective for proximal diffusion model is
0- = arg min E { {Ex, (eMM (Jo(Y;t, A), X2; O))
where X, Pt and YIX,- (X,, Ala) are clean and noisy samples for the forward process, respectively. Intuitively, with sufficient data and model capacity (just as in motre-nathhnn networks), the optimal solution to (7), denoted by fo- (2: t, A), approximates prox- Alpp, (2) as Ç + O.
3.3 Practical considerations and implementation details
Although this proximal matching framework is conceptually simple, translating it into practice requires careful implementation. Herein, we present a few important methodological details.
Sampling of (1, A) pair. Optimizing (7) requires sampling (4, A) pairs, and the sampling distribution should ideally match the distribution to be used after training. As in score-based models, reasonable choice for the time step t is uniform in [0, 7]. The values of is less obvious, however. Recall that this regularization strength depends on the step size (7k) and, thus, also time tx-1 and fx. We adopt a heuristic for sampling (t, A) according to the candidate number of steps to be used at inference time (see Appendix E.2 for further details and motivation), which in practice enables the trained model to be applied for a range of different numbers of steps. We note that this sampling scheme is not universally optimal, and we leave the optimization of sampling schemes to future work.
Parameterization of the proximal network fo(;;; t, A). We parameterize the proximal operator prox An p. using a neural network denoted by fo(:;t,), A). Notably, while the score network is conditioned on a single time scalar [27, 75], the proximal network is conditioned on two scalars, t and A. Similar to score networks, conditioning is implemented by adding learned embeddings for t and À to the intermediate feature maps. Moreover, the score network indirectly parameterizes the score
via s0(x, t)
-ca(a.t)
[27]. where (2 t) is the (learned) score and eø(x, t) is neural network.
V1-ac
Analogously, we parameterize the proximal indirectly by letting fe fz; A) where eo (z; f, A) is a neural network that predicts the residual of the proximal operator: 60(x; t, A) = [x fe-Joole So(2; I, A) This parameterization led to better empirical results at large step numbers in our experiments and coincides naturally with the balancing of objective across different choices of (t, A)'s which we discuss next.
Vāe t, A),
Balancing contributions of different (t, A) in the objective. Score matching uses weighing function to balance the loss at different times (see [75, X(t) in Eq. (7)l). For proximal matching, a proper weighing should also be used for balanced learning, and the weighing should be based on both and X (c.f. (7)). A general rule of thumb would choose these weights such that the objective of the inner expectation in (7) has similar magnitude across different and X. Score model achieves this by weighing according to the square of the magnitude of the target in the objective [75, Sec. 3.3]. which is reasonable when the magnitude of the target can be estimated based on f and when the loss function is (2- Jemooomeuoouss... Unfortunately, none of these are true for proximal-matching Thus, instead of weighing the objective, we transform the target to align its magnitude across different (t, A)ss Specifically, we rewrite the problem in (7) follows (but without changing the objective):
Y - :t
Y - Xx,
0- = arg min E Ex Ex,EVIx PPM
(8)
VA
VX
where : Y-X+ is the new target. Since Y ~ (X,,, AL), we have e ~ NN(, f). Notably, the mag- nitude of is independent of (t, X), allowing for natural balance. Plugging in the parameterization
VA

---

of fo) = Vxcal ('co(2t) (2; t, X), the final objective reduces to:
0- arg min Ee, [xxx le [PM (co (X: vλe; t, 2.2) se; s)]
(9)
Recall that is neural network, (a 1, A) is the learned approximation of prox-A p. (2), A1 are sampled according to a predefined distribution, and Ç is a hyperparameter that should decrease through training. Note that this objective is similar to that of diffusion models with analogous param- cterization [27, Eq. (12)). We include pseudo-code for proximal-matching training in Algorithm 2, noting that it simply boils down to optimizing (9) in stochastic manner, allowing for any solvers of choice (stochastic gradient descent, Adam, etc).
It is worth comparing the resulting method with score matching. First, score matching samples Gaussian noise only once (to compute X²). In contrast, our approach samples noise twice: once to construct X₂ and once for proximal matching training. Second, score matching aims to learn the MMSE denoiser for Xo (more precisely, vaxXxo with noise level v - att, whereas proximal matching searches for the MAP denoiser for X, with noise level vÀ. Lastly, the loss in proximal matching is parameterized by Ç, which decreases via a scheduler, whereas score matching uses a fixed MSE loss throughout.
4 Convergence theory
In this section, we establish a convergence guarantee for Algorithm under the idealized assumption that the proximal operator we use is exact instead of learned. Our results provide a bound on the number of steps needed by the two discretizations
baaksedd-bawwwarddd and hybrid to achieve a KL-divergence of less than E. We consider the canonical Ornstein-Uhlenbeel (OU) process as the forward process, that is (() = in (1), as is standard in both practical implementations of diffusion models and theoretical analyses of convergence. Further, we consider a uniform time grid 4 kT/N /N. With this choice, letting h- = T/N, the step sizes " in Algorithm reduce " = 2h.
We now state an informal version of our main result. The formal statement requires additional technical details that we defer to Appendix C. Recall that po is the target distribution.
Theorem 1 (Informal). Suppose that Ep I|2 is bounded above by M2 < cc. Further, assume for all t 2 O, the potential In Pt is a three-times differentiable function with L-Liipsciiz gradient and --piscchitz Hessian, the moment Ep, Ep, I In pe (enp)((()))) ))? is bounded above by Cp < co, and certain
regularity conditions hold, Set the step size and time horizon to satisfy h S 8L+1 and T 2 0.25. Let 4r be the distribution of the output of Algorithm 1. The following two hold true.
(Hybrid) If the hybrid d flag set to true, then
(L(pooler) S (d + M2)e hTdL2 + hTT [[d
M2 + +x2) H2 L
(Backward) If the hybridflug is set to false and further h S 1, then
(L(pol) S (d + Ma)e + 2T [(d + M2 + C2) + d + I
A few remarks are in order. First, the uniform bounds on H and Cp can be relaxed to a much weaker bound on an appropriate running average on t-we elaborate on the more general assumptions in the formal statement of the result. Second, the regularity conditions mentioned above-made precise in the appendix- are necessary to guarantee that certain boundary conditions are met. Interestingly, similar conditions seem to be implicitly assumed in many existing results [6, 79, 41, 78], and defer a longer discussion to the appendix. Third, note that we do not require the distribution to be log-coneave, nor do we impose that it satisfies the log-Sobolev inequality. Finally, the proof strategy is inspired by the analysis of the Proximal Langevin al algorithm from [79], in combination with techniques from [6, 40, 41, 78]. To better illustrate this result, the following corollary bounds the number of steps necessary to achieve an E-accurate distribution in KI. divergence
Corollary 1. Consider the setting of Theorem 1. Moreover, assume that M2 S d. Cp S dL? The following two hold true.
(Hybrid) If the hybrid d flag is set to true, then (L(poll S provided that (dL* + d' H2H222 1/2
d
and NZ dZdL2
T d d
L2 1/4
T > log
+
8
8
7

---

(Backward) the hybrid flag is set to false and further h S 1, then (poolll S provided that
2 log d and N [d + + + + 1/4
E
The dimension dependence M2 and Cp is standard and is, e.g., satisfied by Gaussian distributions. With this result, the benefits of the full backward discretization become clear: the hybrid discretization requires () steps, while the backward discretization requires only (dds steps. This benefit comes at the expense of the additional condition that h < 1, since otherwise the regularization coefficient of the proximal operator would become negative. This translates into hard lower bound on the number of steps one can take, we discussed above and will see shortly again in the experiments section. In comparison, the most competitive rates for score-based samplers in terms of convergence in KI. divergence are ( for vanilla method [6, 4] and e for an accelerated strategy [81]. Note that these rates are not directly comparable, as the underlying assumptions and algorithmic setups (e.g., smoothness, boundedness, or early stopping) differ across works. In particular, [6] assumed L-Lipschitzness for the score function, similar to our work but without the additional smoothness assumption on the Hessian as used here. On the other hand, [4] lifted the smoothness assumptions entirely but instead analyzed the KL with respect to slightly noise coise-perturbed version of the target distribution. We point interested readers the broader body of work on the convergence theory of diffusion models, including both SDE-based samplers [7, 40, 41, 6, 4, 44, 45, 43] and ODE-based samplers [9, 8, 46]; see [44, Section 4] and [43] for comprehensive review.
5
Experiments
Target distribution
Score Standard Fuler
Scure
Denoising, epe=le-3
10a
Score, Standard
Score, Decoising, epe=le-3
Waserstein 2 Distance
Score, Denoising epa=le-5
PDA-Hybrid (Ours)
(Ours)
Score
Demoising g spa-1e-5
PDA-Hybrid (Ours)
10-a
NFE
Figure I: Sampling from a distribution supported on a discrete set points 2D using exact score and proximal operators. LEFT: the target distribution and samples generated by various samplers using sampling steps; RIGHT: Wasserstein-2 distance between sample and target across varying step numbers (NFE, number of function evaluations). Standard Euler- Marle--aruyymaa without denoising fails when using only 5 steps. Score-based sampler with a denoising step at the end, as is common [75], requires tuning the step size for the last denoising step. Both PDA variants perform well without additional hyperparameters.
Synthetic example. We begin with an example in cases where the scores and proximals can be computed arbitrary precision. We Po uniform distribution over discrete set of points (dataset from [55]), and run score-based diffusion sampling, as well as (PDA) and (PDA-hybrid) (see Algorithm 1). We include both the standard Euler-Maruyama method ("Standard Euler"), as well as the popular modification that includes an ad-hoc final denoising step after the last iterate [75]. As shown by the results in Fig. both variants of ProxDM provide faster convergence with a closer final sample to the ground truth Po. and require less hyperparameters to tune and no ad hoc modification. To further study the impact of the smoothness of the distribution on sampling performance, we perform an experiment to sample from 1D distributions with different levels of smoothness at the boundary of their support (see details and results Appendix H. Although the relation between sample quality and smoothness of the distribution is relatively subtle-likel due to smoothing effect inhere in the forward diffusion process-it clear that ProxDM recovers the support of the distribution more accurately than score-based samplers.
8

---

steps
steps
steps
Som R0G
Score
SDE
7
2
3 8
5
G 2
/ 3 S D 7
tecceeet (then)
B
0
6
3
6 3
7
O
S
8
Soore ODE
3
9 5.
3
/
7
3
9
5
2 7
/ 7
7 6
3 7 5
6
9
9.
2
/ 7 5 0 2
7
5
3
7
7
6
7 5
7
3
7
3
7
Phot Hybrdd
(3urs))
4
6
9
7
7
7
7
6
5 5
6
2
7
6
6
7
4
5
0
6
6
3
7
3
12
Figure 2: Left: MNIST samples generated by the score SDE sampler, score ODE sampler, and our hybrid ProxDM. Right: the resulting FID score as a function of the number of sampling steps.
We move to real, simple case for digits [17]. where we must train score- nettennetworks (through score matching) as well as our learned proximal networks for ProxDM. The 3(t) in forward process is set linear, ((t) Buin + (ámax Bmin) for E [0, 1], with Bmin = 0.1, fmax 20, as is standard [75]. ProxDM is pre-trained with an & loss, followed by proximal matching loss with Ç 5=1 1 in the first half and decayed to 0.5 in the second half. Both networks are trained for the same number of epochs (see more details in Appendix E.3).
MNIST.
We compare with two score-based sampler: Euler-Maruyama for the reverse-time SDE (denoted as "SDE") and the Euler sampler for the probability flow ODE (resp. "ODE") [75]. We show the FID score by NFE (number of function evaluations, i.c. sampling steps) and uncurated samples in Fig. 2 (right and left, respectively). ProxDM provides significant speedup compared to both the score samplers based the reverse-time SDE and probability-flow ODE.
CIFAR10. We perform similar experiment on the CIFAR10 dataset [37]. The score model uses the same U-Net as in [27]. The network for ProxDM is the same as the score network, ex- cept that it has two parallel conditioning branches for t and respectively, with comparable total number of parameters. Score matching follows the standard setup in [27] with learning rate warm-up, gradient clipping, and uniform t sampling as in- troduced in [75]. For proximal matching, the proximal network pretrained with & loss, followed by proximal matching loss with Ç = 2 in the first half and Ç = 1 in the second half (see more details Appendix E.3). Fig. shows the FID score vs. NFE curves. It is worth noting that the score SDE sampler's FID at 1000 steps matches the value reported in the original paper [27]. It is clear that our ProxDM outperforms score-based samplers and enables faster sampling. Notably, the FID of ProxDM at 10 steps matches that of the score ODE sampler at 20 steps and surpasses the score SDE sampler even at 100 steps. Uncurated samples are shown in Fig. 4, demonstrating sharper and cleaner samples by ProxDM with fewer sampling steps. Furthermore, when trained specifically for small step numbers (5, 10 and 20), the advantage of ProxDM over score-based samplers becomes more significant (see Fig. 3, dashed lines). If one is specifically interested in a very small number of steps, disregarding some of the heuristics described in Section 3.3 can lead to even better results (see Fig. 9). Additionally, we report the precision and recall metrics [38], which measure the quality of generated samples and coverage of the real data distribution, respectively (see detailed definitions Appendix E.3). The results are presented in Fig. 10, further illustrating the improved sample quality and coverage achieved by ProxDM. Finally, the FID of ProxDM 1000 steps appears worse than that 100 steps-we discuss on this in Appendix A.3.
Score SDF
Scoce ODE
Hyhrid (Ours)
10a
10t
10
10?
10
NFE
Figure 3: FID vs. number of sam- pling steps (NFE) on CIFARIO. The dashed lines correspond to models trained specifically for 5, 10 and 20 steps. as opposed to the full- range models in solid lines.
Ce (256 x 256). We further evaluate our approach for high-resolution image synthesis using the CelebA-HQ dataset at 256 x 256 resolution [34]. Because of the higher dimensionality of this data, both the score and proximal approaches operate in the latent space of a pretrained VAF, following the standard latent-diffusion setup [64]. As shown in Fig. 11, the hybrid ProxDM achieves improved FID compared to the score SDE sampler and performs competitively with the score ODE

---

10
1000
Score O0E
Dres Relinid (Ohor))
Figure 4: CIFARIO samples from score SDE, score ODE, and hybrid ProxDM samplers.
Score SDE
Soore ODE
PreeDM
Figure 5: CelebA-HQ (256 x 256) samples from score SDE, score ODE, and hybrid ProxDM samplers using 20 sampling steps.
sampler across all step counts. Uncurated visual samples using 20 sampling steps provided in Fig. 5; samples at other step numbers can be found in Fig. 12.
Conclusions and limitations
Diffusion models represent a paradigm shift in generative modeling. Since reversed-time SDE drives sampling, discretization strategies are needed to turn these into practical algorithms. To our knowledge, all diffusion models have far relied on forward discretizations of these SDEs, relying on the distribution score at each step. Here, we have shown that backward discretizations can be greatly beneficial, both in theory and practice. For the former, our full backward ProxDM requires only Ödd VE) steps to produce sampling distribution E-away from the target in KL divergence (assuming perfect proximal steps). In practice, show that neural networks trained with proximal matching allow for practical methods that generate accurate samples far faster than alternative score-based algorithms, and even faster than score-based ODE samplers.
Our work has several limitations. First, our thcoretical analysis assumes access to an oracle proximal operator. While we demonstrated that such models can be approximately obtained in practice, incorporating this discrepancy (say, at an a-accuracy level) as done in other works [7, 6, 31] remains open. Second, convergence results also require additional regularity conditions. While these arguably mild (see discussion the appendix), we conjecture that these could be removed with a more careful analysis. On the practical side, provided first instantiation of learned proximal networks for ProxDM, but many open questions remain, including how to choose the optimal weighing strategy for different sampling points across training, optimal ways to schedule the Ç hyper-paramettr the sampling time.
loss, and employing networks that provide exact proximal operators for any
10

---

Acknowledgment
ZF acknowledges support from the Johns Hopkins University Mathematical Institute for Data Science (MINDS) Fellowship. MD was partially supported by NSF awards CCF 2442615, MD and JS were partially supported by DMS 2502377. ZF and JS have been partially supported by NIH Grant P41EB031771.
References
[1] Luigi Ambrosio, Nicola Gigli, and Giuseppe Savare. Gradient flows: metric spaces and in the space of probability measures. Lectures in Mathematics. ETH Zürich. Birkhauser Verlag AG, Basel, Switzerland, 2 edition, December 2008
[2] Brandon Amos, Lei Xu, and J. Zico Kolter. Input convex neural networks. In Proceedings of the 34th International Conference on Machine Learning (ICML), pages 146-155, 2017.
[3] Amir Beck and Marc Tehoulle. A fast iterative shrinkage-threshoolinng algorithm for linear inverse problems. SIAM journal on imaging sciences, 2(1):183-202, 2009.
[4] Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly d-linear convergence bounds for diffusion models via stochastic localization. arXiv [stat.ML], August 2023.
[5] Espen Bernton. Langevin monte carlo and JKO splitting. arXiv [stat.CO], February 2018.
[6] Hongrui Chen, Holden Lee, and Jianfeng Lu. Improved analysis of score-based generative modeling: User-friendly bounds under minimal smoothness assumptions. arXiv [cs.L.G), November 2022.
[7] Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv [cs.LG], September 2022.
[8] Sitan Chen, Sinho Chewi, Holden Lee, Yuanzhi Li, Jianfeng Lu, and Adil Salim. The probability flow ode is provably fast. Advances in Neural Information Processing Systems, 000668552-8557-- 2023.
[9] Sitan Chen, Giannis Daras, and Alex Dimakis. Restoration-degraaatioo beyond linear diffusions: non-asymptotic analysis for ddim-type samplers. In International Conference on Machine Learning, pages 4462-4484 PMLR, 2023.
[10] Tianyu Chen, Yasi Zhang, Zhendong Wang, Ying Nian Wu, Oscar Leong, and Mingyuan Zhou. Denoising score distillation: From noisy diffusion pretraining one-step eigh-quality generation, 2025. URL https: arxiv. org/ abs /503 07578.
org/abs/
[11] Yongxin Chen, Sinho Chewi, Adil Salim, and Andre Wibisono. Improved analysis for proximal algorithm for sampling. In Conference on Learning Theory, pages 2984-3014. PMLR, 2022.
[12] Lee-Shin Chu, Sudeep Sarma, and Jeffrey J Gray. Unified sampling and ranking for protein docking with DFMDock. bioRxivorg, page 0024.09.27.615401, September 2024.
[13] Hyungjin Chung and Jong Chul Ye. Score-based diffusion models for accelerated MRL Med. Image Anal., 80: 102479, August 2022.
[14] Patrick L Combettes and Jean Chnn-Chisttppe Pesquet. Proximal splitting methods in signal processing. Fixed-point algorithms for inverse problems science and engineering, pages 185-212, 2011.
[15] Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. DiffDock: Diffusion steps, twists, and turns for molecular docking, arXiv [q-bio. BMJ, October 2022.

---

[16] Mauricio Delbracio and Peyman Milanfar. Inversion by direct iteration: An alternative to denoising diffusion for image restoration. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https: openreview wwwp/eeeeeeneneeee nett:foruum..dddWyyPPP.......... Featured Cer- tification, Outstanding Certification.
[17] Li Deng. The MNIST Database of Handwritten Digits. http://yann... lww.nneenn.. www.encoo/eeeeeeeeeeeeeeeeeeeeee www.cooo//cceeeeeee mnist/,200 2012.
[18] Dmitriy Drusvyatskiy and Adrian S Lewis. Error bounds, quadratic growth, and linear conver- gence of proximal methods. Mathematics of Operations Research, 43(3):919-94,, 2018.
[19] Dmitriy Drusvyatskiy and Courtney Paquette. Efficiency of minimizing compositions of convex functions and smooth maps. Mathematical Programming, 178:503-5555 558, 2019.
[20] Alain Durmus, Eric Moulines, and Marcelo Pereyra. Efficient bayesian computation by proximal markov chain monte carlo: When langevin meets moreau. SIAM J. Imaging Sci., 11(1):473-566 January 2018
[21] Efron. Tweedie's formula and selection bias. J. Am. Stat. Assoc., 106(496) December 2011.
[22] Jiaojiao Fan, Bo Yuan, and Yongxin Chen. Improved dimension dependence of a proximal algorithm for sampling. In The Thirty Sixth Annual Conference on Learning Theory, pages 1473-1521, PMLR, 2023.
[23] Zhenghan Fang, Sam Buchanan, and Jeremias Sulam. What's in a prior? learned proximal networks for inverse problems. In International Conference on Learning Representations, 2024.
[24] Kevin Frans, Danijar Hafner, Sergey Levine, and Pieter Abbeel. One step diffusion via shortcut models, 2025. URL https: /tarsi///////v//// org/ hbtgg//b///// abs//a/s//2 12557.
[25] Zhengyang Geng, Mingyang Deng, Xingjian Bai, Zico Kolter, and Kaiming He. Mean flows for one-step generative modeling, 2025. URL
org/abs/ 2bs/ 13447.
[26] Andreas Habring, Martin Holler, and Thomas Pock. Subgradient langevin methods for sampling from nonsmooth potentials. SIAM Journal on Mathematics of Data Science, 6(4):897-925, 2024.
[27] J Ho, A Jain, and P Abbeel. Denoising diffusion probabilistic models. Adv. Neural Inf. Process. Syst., 2020.
[28] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet. Video diffusion models. Neural Inf Process Syst, abss22004.03558:8833 8646, April 2022.
[29] Liam Hodgkinson, Robert Salomone, and Fred Roosta. Implicit langevin algorithms for sampling from log-concave densities. Journal of Machine Learning Research, 22(136):1-30, 2021.
[30] Samuel Hurault, Arthur Leclaire, and Nicolas Papadakis. Proximal denoiser for convergent plug-and-pla optimization with nonconvex regularization. In International Conference on Machine Learning, pages 9483-9505. PMLR. 2022.
[31] Samuel Hurault, Matthieu Terris, Thomas Moreau, and Gabriel Peyré. From score matching to diffusion: A fine-grained error analysis in the gaussian setting, arXiv preprint 1611 2025.
[32] Aapo Hyvärinen. Estimation of non-normalized statistical models by score matching. J. Mach. Learn. Res., 6(24):695-709, December 2005.
[33] James M Joyce. Kullback-leibler divergence. In International Encyclopedia of Statistical Science, pages 720-722. Springer Berlin Heidelberg, Berlin, Heidelberg, 2011.
12

---

[34] Tero Karras, Timo Aila, Samuli Laine, and Jaakko Lehtinen. Progressive growing of GANs for improved quality, stability, and variation. In International Conference on Learning Representa- lions; 2018. URL https: WtRLhhtts:///ppeee //tppr/eee//eereeeee... www.:/ooouuddddlloooaaaaa.
[35] Tero Karras, Miik Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generative models. 2022. URL htttt::/taaatae /rg//bs///////// ab2//20/.222.... 020.//226000004....2 00364.
[36] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. DivfWave: A versatile diffusion model for audio synthesis. arXiv [eess.AS), September 2020.
[37] Alex Krizhevsky. CIFAR-10 and CIFAR-100 datasets. https: https://www/ //wp::://ww...../ www.ontttttooooooooooooooooo.... edu/ "kriz/ciffe cifar. htm., 2009.
[38] Tuomas Kynkäänniemi, Tero Karras, Samuli Laine, Jaakko Lehtinen, and Timo Aila. Improved precision and recall metric for assessing generative models. Advances in neural information processing systems, 32, 2019.
[39] Tim Kiz Lau, Han Liu, and Thomas Pock. Non-log-concen concave and nonsmooth sampling via langevin monte carlo. Advanced Techniques in Optimization for Machine Learning and Imaging, 61:83, 2024.
[40] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence for score-based generative modeling with polynomial complexity. arXiv [cs.LG), June 2022.
[41] Holden Lee, Jianfeng Lu, and Yixin Tan. Convergence of score-basee based generative modeling for general data distributions. arXiv [cs. September 2022.
[42] Yin Tat Lee, Ruoqi Shen, and Kevin Tian. Structured logconcave sampling with a restricted gaussian oracle. arXiv [cs. October 2020.
[43] Gen Li and Yuling Yan. O (d/t) convergence theory for diffusion probabilistic models under minimal assumptions. arXiv preprint /895994499//959, 2024.
[44] Gen Li, Yuting Wei, Yuxin Chen, and Yuejic Chi. Towards faster non-asymptotic convergence for diffusion-based generative models. arXiv [stat. June 2023.
[45] Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating convergence of score-based diffusion models, provably. arXiv preprint arXiv:2403.03852 2024.
[46] Gen Li, Yuting Wei, Yuejie Chi, and Yuxin Chen. A sharp convergence theory for the probability flow odes of diffusion models. arXiv preprint arXiv:2408.022 022320, 2024.
[47] Lingxiao Li, Noam Aigerman, Vladimir G Kim, Jiajin Li, Kristjan Greenewald, Mikhail Yurochkin, and Justin Solomon. Learning proximal operators to discover multiple optima. arXiv preprint
2022.
[48] Jiaming Liang and Yongxin Chen. proximal algorithm for sampling. arXiv preprint arXiv:2202.13975 2022.
[49] Jiaming Liang and Yongxin Chen. Proximal oracles for optimization and sampling. arXiv preprint arXiv:2404.00223, 2024.
[50] Jiaming Liang, Siddharth Mitra, and Andre Wibisono. Characterizing dependence of samples along the langevin dynamics and algorithms via contraction of d-mutual information. arXiv preprint arXiv:2402. /7067, 2025. URL https:// org/ abs/ 17067v2.
[51] Yaron Lipman, Ricky T. Q. Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023. URL ptttpp:///ppereeeeee wtt.eerrrieeeeene/fruumuummme
[52] Xingchao Liu, Chengyue Gong, and qiang liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023. URL https:
13

---

[53] Cheng Lu, Yuhao Zhou, Fan Bao, Jianfei Chen, Chongxuan Li, and Jun Zhu. DPM-solver: A fast ODE solver for diffusion probabilistic model sampling around 10 steps. arXiv [cs.LG]. June 2022.
[54] Yi-an Ma, Niladri S Chatterji, Xiang Cheng, Nicolas Flammarion, Peter L. Bartlett, and Michael I Jordan. Is there an analog of nesterov acceleration for gradient-based memc? Bernoulli, 27(3): 1942-1992, 2021.
[55] Justin Matejka and George Fitzmaurice. Same stats, different graphs: generating datasets with varied appearance and identical statisties through simulated annealing. In Proceedings of the 2017 CHI conference on human famtofactors in computing systems, pages 1290-1294, 2017.
[56) Tim Meinhardt, Michael Moller, Caner Hazirbas, and Daniel Cremers. Learning proximal operators: Using denoising networks for regularizing inverse imaging problems. In Proceedings of the IEEE International Conference on Computer Vision, pages
[57] Peyman Milanfar and Mauricio Delbracio. Denoising: A powerful building- block for imaging. inverse problems, and machine learning. arXiv preprint arXiv:2409.06219, 2024.
[58] Siddharth Mitra and Andre Wibisono. Fast convergence of <-divergence along the unadjusted langevin algorithm and proximal sampler. In 36th International Conference on Algorithmic Learning Theory, 2025.
[59] Tianyu Pang, Kun Xu, Chongxuan Li, Yang Song, Stefano Ermon, and Jun Zhu. Efficient learning of generative models via fmite-difference score matching. arXiv [cs.LG], July 2020.
[60] Marcelo Pereyra. Proximal markov chain monte carlo algorithms. Statistics and Computing. 26: 745-760, 2016.
[61] Herbert Robbins. An empirical bayes approach statistics. In Proceedings of the Third Berkeley Symposium on Mathematical Statistics and Probability, Vol. I, pages 157-163. University of California Press, 1956.
[62] Tyrrell Rockafellar. Monotone operators and the proximal point algorithm. SIAM journal on control and optimization, 14(5):877-895 1976.
[63] Tyrrell Rockafellar and Roger J-B Wets. Variational analysis, volume 317. Springer Science
& Business Media, 2009.
[64] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björ Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10684-10695, 2022.
[65] Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In International Conference Medical image computing and computer-assisted intervention, pages 234-241. Springer, 2015.
[66] Adil Salim, Dmitry Kovalev, and Peter Richtárik. Stochastic proximal langevin algorithm: Potential splitting and nonasymptotic rates. arXiv [stat.MLJ, May 2019.
[67] Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusson models.
In International Conference on Learning Representations, 2022. URL net net/forumiid-ie iitrrummPdddIIdIIIhhhho.
[68] Ruoqi Shen, Kevin Tian, and Yin Tat Lee. Composite logconcave sampling with restricted gaussian oracle. arXiv preprint arXiv:2006.05976 2020.
[69] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. In 3rd International Conference on Learning Representations. Association for Computing Machinery, 2015.
[70] Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. arXiv [cs.LG), October 2020.
14

---

[71] Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models.
In The Twelfth International Conference on Learning Representations. 2024. URL https:
/ //openreview.
[72] Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. arXiv [cs. July 2019.
[73] Yang Song and Stefano Ermon. Improved techniques for training core-based generative models. arXiv [cs.LG], June 2020.
[74] Yang Song, Sahaj Garg, Jiaxin Shi, and Stefano Ermon. Sliced score matching: A scalable approach to density and score estimation. arXiv [cs, May 2019.
[75] Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-hased generative modeling through stochastic differential equations. arXiv [cs.LG], November 2020.
[76] Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. In Proceedings of the 40th International Conference on Machine Learning, ICML'23. JMLR.org. 2023.
[77] Michalis K Titsias and Omiros Papaspiliopoulos. Auxiliary gradient-based sampling algorithms. Journal of the Royal Statistical Society Series B: Statistical Methodology, 80(4):749-767, 2018.
[78] Santosh S Vempala and Andre Wibisono. Rapid convergence of the unadjusted langevin algorithm: Isoperimetry suffices. NeurlPS 2019, March 2019.
[79] Andre Wibisono. Proximal langevin algorithm: Rapid convergence under isoperimetry. arXiv preprint arXiv: 2019.
[80] Andre Wibisono. Mixing time of the proximal sampler in relative fisher information via strong data processing inequality. arXiv preprint arXiv:2502.056: 2025.
[81] Yuchen Wu, Yuxin Chen, and Yuting Wei. Stochastic runge-kutta methods: Provable accelera- tion of diffusion models. arXiv preprint arXiv:2410.04766 2024.
[82] Qinsheng Zhang and Yongxin Chen. Fast sampling of diffusion models with exponential integrator. arXiv [cs. April 2022.
[83] Yasi Zhang, Tianyu Chen, Zhendong Wang, Ying Nian Wu, Mingyuan Zhou, and Oscar Leong, Restoration score distillation: From corrupted diffusion pretraining to one-step high-quality generation. arXiv 2505. 13/77, 2025. URL htttts://aariiie http:///aaaww/aaaaaaa/aa two authors contributed equally).
13377. (the first
[84] Mingyuan Zhou, Yi Gu, and Zhendong Wang. Few-step diffusion via score identity distillation. arXiv /2674 2025. URL https: /ttp/v/axvvvv/vv//////eeeee//avl orggabs.// abs./50........../../ 12674.
[85] Xinkai Zhou, Qiang Heng, Eric C Chi, and Hua Zhou. Proximal meme for bayesian inference of constrained and regularized estimation. The American Statistician, 78(4):379-390, 2024.
[86] Bernt Øksendal. Stochastic differential equations: An introduction with applications. Universi- text. Springer, Berlin, Germany, July 2003.
15

---

Contents
1
Introduction
2 Background
3
Proximal-based generative diffusion modeling
3.1 From backward discretization to proximal-based sampling
3.2 Proximal matching for proximal operator estimation..
3.3 Practical considerations and implementation details
4
Convergence theory
5
Experiments
6 Conclusions and limitations
10
Discussions
17
A.1 Comparison to Existing Accelerated Methods of Diffusion Models
17
A.2 Connection Maximum Likelihood Training
17
A.3 Bottleneck of Convergence at Large Step Numbers
18
B Notation
18
Formal statement of the convergence guarantee
19
C.I Proof of Proposition
23
Proof of Lemma
26
Proofs of auxiliary lemmata
27
D.1 Proof of Lemma Interpolation SDE for the Backward Algorithm
27
D.2 Proof of Lemma Interpolation SDE for the Hybrid Algorithm
28
D.3 Proof of Lemma Second Moment of OU Process
29
D.4 Proof of Lemma 4 Bound on Initialization Mismatch
29
D.5 Proof of Lemma Bounds on the Diffusion Matrix G3 for Backward Algorithm
30
D.6 Proof of Lemma Bounds on the Diffusion Matrix G₁ for Hybrid Algorithm
30
D.7 Proof of Lemma - Bound on the Drift Difference μ4 fūy for Backward Algorithm
30
D.8 Proof of Lemma 10 Bound the Drift Difference 14 ÿāg for Hybrid Algorithm
32
D.9 Proof of Lemma 8 Bound the Divergence of Diffusion Matrix V Gr for Backward Algorithm
32
D.10 Proof of Lemma 11 Bound on the Divergence of Diffusion Matrix V G2 for Hybrid Algorithm
34
D.11 Proof of Lemma 13
34
D.12 A General Lemma for Time Derivative of KL
35
16

---

E Method and Experimental Details
36
E,1
Pseudo-code for Proximal Matching Training
37
E.2 Sampling and À
E,3 Image Generation
38
Extension of ProxDM Other SDEs and ODEs
38
G Connection between Backward Discretization and Proximal Algorithms
39
H Additional Results
39
H.1 Ablation on Smoothness Target Distribution in ID
39
H.2 Additional Results..
A Discussions
A.1
Comparison to Existing Accelerated Methods of Diffusion Models
Since the introduction of diffusion models, numerous attempts have been made to developing accelerated methods, including those based on deterministic samplers [70, 8). higher-order solvers [53, 82, 45, 81], consistency models [76, 71], distillation methods [67, 10, 84, 83), flow matching [2], and shorteut models [24]. Our aim not merely to accelerate inference within score-basee based models, however, but to propose a new and principled framework for diffusion models based on new discretization schemes with convergence guarantees.
In particular, we demonstrate that backward discretization schemes lead new class of diffusion models based on proximal operators (instead of scores), and show that this formulation leads to immediate gains in sampling efficiency compared to vanilla score-based SDE samplers. Given its simplicity and generality, the same principle can be naturally extended to existing accelerated methods-for example, by applying backward discretization to probability flow ODEs and flow matching. Our approach thus represents orthogonal direction prior advancements, offering general pathway to improve existing methods. Notably, our approach requires no distillation and is trained directly from data. Moreover, method enjoys provably favorable convergence properties over score-based counterparts.
Finally, we believe the findings of ProxDM could lead to novel understandings of recent empirical heuristics. For example, recent works have shown the benefits of training generative models using MAP-promoting losses rather than the conventional MSE loss [71, 25]. Our framework employs the proximal matching loss, which also seeks a MAP estimate, but arises naturally from the perspective of learning proximal operators. This connection offers new theoretical insights for understanding the effectiveness of such losses and may inspire future methodological improvements.
A.2 Connection to Maximum Likelihood Training
has been shown that minimizing the denoising score matching objective equivalent to maximizing lower bound on the likelihood of the observed data [27]. In contrast, our proximal matching loss does not naturally admit such a Maximum -ikiiumolikeiooodd (MLE) interpretation, since it is derived from a continuous Caati SDE perspective rather than the discrete Markov chain formulation. It remains open question whether such an equivalence exists between proximal matching and MLE training. Notably, since our ProxDM differs fundamentally from score-based diffusion models in both the transition kernel mnon-Gaussian vs. Gaussian) and training objective (non-MSE vs. MSE), we anticipate that establishing such equivalence is nontrivial-B at all possible.
Despite this, the samples generated by our method enjoy strong theoretical guarantees. In this work, we bound the KI. divergence between the samples generated by proximal diffusion algorithm and the true data distribution. Combined with existing theoretical results for proximal matching [23), our ProxDM provides provable guarantees for its generated samples.
17

---

A.3 Bottleneck of Convergence at Large Step Numbers
Our results showed that the FID score of ProxDM deteriorates at large step numbers (see Fig. 3). We conjecture that this behavior arises because ProxDM, unlike score- based methods, relies on solving an implicit equation that in turn defines proximal map. practice, we use proximal matching to train networks to approximate the corresponding proximal operators, rather than implementing exact proximals each step. As a result, the computed updates may violate the optimality conditions that define such updates, introducing inaccuracies that accumulate. Empirically, we observe that these inaccuracies appear more problematic than the approximation errors of scores in score-based samplers. Better network architectures and training strategies will provide more accurate and stable approximation of the proximal operators, leading to even higher empirical accuracy at high step counts.
B Notation
In this section, we introduce the notation we use and some necessary technical background.
Glossary.
Let us summarize the symbols used throughout the appendix:
Discretization step size.
T: Time horizon of forward process.
N: Number of discretization steps N = T/hh.
pr: Distribution of the forward process at time
Pr: Distribution of the reverse-time process at time t, i.c., P;= PT-t-
g: Distribution of samples from the algorithm time t. The density of the samples after steps is Skb- The density of the final samples is @r. The density of the initial samples go is standard Gaussian.
p(h): Distribution the reverse-time process within the k-th step of the algorithm, where [0, h). That is p(t) - Ph-k
k (0, 1} and
(k). Distribution of the interpolating process of the sampling algorithm within the k-th step (see Lemmas and and surrounding discussion). Note that q: (k) 4kk++-
Remarks on terminology. We abuse notation by using the same symbols to denote distributions and their probability density function. We call the distribution we wish to sample from the target or data distribution. We use the word "process" to refer to a stochastic process.
Linear algebra. The symbols Rd, Rdxdd St and Rddddd denote the set of real d-dimensional vectors, d d matrices. and d x symmetric matrices, and d d d tensors. We endow Rd with the standard inner product (x, y)
y. which induces the l2 norm via || / Var, x:). For matrix E IRdxd we let || Allop denote the operator norm induced by the lg norm and ||A|||r the Frobenius norm. In particular, if A is symmetric with eigenvalues X(A) E Id., then || max; |Ail and |||||r ||{())|-- Thus, ||А]| s VdliAllop- For symmetric matrix E S with positive eigenvalues (i.e.. positive definite matrix), we write || = / (rr Qx:) for the associated norm.
Calculus. Given differentiable map Φ: Rd R use Vi((x) to denote its gradient Jacobian at æ if = 1 or A 1, respectively. Similarly, given O3 function IRd + R, we use the symbols V=v(x) ddd and V²r((a) to denote its Hessian and tensor third-order derivatives at whose components given
V2y(F); 0 p(r) and
0 (7((z))
V
dx,0x,
Ōr,0xx Ozz
respectively. We denote the Laplacian of via
8 P(p(z)
Ar(x) = Tr Tr(V)(((x)
Σ
R..
du
18

---

The symbol V. denotes the divergence operator, which acts on vector field v: R - I2 via
Σ (x)
V v(x)
R.
dz,
Note that with this notation in place we have the identity (Vp) Ap. For a matrix-valued function G: _² [gddd and differentiable function Rd + let Tr(V)p(2)G(z)) e R
denote the vector whose i-th component is T((V.9yp(r)G(m)) where
a Or,
@)p((z) Oz,0z,0rk jk
V,
ERaxd
We use the symbol V - G' to denote the divergence operator broadcast to the rows of G, i.e.
V G =
-V- -I V.a(t) Rd
where Gx(x)= E _d is the i-th row of G(x). For brevity, write G in place of G(z) when the argument clear from context.
Finally, given a function p: X' + y between finite dimensional vector spaces and a norm-one vector we define the directional derivative of & with respect to v the point x as
<p(x: + sv) ppxx)
Vo((ax)
lim sJ0
u
Further, given a canonical basis fei,..., Cdim x) of X we let Vap be shorthand for Ve, 4 for k E (1.2.... dim }
Lipschitzness of the Hessian. For a function f: IRd + RR we define
Hf(x) max [<v.v> op
ii= 1
ER.
(10)
The condition (x) < oo is weaker than a global Lipschitzness for the Hessian." Indeed, if
[(x) v - H| vll for all x, Rd.
then for all Rdd
v.)' < H for all Rd ||aa| 1,
which implies
(()) H for i 1, 1.2.....d.
Comparing the above with (10), we can see that Hx) is weaker than in two aspects: (i) the quantity H (x: is local, and (ii it only concerns the directional derivative with respect to canonical basis vectors e;, not all directions. We use this definition because it suffices for our convergence analysis. We define the shorthand
(x(x) Hin (x)
for t IR.
(11)
where Pt is the density along the forward process at time t.
C Formal statement of the convergence guarantee
Before introducing the formal statement, we sketch the proof strategy, which motivates some of the technical assumptions. We present the sketch here regarding the fully backward discretization, but the strategy for the hybrid method is analogous. This proof strategy is based on the analysis of Proximal Langevin Algorithm [79].
6179 assumes that the Hessian of the target potential has global Lipschitz constant.
19

---

We consider the OrsstennUUleenbee (OU) process as the forward process:
dxX; -X,dt + 22dW,, Xo ~ Po. t [O,T],
(12)
and use Pr to denote the distribution of X, For simplicity, we will consider reverse-time processes with time flowing forward. That is, although the original reverse time process evolves via
dx, I-X, 2V In pr (X.)) + v2dW, where & flows from ''
(13)
and We is the backward Brownian motion, we consider the more intuitive
dX; Xx + 2V In pr dl + /2dW.
(14)
where flows from 0 to T. We use P.= pr-t for the distribution of the process in (14) distinguish from that of the process in (13). We consider the backward discretization XEΛ evolving via the recursion
(K+1I - RRh = 1& 1-18&((42) (K+1J + 2hV In PT-(k+1) (R(-11)) + 2hzzk
(15)
here goes from 0 to N -. 2 ~ (0, Ia) and Xo ~ (V), This is equivalent to ProxDM with the hybrid flag set to false (with the only difference that & traverses from 0 to N - 1 as opposed to the reverse). When the hybrid flag is set to true, the discretization reads
{(K+IJA
NEs = h&rr + 2hV In (K++)) (X(++1))) + 22h2k.
(16)
Our goal is to bound the KL divergence (L (l between Pr, the target distribution (note that pr= Po), and @r, the distribution of Ny. If Xy defined by continuous time process with sufficient regularity, then one could apply the fundamental theorem of calculus to get
! OKL(Pp l10). dt.
KL.(Pr 11 KT. (1
(17)
Ōt
In turn, bounding the derivative of the KL divergence time evolves would yield bound on the quantity of interest. To execute this plan in our discretized setting, we construct interpolating stochastic processes flowing from ARt. X(K+1JA. Define the stochastic processes
(X/(*)
[EJO.N]
given by
= = x800 +x((() + 21V W (x.-
(8((
V2W,
(18)
starting at X2/00) = Ses for all k. The next lemma gives an explicit formula for the SDE governing the dynamics of this stochastic process; we defer proof to Appendix D.
Lemma 1 (Interpolating Process for Backward Algorithm). Suppose that h S 2L+1 discrete algorithm in (15) has continuous interpolating process evolving according to the SDE
11 Then, the
dx(()) - (8)( ,1) di v 20x (8!1)) .t) dW.
10,
(19)
where
(x( = [11 -t)la- 2 InpT 11
and
① v (n Pa(x,t) = (o((z,t) x 2 + 28 at
+
We shall compare the distributions of these processes with the equivalent pieces of (14), i.e.
(x(()) = + 2V In pr pT -h-
X (k)
dt + /2
dW,
LE h0
(20)
To derive a bound on the time derivative of KI. divergence in (17), we leverage the Fokker-Planck equation in tandem with boundary condition that allows us to apply an integral by parts. This motivates the assumptions we impose next.
20

---

Assumption (Regularity of the potential). The following hold.
(Smoothness) There exists L 1, such that for all t Z O. In three-times differentiable and L-smooth, i.e. the gradient V lnpp pt is L-Lipschitz.
(Data second moment) The expected energy M2 == Epo |2]1|1] < co is finite.
(Score second moment) The running average Cp == + S Ep Vl (2) < co is finite.
(Lipschitz Hessian) For any x and I, let He(ar) = max [VV I /
1..,.....
The running average 2² :
+ Jo' (x) co finite.
We highlight that we do not require the distribution to be log-concave, nor do we impose that it satisfies the log-Sobolev inequality. Smoothness of the potential and bounded second moments standard assumptions in the literature [6, 81J. Lipschitz continuity of the Hessian is not as common for diffusion models, but has been used to study Langevin-type methods [54, 79].
Additionally, we assume regularity conditions for the analysis of time derivative of KL. We first state it in its general form.
Assumption 2 (Regularity of KL time derivative). Consider the SDEs
dY: = ((Y, + (2,(Y,,))W,
and
Ny. = (Prtt) di + 2 (Pr,t) (Pr.t)at,
where Y, Y, E Rdd and f.f: Rd x R - Rd, and Rd S Let Pe and be the distributions of Y, and Y, respectively. The following two conditions hold.
(Boundary condition) For all t,
Jes
In Pt V (p:J) - pr]] do= =0
Pt
(v:J) - v18 dr = O.
and
Here, we write Pt. Ub, S. J, and J in place of (x),, vs (x), (xx 1), Ma, 1), Me, 1) and M(r 1). respectively. (Dominated convergence) There exist two integrable functions 0, R: Bd - R such that
(c(x)) In p((x) 24 (x(x)
S ((x)
and
ot Pc(c) S k((x)
Or
for all I and almost every xr.
To analyze the backward algorithm (ProxDM with hybrid flag set to false), we instantiate Assump- tion 2 for the eeverse-time and interpolating processes.
Assumption 3 (Regularity Condition for Interpolating Process of Backward Algorithm). For
each {O,...,NN 1}, let μas Ga, and x((k) be defined as in (20) and let Pe be the density of x( Similarly, let fux. Gk, and (k) be defined as in (19) and let q;
=(k)
(k) be the density of X
Assumption 2 holds for f = pk, fi.. J Gr. Gt. -G ((K) and
The boundary condition required for applying an integral by parts the analysis of time derivative of KL (see Appendix C.2). Similar conditions are implicitly assumed in proofs in prior works [6, 79, 41, 78]. These conditions are met when the functions inside the divergence operator decay
RReaal that St denotes the set d x symmetric positive semidefinite matrices. By integrable mean f ((z)dr co and f k(x)dx= 30.
21

---

sufficiently fast at infinity, which is satisfied by distributions with regular tail behaviors. The domination condition is required apply the Leibniz integral rule in the analysis of time derivative of KI. (see Lemma 20).
To state analogous assumption for the hybrid algorithm we require a slightly different interpolating SDE, which we introduce in the next result. The proof of this lemma is deferred to Appendix D.2.
1 Lemma 2 (Interpolating Process for Hybrid Algorithm). Suppose that h < 21," Then, the discrete algorithm in (16) has continuous interpolating process evolving according to the SDE
dxaa = - t; x00)) dt 20h (x/1) .t) dW
t [0,
(21)
where
Gx(2,t) = -d- 21
- In in and
(=));;;a) ((n(z,))) (a
en- 2a+2VnnpT 112VVnnTTaa + 21 v In
in cl))
01
-2 Tr (e
With this hybrid interpolating process in place, introduce analogous regularity conditions.
Assumption 4 (Regularity Condition for Interpolating Process of Hybrid Algorithm). For each - 1}, let Pk, Gk, and (K) be defined as in (20) and let Prjo ((k) (. 1 a) denote the k E NO....,N-1 conditional density of x(x) given x(k) = Similarly, let Jug, GR, and x( be defined as in (21) and let 4/jo (.) a) be the conditional density of ((()) given 2( = a. Assumption 2 holdsfor f = pkk
t
fax, 1) (3, fo a), J = Gk, J Gr,, Pt 0 (11a) and v Gej0 a).
(k)
Equipped with these assumptions, we are ready to state our main result. Theorem 2 (Convergence guarantec). Suppose that Assumption I holds. Set the step size and time horizon to satisfy BL+4 and T 0.25. The following two hold true.
(Hybrid) If the hybrid d flag is set to true and Assumption 4 holds, then
KL (Pr ((rr) I (d+ M2)e -T + hTdL? 1T [(d + M2 + Cp) + & 1.
(22)
(Backward) the hybrid flag set to false and Assumption 3 holds, then
KL (Pr ((r)) & (d M2)e +12 [(d M2 + C,) + H2 + +
(23)
Proof. The proof strategies for both hybrid and backward are completely analogous and follow from the identity (17). The two proofs diverge slightly in how we upper bound the derivative of the KL divergence. The core of the argument appears the next proposition, whose proof deferred to Appendix C.I.
Proposition 1. The following two hold for any (O....,NN N 1}.
(Hybrid) If the hybrid d flag set to true and Assumption 4 holds, then
(<() KL.
Ph |/4h KL
(k))
(K)
Po
S h2 122 1-' Egg ||^||| + || In p((a) ||2] dt
+ (h2dd + n2dd /''' Byt -n x 2 dt.
(24)
22

---

(Backward) If the hybrid flag is set to false and Assumption 3 holds, then
KL
(=) Ph
KL
Po
(k)
110²))
2 L²
[ E7 ||=||2 + Vin (k)
Pe 11a
df + 11 / Bppe ( dt.
(25)
The main difference arises from the first term in the upper bound of the hybrid case. From now on, we will continue only with the hybrid scenario, as the rest of the argument analogous to purely backward discretization. Taking the sum of the inequalities in Proposition I from 0 to N 1 yields
N-1
KL (Pr |(rr) KL. (Po 1140) =
Σ k=0
KL
(k) Ph
14
KI. (TO-) 1190 (
S Nh?dL + t1? 1'" BF, IV P, [1 dt (26)
+ ( + +1) J. Ez. [H(())]] dt.
To finish the argument, invoke the following two results, whose proofs standard, but include them for completeness in Appendix D.3 and Appendix D.4, respectively.
Lemma 3. The second moment of Pe bounded
Ep, |= 2(d + M2) for all t E [O,T].
Lemma 4 (Lemma 9 in [6]). Let qo be the standard normal distribution in d dimensions. For 2 0.25, we have
KL (19o) s 2(d + M2)e 2I
Substituting these two into (26), using N T/h and Assumption I and reordering yields
KL (Pr ((T)) $ (d + M2)e 2T hTdL? h?TL² (d M2 C2) (h²d² htd TH?
Thus, (22) follows. The argument for (23) is analogous, which completes the proof.
The rest of this section is devoted to proving Proposition 1.
C.1 Proof of Proposition
Once more, the goal is to use the identity (17) by bounding the time derivative of the KL discrepancy, The crux of this proof lies in the following lemma, which we will instantiate twice, once for the hybrid discretization and once for the backward discretization. We defer proof to the end of this section (Appendix C.2).
Lemma (Time derivative of KL between SDEs). Suppose Assumption 2 holds.
Then, for all we have
d [L(pll)er) = -Ep
Vn P:
(s-J.vin P 1/
+ Ep
dl
V/
(J--)VVnp (J- V ( |
+ Ep,
(27)
where the arguments for f. F. J and are omitted.
We consider two cases. We start with the fully backward discretization since the argument is slightly more straightforward.
23

---

Case 1: Suppose that hybrid flag is set to false. Fix k E (O. defined in Lemma 1. First, we state some properties of Gk with their proof deferred Ap- pendix D.5.
- 1]. Let X, Th and Gk
Lemma 6. Let Gn defined as in Lemma 1. Then, we have that
(1- Gx 3 (1 + 18tL) and (1- (1-2LL)1= = V Gh (1+ 6tL)1.
We instantiate Lemma 5 with f = Mk. f ik J Gk. and J = Gg.. In which case, Pt =p; and and the lemma yields
(A)
Vt = q:
KP
(k)
ll
(
dt
(2)
2
K M4 (mm- Ma: Vin
(K)
P
P.
-Bp(4
V
In
+ p/
qt
6
Oi
(k)
+ E
(Gxx - 1)V In (k)
+V (Gxx - 1), v
In
Pi
(k)
((a)
1 E
||
< P(N)
V In
+ 2E 7(00
+ 22 32Bp( (N))
1)V In
2
gr
(A:)
|| (Gx 1111
P1
+ 2B700
V
In
8
91
2 2
[1o - أ|| + CG - 1)
p |1
+
Gx
(28)
where the second relation uses
Schwarz and Young's inequalities to bound (a, b) S 2||a||2+ |16| and the fact that || ||2 < [| since _ T Gj thanks to Lemma and the constraint on
t S h < +. The following two lemmas will help us bound the first and last terms in the expected value in (28). Their proofs deferred to Appendices D.7 and D.9, respectively.
Lemma 7. For all k E {O......A N 1}. the difference between and Pa is bounded by
||m p < 2 363 | 2 2187 | P( (w)))P) 24300 1 L (z)
Lemma 8. For all is E {O...., N 1}. the norm of V Gg is bounded by
||V Gx(x)) s 364 Hd xx-())))(((x))
Furthermore, Lemma 6 gives ||6 T1I2p - (18tL)² = 324 122 Substituting all of these bounds into (28) gives
d ( dt
(k) l1 ) E | V p. (z) + ( HL
(1)
Upper bounding I and integrating from to yields the desired inequality (25); finishing the proof for the fully backward discretization.
Suppose that hybrid id flag is set to true. Let {() Pr and G. defined as in Lemma 2. The proof of this case is slightly different: we use the following chain rule [33, 6], before applying the fundamental theorem of calculus
Case 2:
KL Ph
(k) 11 - KL
(K) (pb
1140
[
pjjo (-1 (1 I| 1 (- a))]
(29)
(R) (- a) denote the conditional density of (k) given {2(R) a, and p/10 (- denote the where 4:j0
conditional density of x in (20) given x((²)
Notice that can express the KL divergence
24

---

on the right-hand side by applying the fundamental theorem of calculus in a similar fashion as in (17). (
However, in this case the value at t = 0 satisfies KL
(8) (- a) ll
Yojo (- a)) = KL (6% || 6g) = 0.
Thus, the right-hand side is bounded by E an ['t'
d KL Pdoo (1aa) ll ( (- a) at Next, we
dt
bound the time derivative. Paralleling the proof in Case 1, we introduce relevant properties of Gs: their proof is deferred to Appendix D.6.
Lemma 9. Let GR be defined as in Lemma 2. Then, we have that
(1 (1-411) (1-4IL)= Gk (1 + 12IL)T and (1-211)] 3 6ox (1 + 4IT)T.
We invoke Lemma 5 with f = Mk. f(x, t) ta(x,t;a), J = Gk, J = G2,, Pt -P:0 (- a) and (k)
14- 2: (- aa) The lemma yields
KL (A)
Pilo (- a) ||
(k) 9ejo (- a))
dt
(k)
P0jo (- a)
= =- (-1a)
V In
4:jo (- a)
(b)
P Pijo
+ (t)) a)
Mk msa.
V
In
(8)
(- a)
9:10
((0a-
(K)
Pejo ( (1))
(()) do
(1)a) + v (ōx (Gk-Gi) Vin
+ pitt (1a)
(1a)
v
In
40 (.1 a)
(k)
Applying 1 3 G. thanks to Lemma 9-in tandem with the analogous argument we used for the purely backward method (28) we derive
d dt
KL
(a)))
(-1a) ll 200 (- a)) Pelp
(-
|l - 12 + (Gv- b/0 (. (1a)]] +
+/4 a.ll]
$ 2E 2E (8) (-1a)
(k) Integrating over E [0, hi] and taking expectation over ~po we have
KT. (. a)) 4 (-)a)
(k)
(8) [a))]
aa P p
h E
[lex 2al1 + +1(Gx-1) V In P/o (x a)
[a)²
+ || Gn 1 dt.
$2 0
2n (pp
(30)
The following two lemmas will help us bound the first and last terms in the expected value above. Their proofs are deferred to Appendices D.8 and 1.10, respectively.
Lemma 10. For all l: E (O....,N N - 1},, the difference between μr and fig is bounded by
Ilue= Tull < IT (8|=1) 20 V (np (x) + 48 4 --HT-R--- + 21 20
Lemma 11. For all & E {0.
N - 1}, the norm of V G. is bounded by
V Gn(w)|| < 192 192HH
25

---

Gx - < 12tL. Substituting all of these bounds into (30)
Furthermore, Lemma 9 gives that
gives
((k) (l1a)
(k) Unjo (- a))
E P
KL
1 L mNDt [||0|| + V kh- dt
11' 12 3()) Hh- (x)
+
(a.a)~ (tt) [lla- -|||=
dt
+ J. 1 (a.z)~ (1.))~(66")
In F
(xr a)) dt.
(31)
We now establish two lemmas that bound the last two terms on the right hand side above. Lemma 12 adapted from [6, Lemma 10J. We defer the proof of Lemma 13 Appendix D.11. Lemma 12. For all & {0, ,O.
N - 1} have
(a.a)- (P! (p [lla x1(3] I 4td + 2 7 [|<=|1²]
Proof of Lemma 12. For case of notation, we let 11 -k)
1/4 =P; and Mole =Po within this
(k)
proof. Applying the triangle inequality and (a + b)² < 2a² + 2b yields
E ((a.))((b).1))) [||a- z||3] < 2E Earsoj (-1 (-1r) [[]a - e 'xx1''] 3] + 2EE [|le x=13]
Next we bound each term this sum. Since Voje( (. x) = NN(e* xx - e e-2 we have
Exgny (o-rop((12)) (-1x) [lla - e 'x [ = d(1 - -e-2 S 2td.
For the second term, factorizing and using the fact that (1 - 0- < 1 for all 2 0 gives
E Egwar |le ' 2||
t)² E Egwaer
[||:|||| S ta Enwer [||=:|]]]
Combining the previous three bounds finishes the proof of Lemma 12.
Lemma 13. For all l {0. .O....,NN . 1} we have
2 E
| p/o (or a]I*] S IdL +
20000 [|] ih
(a.)~ (78 770)
Applying Lemmas 12 and 13 into (31), using the chain rule (29), and recalling that t S h and L 2 yields inequality (24), finishing the proof Case 2.
This concludes the proof of Proposition 1.
C.2 Proof of Lemma
Proof. By the Planckk equation,
dv:
V (vJ) vef
Øt
0p:
V AV (ptJ) ptf]-
We now use result on the time derivative of KL from prior work [78, Section 8.1.2]; see also [79, 6]. For completeness, provide full statement of the result in Lemma 20 and a proof in Appendix 1212, which relies the domination conditions. By Lemma 20, the time derivative of KI. satisfies
dt kI(P)))pt) = = =ʃ(
Oppt In Pt Ou: Pt dr.
Øt 2t
Øt vs
26

---

Plugging in the Fokker-Planck equations,
d KT(pellus) dt
/ v V (ptJ) Ptf]In Pr
-v. v. V (v.))) vvs]
Pt
dx
V
24
(p:J) p.f. V + v. (unJ) v V P: 1/.
} dx
(32)
where in the second equality we used integration by parts, which holds because of the boundary conditions. In the rest of the proof, we will transform the integrand in the equation above to show the desired result. For better readability, we omit the subscript Ps and v4 from this point onward.
For the integrand, we have
(DJ) pl, Vn 9 + V (wJ) us. v)
(DJ) pf. (v (vJ) uf. P V
(v (v- (pJ) Pf. / + ( VVV-vV) pf. Vn P
1/
1
(p ( 1) Vnn ( + ( (Lv (vJ) V (oJ), Vn L
(33)
v
V
where the first equality is by V = - The derivation above splits KT. into two terms: one 2t
depends the drifts f and while the other depends the diffusion terms J and J. Observe that the drift-related term matches the second term in the final result, which vanishes when f = f.
We now focus on the remaining. diffusion- related term:
L (1j) V (pJ), Vln ())
P (Jau+ -) JVp vV J. ")
(tvmu J JVhp J.V. e))
JV 0 JJ-JJV
(J-J), JJ, Vln ()
=p
= p |
P e
v f p (5 JJVmn++ (j- J). Vn ()
p
(34)
Here, the first equality is by V (w)) -vv+ + J and V (pJ) JVp + Pv and the second equality uses 1 Vv Inv and /v = V Plnp. Combining Eqs. (32) to (34) yields the desired result.
D Proofs of auxiliary lemmata
D.1 Proof of Lemma 1 Interpolation SDE for the Backward Algorithm
For ease of notation, within this proof omit the step index i: in (kk Tha and Gg. For t 2 0, let
X = (1- X, 2tV In - (x.)
so Xo = -No, and we can write (18) as
X, = X, V2W,
That is. X, evolves following the SDE
dx; V2dWe.
(35)
Suppose that the process (X.) evolves via an SDE of the form
ax: f (Xe.t) dt + + 2 (xe,t) dW,
27

---

for some functions f and G. In what follows, we solve for the form of these two functions. Let (=) (1- 2 In pr-pr-hha (h- R. Then, have
t)
UT:
8 V Inppr-t-((z))
01. (x) -- 2V In Pr pr- 28
01
2 xt (TT,(x) -2/V pT-
VT,(z) (1 (1-t) 1)
and
We now apply high-dimensional Itô's lemma [86, Chapter 4.2] to X, = T, (8.) and obtain the following
dx, dT, (X:)
PT. (8.) 2 dr,0z, dR. dx?
ÖT: (X.) dt VT (8))
dx, + Σ
0t
1,j=1
a 8.. 81
oT:
{x, dt VT, (x,) jutt
VT. (x.) v 2G dw, +
G
0t
dz,0o,
i,5=1
Oli
(2))
(T.( Tr ( (x.) G)) di + V2VT, (x.) V@aw.. (36)
at
where the third equality is by plugging in dX, = jut + v 2GdW, and noting that
dRidXi = dtt,
and the last step uses Tr(AB") = Eig Ar, Bij. Note that in the equations above, Tτ(-) E Rd is da E Rd, (-) E Ruxdxd, and TH(V"T(--G] e Rd are obtained by broadcasting the corresponding operator each entry in T((-).
vector and therefore, dT((-))
Matching (35) and (36), we must have
ar, (x) VTi(a)p + Th [2T](1)) 0,
V
which implies
G(z,t) = [VT(())]-
Mxx () = V
{
- Øt (x) (x) -T Tr [v
O..
Plugging in 0T (z). (T) and (T) completes the proof.
D.2 Proof of Lemma 2 Interpolation SDE for the Hybrid Algorithm
Observe that the next iterate of (16), & 8(++1)A, is the value at time / = h of the stochastic process
(X (k)
]ee[o.n] given by
1 xx(0) + + + 2tV In Inpr- V2W,
(k)
8 Nah-
(37)
For t 2 0, let
X, x((1 (((()))) 2tV Inpr
so Xo 8( and {(())
x
(k) /2W,
The rest of the argument is analogous to the proof of Lemma (see Appendix D.2), but with (x) xx 2 (mpr- Rd.
28

---

D.3 Proof of Lemma Second Moment of OU Process
This is standard property of the OU process. We include proof here for completeness. Along the OU process (12), we have
X, _ Xoo + V1- ~~N(0,11
Therefore,
=B ||e- Xo + V1 e e221
B
[1]
2e E || |2]
+ e " |211
2e "E [| 11=] + e
v 2E [1*] + 2d
where the second line follows by Young's inequality, completing the proof of Lemma
D.4 Proof of Lemma 4 Bound on Initialization Mismatch
Proof. Recall that p. denotes the density of the OU process at time t, thus pr =Po. We have
[ pe(x)) In (n
S (n
<T(p:l140)
Rd
[ pe(x) (npe(c)de / pr(x)) ||=|1 d
#2 2l((2T) dx
< J pe(x) In (x)dd En. In(2x).
d
For the first term, letting Pijo be the conditional density of X. given Xo. we have
J (r(x) (((x)lup((x) do Jec (fre ((o(e()))ooo)dd In (Le Pho(ale)) po(s)aty) de
< l U In polu)dy dr
J (ffz p(o(zly) In pajo dx po (o(yddy
where the inequality is by Jensen's inequality and that + x In x is convex function for x > 0.
Since X, Xo ~ (e ( e e-2() by entropy of Gaussian distributions, we have
[ (x)(()l)) In (x)()))) dr d In -221)
d
21) 2
2
Therefore,
d KL(Pdl10) - In(2 - 21)
d
d In(2x) Ep.
|||==11
2
2
1 Ep. |||=r|]
1
In
26 -1
c
$ M2 +
2
- e-
where the last inequality is by Lemma 3. Taking t' = 1/2 Inn --= thus
F 0.229, we have In
1
1-6 20
KT(pe (lq0) S M2 + d. By exponential convergence of the OU process, we have for all >
KT(pr(190) I e -2T -(T se -((--() (M₂ d) < 2e 2T (M2 + d)
as desired.
29

---

D.5
Proof of Lemma 6 Bounds on the Diffusion Matrix G₁ for Backward Algorithm
Since p is L-smooth,
L 14) LI.
So
(1 - 1 1 -n
3 (1 (1-t+ 2 2L = - 2+ since - 2(2L+1) Therefore,
Check that 0 1-11--2t
(1-1+211) 21 G [(1 - - In PT Pn 3 (1-1-211)
and
(1 21 t)Ta 2/ 10- 3 (1-1-211)-11
We collect some useful facts:
(1 x) $ -1
60 if0 x 2'
(1 - 21-1) 2 11 22 i
1
(1-x) (1=1)"'s1+2) 11+2r #050$5 2'
(1 (1-z)-1214
Since / 21T. 1, < + 21 S }, and I, 1
(1-4 = [] + 2t(1 21) = G N [1 + 6t + 2 = 18tL)
and
(1 - 11-) < [1 + (11 - 2 V G [1 + 2/ / 3 (1 + 6LL)
as desired.
D.6
Proof of Lemma Bounds the Diffusion Matrix Gg for Hybrid Algorithm
The proof is analogous to that of Lemma 6 (see Appendix D.5). The resulting bound is identical up changes in constant factors.
D.7 Proof of Lemma - Bound on the Drift Difference μa Íig for Backward Algorithm
We start by establishing a few auxiliary lemmas (Lemmas 14 to 16). We let [b,] E Rd denote vector whose i-th entry is bi. Lemma 14. The following two equalities hold true.
0
In = In + V In x) + A(Inpi)+ + d, and
Ot
In PT Wh- wh-r -Tr( -
(IVinpr-ua-ell)
((V + Inw- - [
-1)) 1=1
Proof. Applying the Fokker-Planck equation to (12) yields dy. dt
V [p, (Vip: + +x)] Therefore,
Op:
In
Ot
Pt Öt
V [pt (VInpt +z]
Pt
((Vpt, Vin + + (Vinp +x))
Pt
-V In, Vinpt) + Φ) + A(np:)) +
30

---

where the second equality uses V (V for f: R + R and G: R - R.. We have shown the first result.
Then, note that
In PT (n-
In Pr (I)
01'
M-T-khht
= VIIVinpr-Aa (Vimpr---n.)) A(mpr-Ah-o)
Therefore,
0 In (n - (W (A(InTT-h--))
V
Ir V Impr -h-
Recall that
IV (2) i=1
Adding the two equations above yields the second result.
Lemma 15. The following inequality holds.
S H In 11) i=1
Proof. For each entry, we have
Tr IIV
[ -ll
16x - =1(
d V. In In
< d &
Summing over the entries yields the desired result. Lemma 16. The following inequality holds.
i (h Tr(V" IIV- S 122 + L||=|| + 18dd
a
Proof. By Lemma 14, we have
v In pT PPr- (x) + Tr Inpr
Ot
In pr - 2))1 + [ - -1)] i=1
$ V V \npr- + V
From this point onward, we let = pT - for ease of notation. Next, we bound the three terms in this sum For the first term, we have
V In
V ((V In 1/, V VInn)) = 2V V
where the second equality uses v((f.g)) = JJ g + J where f.g: R + R", and Jy is the Jacobian of f. Therefore,
19) 2 2 1/ 11 V 1n < 2L Vn
To bound the second term, note that ((V In U2 x)) = Inn + 2 In 4 x. So
I ((Ving, x)) In m + V In v Il
| S V In m + L|
Lemma 15 gives a bound on the third term. Combining the three upper bounds finishes the proof of Lemma 16.
D
31

---

We are ready to prove Lemma 7. We start by decomposing fi
M4 as
Pk pik = (Vāx- -1)
(x 2V In Pn
+ 21
8
V Vinyy- + Tr ( (o
at
Next, we bound the two terms on the right-hand side. For the first term, we have
(x+ (Vox-1) (x 2V 2V-- (0+2 - /@. -
- 110p 1< + 2V Inpr-
S G177 (||=|| + + In -n
where the second inequality uses the bound || - 1
op
6t7 by Lemma 6.
For the second term, we have
21
G
Vin VInpr-m--(()) Tr (F
0t
V
Gx
i
(2) Tr ( hnp-r (2
s
Øt
< 56
(2L + 1) WinP--M--l
L||=||| + 18dd L L
where the second inequality is by Lemma and V Gk
S
from Lemma 6.
The desired bound follows from combining the above, noting that L 2 1 and using (a + bb 3a² + 3b² 3c for a, b, c R.
D.8 Proof of Lemma 10- Bound on the Drift Difference μ4 jb for Hybrid Algorithm
The proof follows the same structure as that of Lemma' (see Appendix D.7). We start by decomposing
Ak
Mk
Pk Ma = GV
(x 2
+ 21 G,
0 VInpT-m-- + Ir 01
( (v'mprr-mm
The first terms bounded analogously to the proof of Lemma (see Appendix D.7). The last
V
Gx
term is bounded by
< 2 thanks to Lemma 9, This completes the proof.
lop
D.9
Proof of Lemma Algorithm
Bound on the Divergence of Diffusion Matrix G for Backward
In this section, we prove Lemma 8, which provides bound V G+lll The derivation in this section is largely inspired by [79, Proof of Lemma 6), but we improved the technique to obtain a tighter bound in terms of d and with weaker assumptions.
Before proving Lemma 8, we show a few general lemmas related to matrix perturbation (Lemmas 17 and 18) and divergence of matrix-valuee functions (Lemma 19). These results then used to prove Lemma 8. The next two lemmas are folklore; we include their proofs for completeness. Recall that a matrix norm || is sub-mun mubtipllliicaaiie if for all A, B we have
|AB|| < ||].
32

---

Lemma 17. Let A and B be real invertible matrices. Let ||-|| be sub-multiplicative matrix norm. Then, the following three inequalities hold true
||A² - B² < 2 A B|| (I) + [B(),
|| - B- - - B|| ||1 |
and
-A-2-B-i 2 B|| (|| + ||B|] || B-
Proof. We start with the first inequality. By the triangular inequality and sub-multiplicatiiity. have
| - Ba| [ A - B) B²
(A B) B(A B) + (A- B)B]
< || - B1I + 2 B||
A combination of symmetry and the triangle inequality yields
||A- - B²1
< A - 3II + 2 || - B|| min (II ([B))
< A - B|| || + [B|| + min[[]]4]] II
S 2 2||A - B|| (1 All + ||
as desired.
The second inequality follows from the identity A-1 - B-1 = A- A)B-1 and the sub- multiplicativity of the operator norm.
Iteratively applying the first and second results gives the final inequality
- = || - B || s 2 ||A B (IIAII + (|BII) ||- |= |-
this concludes the proof of Lemma 17.
Lemma 18. Let A: IR" + R be a matrix-valued function. Assume at point x, A is continuous and differentiable, and [A(x)] exists within a neighborhood of Let II-Il be mub--uuliplicatiie matrix norm. Let (p(x)= = A(x) and fix norm-one vector V E R", then
[V)e((>)]| ||)(() [[((())) [[A{)))
IV.((=))
Proof. Let ((x)= = [x(x)] -2. We have
(p + uv) - ((x)
lim
u=+0
u
||x + wv) - (((x)]
lim
m00
u
21IA(x + uv) A(x) (IA(x + ue)ll + [I(())))) [l + uv)ll ||x(())
lim
||(( + uv) A(x)
4 (z))l [x])
lim u
= 4 ||A())]] II [[()))] |[x]]
where the inequality is by Lemma 17.
Next, establish general result that bounds the divergence of matrix-valued function by the norm of its derivatives.
Lemma 19. Let M: Rd + Rdxd be a matrix-valued function. (Note that M (xx does not need to be symmetric matrix.) Then, |/V (x)]? < d - (
33

---

Proof. Denote the j. entry of a matrix A by Aij. Recall that the & norm of each row or column vector of matrix is upper bounded by the operator norm of the matrix. Thus,
(My(aa)
d E (V,M())) & (V,
Ory
Hence, applying I|ull? z d|lu|l3 for all IRd
yields
2
(M;((z)
OM1;(x) dr,
V M(
Σ
<d
Ony
j=1
desired. This concludes the proof of Lemma 19.
Finally, we prove Lemma
Proof of Lemma 8, By Lemma 19, we have | S d v for each j = l,...,d. d,
We now claim that
V Gr op < 364 (M
which will imply the desired bound |/ G+ll S 364 -d Let T(x) (1 t)
21 In pn-p----(22 so that (xr = [T(x)] -2 We invoke Lemma 18 with v Ej (note that I1-Ilop sub-mulliplicative) to obtain
V, <1 | (x)
[V)T(>)
Note that both [[T and (z)lo)) bounded by constants. On the one hand, Lemma 6
ensures ||px|lop s 2.
11 On the other, since 2 pn sh- < T. we have
|T(x) S -1-t+ 21L S 1 + 21L <
where the last inequality follows since t 2(2++1) and so tL < 1/4. Plugging in the bounds for (x))l and [[)lo= gives
(x() s 182 (V 364t In S 364tH -
as desired; concluding the proof Lemma 8.
D.10 Proof of Lemma 11 - Bound on the Divergence of Diffusion Matrix V G. for Hybrid Algorithm
The proof follows the same structure as that of Lemma 8 (see Appendix D.9). The key difference is that we let
T(x) La 2/ In (np)-ss-(())
so that Gz(z) [x)]] The resulting bound only differs in the constant factor.
D.11 Proof of Lemma 13
(k) For ease of notation, let vo -p(h) Mi -P(() Mijo = and Moje =Polt within this proof.
By Bayes rule, we have
V (x (
a) = -vz ( (al x)) +n)++lmm(e ln))
34

---

Let p(2; Σ) denote the density of (p, Σ) at Then,
Poje(a z) Pr -hilT - gla;e (la;;e''x,
xx (1 E 1).
Therefore,
|| - e 21 2(1 e
V Vojoaa 2)) = V,
a e 'r)
Plugging the above into the Bayes rule identity, we have
(a
Va In a)
+ Va lnv(())
- e
Therefore,
((ox))((oiie)) [ In vnj( a)]l²] ((ax))((0))))
aa xx)
V In (1)
2
e 2
c (a E (x)
S ((0x))((00)1))
| In 12
2-
2
(1 - 11-e-2r e-22 Blaxs)(oa)))
a1a---001l + Ey, [1V In
Then, since vo x- x) = -N(e (1 - e I)) we have
((0,2)((vo11) |1a-c ela-e"'x''"" Erwny (a e|a---'x11""
= d(1 e -2)).
Therefore,
[[[V In ayo aa)l S L2 ʃ -t
L2 (((,z))(vo
Ey, (x)
1 - e-22
L 21
dL2
+ 12 L2, [| (x)]1
1 - e
Note that for < 2585 S SE S 3. it holds that
21
Therefore,
-2-
f2
2
The claim has been proved.
D.12 A General Lemma for Time Derivative KL
We restate the following lemma from prior work [78, Section 8.1.2]; see also [79, 6]. and include proof for completeness.
Lemma 20 (Time derivative of KL). Let Pt, : Rd + be probability density functions and 2 Suppose the following conditions hold.
1. There is an integrable function 0: Rd + R, i.e. f [((x)dx < oc, such that for all and almost every x.
@ Pp. In pr: < ((x)
at
There an integrable function k:: Rdd R. Le. f (s(x)do oc, such that and almost every
0 S xxxx) all
35

---

Then,
d (L(pelle))
- 1
(
In Pt
dq:
Ot dr.
dt
0t
qa
Proof. We have
d
(Tppell
dt ]
d
pt
In
Pt
de
dt
qt
0
Pt
Pr
In
dx
91
Pr Ppt
0
Pt
In
Pt 0t
In
da
Qt
9t
/
Pe 0pt
Pr 0q:
0pi
In
dx
qt. Øt
0pt
Pt
dq:
In
dr.
01
Ōt qr
Here, the second equality uses the Leibniz integral rule, which holds because of the first domination condition. The last step is by
Opc dr d Öt
dt / ppdx 0
which applies the Leibniz integral rule on pa thanks to the second domination condition.
E Method and Experimental Details
1.8
fr 0 0.2)
MSE
(2, 0.5)
61
(x, 0;5 1.0)
1.2
Loos value
1.0
0.0
Figure 6: One-dimensional visualization of the proximal matching loss lpM, in comparison with the mean squared error (MSE) and mean absolute error (E,) losses.
36

---

10
10
10-2
steps
-3
10 steps
100 sleps
1000 steps
0.0
0.2
0.4
0.6
0.8
1.0
Figure 7: Values of and À used by ProxDM-hybrid under varying numbers of sampling steps. Here, set ((t) = Bunin (Bmax Bmin) with 8min 0.1, fmax 20, as used in experiments.
E.1
Pseudo-code for Proximal Matching Training
Algorithm 2 Proximal Matching Training
Input:
Forward SDE noise schedule B(t) (c.f. (I)); sampling distribution q(t, A) over (4, A); data dis-
tribution po: annealing schedule ( for Ç; optimization algorithm Optimizations zontep (e.g., SGD, Adam)
Initialize network parameters 0g
repeat
k &- k+
C Anncalç(k))
D Update according to the annealing schedule Sample t, \ ~ qt, A), Xo Pu. and n. No0, I) independently
+ exp - 50 B(s) ds)
X, + / 04 Xu + v1 (4
Generate X, Pe
Y x Vīe
Compute loss frM t EPM (co, (Y; A), e: (k) 0x+1 OptimizationStep(( 0r,, VolPM) until convergence
Proximal matching loss (cf. Eq. (9))
Output: fo(2, = ((e((z;12))
b Approximated prox Aln (x)
E.2 Sampling and À
We provide more details of the sampling strategy for and used for proximal matching. To motivate our design, Fig. 7 displays the (t, X) pairs encountered by ProxDM-hybrid under different numbers of sampling steps: 5, 10, 100, and 1000. The values are evenly distributed [0, 1) regardless of the step count, whereas the range of À varies by the number of steps. Based on this observation, we design the following sampling scheme. We first sample step number N from predefined set of candidates. Then, for the chosen N, we uniformly sample one of the possible (t, X) pairs associated with that step number. This procedure repeated independently for each sample in the training batch and performed on-the-fly. For ProxDM-hybree hybrid, the step number candidates {5, 10, 20, 50, 100, 1000). For
we exclude and 10 due to the instability associated with small step numbers. To promote balanced learning across different step numbers, apply weighted sampling, with weights proportional to log(N) for MNIST and 11/3 for CIFAR-10.
37

---

E.3 Image Generation
For CIFAR-10, we adopt the same U-Net architecture as [27]. The ProxDM is trained for a total of 375k iterations with batch size of 512. We use the & loss during the first 75k iterations pretraining step, followed by proximal matching loss with = for 150k iterations, and then Ç = for the last 150k iterations. The learning rate is set to 10 For score models, follow the standard practice [27, 72, 75] of using a batch size of 128, and train for 1.5M iterations to ensure the number of training samples (i.e., effective training epochs) matches that of the proximal model. We set the learning rate to 2 x 10-4 and use gradient clipping and learning rate warm-up, following [27, 75]. For both score and proximal models, the last checkpoint is used, without model selection.
For MNIST, we halve the number of filters in the U-Net, following [72]. The ProxDM is trained for total of 225k iterations with a batch size of 512, using & loss with learning rate of 10 during the first 75k iterations, followed by proximal matching loss with Ç = and ( = 0.5 for 75k iterations each, with learning rate of 10-5. For the score model, we use a batch size of 128 following the standard practice [27, 72, 75] and train for 900k iterations to ensure the number of training samples (i.e., effective training epochs) matches that of the proximal model, using learning rate of 2 10"
Precision and recall metries. Following [38], precision measures the proportion of generated samples that fall within the manifold of real data. Specifically, real images are first embedded into feature space via a pre-trained VGG-16 network [69]. Then, a hypersphere is defined around each feature vector with radius equal to the distance to its ktth nearest neighbor, forming a volume that approximates the data manifold. Precision is then computed as the fraction of generated samples located within this volume, intuitively capturing sample quality and fidelity. Analogously, recall is defined as the fraction of real samples that lie within the manifold of generated samples, reflecting the diversity and data coverage the generation.
F Extension of ProxDM to Other SDEs and ODEs
In this section, we demonstrate conceptually how the proximal diffusion model can be extended to other SDEs and ODEs, such as the variance exploding SDE and probability flow ODE [75]. We leave the theoretical and empirical study of these extensions to future work.
Besides the variance preserving (VP) SDE in (I), alternative forward process used in diffusion models the variance exploding (VE) SDE [75], given by
dX2, = dW,
(38)
Here we omit the time-dependent diffusion term for simplicity. This corresponds to the standard Brownian motion initialized at the data distribution. The associated reverse-time SDE takes the form
dX,- = VV In ddW,
where time flows backwards from 7 to
Given a time grid (0 = to < h ... < tN = T}. applying forward discretization yields the following score-based algorithms widely used in generative models [72, 73]
Xx-1 Xx+ + Y Ptp (Xx) + T
where "/ t4 - and z NO0, I). Analogous to the derivation of ProxDM for the VP SDE, we can apply backward discretization to the reverse-time VE SDE:
X = X₂ + Th In Pta-, (Xe-1) + VTzz,
and obtain proximal sampling scheme:
X prox In 1 (Xx + vx2)).
Furthermore, consider general forward SDE of the form
dX, f(X,, dt + g())
38

---

where f: Rd x R - Rd denotes the drift coefficient and g: R - IR the diffusion coefficient. [75]
introduce the probability flow (PF) ODE, given by
dx; - (1200)---00 (00=((()))--000//mmm))))) di.
where time flows backwards from 7 to O. The PF ODE shares the same marginal probability densities as the forward process, and therefore, admits sampling algorithms analogous to those based on the reverse-time SDE. Notably, the DDIM sampler [70] has been shown to correspond to discretization of the PF ODE [67, 53].
In particular, for the VP SDE (1), have (ar,t) ()(())x and g(() /3(1), which leads to the following PF ODE
dx, - 8())x, - dl
1
Applying backward discretization yields
7k 2
X Xx X
v In Pta (Xk-1);
where = B(s)ds. This, in turn, gives rise to proximal-base based ODE sampler:
XR-1 = prox I-TH 1n
Xx
7/k
G
Connection between Backward Discretization and Proximal Algorithms
In this section, we provide the detailed derivations for previously explored connections between proximal algorithms and backward discretization, using two continuous-timg processes as examples: gradient flow and Langevin dynamics.
First, consider the gradient flow ODE
dr. = - (xx))d
whose forward (Euler) discretization gives the gradient descent algorithm
I At fxx)
where x and + denote the current and the next iterates, respectively. On the other hand, backward discretization yields the proximal point method [62]]
= x At (2++t)
x+ prox (x).
In the context of sampling, forward discretization of the Langevin dynamics
dX, = (XX,))d + V2dW,
leads to the unadjusted Langevin algorithm
X - + V2Atz, N(01))
while backward discretization leads to the proximal Langevin algorithm [60, 20, 79]:
x= A 22tz
V2A2z
x+ prOXA: V2Atz)
H
Additional Results
H.1 Ablation on Smoothness of Target Distribution in ID
We further investigate how the smoothness the target distribution affects sampling performance. Specifically, consider one-dimensional target distribution given by p(z) (0.5 + (0.5
39

---

x)",2 E [-0.5, 0.5], with varying smoothness parameters E 0., 2.0. This family starts from uniform distribution over [0, 1] with non-smooth boundaries (a 0) and becomes increasingly smooth discretization (without the additional final denoising step), while the proximal sampler follows the hybrid update in Algorithm 1. As shown in Fig. 8, both samplers recover the target distribution when given sufficient sampling steps. The dependence of sample quality on the smoothness parameter is relatively subtle-iieel due to the smoothing effect inherent in the forward diffusion process. Nevertheless, ProxDM provides samples that consistently fall within the true support and achieves better alignment with the target density, especially at smaller step counts.
grows. The score-based sampler uses the canonical Fuler-Maruyama
(a) 5 steps
1.40
(b) steps
(c) 100 steps
Figure 8: Samples from score-based and proximal samplers using 5, 10, and 100 sampling steps for distributions with varying level of smoothness at the boundary (controlled by a).
40

---

H.2 Additional Results
Score SDE
10a
Scure SDE (Spec.)
Score ODE
FID
Scure ODE (Spec.)
Prox Hybrid
Prex Hybrid (Spec.)
10
Prox Hybrid (Spec., No heur.)
10a
10a
10
NFE
Figure 9: FID vs. number of sampling steps (or number of function evaluations, NFE) on CIFAR1O. Dashed lines indicate models trained specifically for 10 and 20 steps ("Spec."). while solid lines represent full-range models. The result labeled "Prox Hybrid (Spec., No Heur.)" is obtained without applying the heuristies for network parameterization and objective balancing described Section 3.3.
0.8
0.7
0.5
0.5
Fresison
0.4
0.4
Score SDE
Scoze SDM
Score ODE
Scoce ODE
0.1
Proor Backward (Ours)
0.1
Prox Backward (Ours)
0.0
Prox Hybrid (Ours)
Prox Hybrid (Ours)
10
10
10
10
10a
10
NFE
NFE
(a) MNIST Precision
(b) MNIST Recall
Score SDE
Score ODB
0.7
Prox Hybrid
0.6
0.5
Pretiion
0.4
0.4
0.3
Score SDE
0.1
0.1
Score ODE
0.0
Prox Hybrid
10
10
10
10a
10a
10a
NFE
NFE
(c) CIFAR-10 Precision
(d) CIFAR-10 10 Recall
Figure 10: Precision and recall metrics [38] for MNIST and CIFAR-10 10 datasets. Precision measures the sample quality while recall reflects the coverage of the data distribution. Across most sampling steps (NFE), ProxDM achieves higher precision and recall than score-based baselines, indicating improved sample quality and diversity.
41

---

Scure SDE
Score ODE
ProXDM
10a
FID
10a
10a
10
NFE
Figure 11: FID vs. number of sampling steps (NFE) on CelebA-HQ (256 x 256) from score SDE,
score ODE and hybrid ProxDM samplers. FID is computed 30,000 generated samples.
42

---

Score SDE
Socce ODE
ProdD
100 steps
Scoce SDE
Sooce ODE
ProdDM
(b) 20 steps
Sizre SIE
Souce OIB
ProDDM
(c) 10 steps
Figure 12: Uncurated CelebA-HQ (256 x 256) samples from score SDE, score ODE and hybrid ProxDM samplers. ProxDM produces cleaner and more coherent samples with fewer sampling steps.
43

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
The abstract and introduction clearly state the scope of the paper, its contribu- tions, and results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We discuss limitations in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
44

---

Justification:
The full set of assumptions and proof are provided in appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide details on network architecture, loss function, training setups in Sections and 5, as well as in appendix.
Guidelines:
The answer NA means that the paper does not include experiments. If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
45

---

Answer:
[Yes]
Justification:
All data and models used in this paper are public. We make our code available.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
public/guides Sasiidde//odeSSumisssssse...s cessooPolll oolicyliyy for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //niss.. wtt.cc/ppbb. wwtpccccpp/paaecccccccccccccceet Seliiiigguidee/CCooe htt////ooSuuiisseeeeeee iinniissioonn onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Yes, we provide all the details necessary to understand the results presented in the main text. We include further details in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Experiments require significant amount of GPU time to run, so it is quite costly to compute error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
46

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Yes, we specify the hardware details used to rum experiments in appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: /ttp:///eeuueeeeeeeeeeeee http://wuurrsscccppiiiiiiee/e/ee httpe////////tttttttteeeeeeeeeee csGuidel ines?
Answer: [Yes]
Justification: Yes, our research complies with the Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We discuss societal impacts of our work in appendix.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no socictal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
47

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification:
We will implement safeguards when releasing our model.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
The sources of datasets are properly cited.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com /atasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
48

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
No new assets introduced.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
No crowdsourcing or research with human is conducted.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
49

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The method does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy (https /ttp:r/://.u...///eee /ttps/////oooeeeece//////////// for what should or should not be described.
50


Reviews:
[{"review_id": "wipliUoZAk", "reviewer": "Reviewer_uiu8", "summary": "Proximal Diffusion Models propose learning the proximal operator of the reverse‐time SDE in diffusion generative models, rather than estimating its score (gradient of the log probability). This operator directly produces the MAP denoised output under a Gaussian corruption model, simplifying each sampling step. The authors introduce a “proximal matching” objective to train this network and back it up with rigorous convergence guarantees, showing that Proximal Diffusion Models reaches high‐quality samples in fewer iterations than traditional score‐based methods.", "questions": "- Proximal matching trains a MAP denoiser under Gaussian corruption (i.e., effectively targeting the mode of the posterior) whereas score‐based diffusion methods optimize a variational bound that (at least loosely) increases the model’s log‐likelihood, this can be more explicit with MLE diffusion weighting. Do you have any theoretical guarantees that minimizing the proximal matching loss will also improve—or even converge to—the true data log‐likelihood? In other words, is there an ELBO‐style bound or other analysis showing that training with a MAP denoiser yields maximum‐likelihood consistency?\n- The paper reports faster convergence, but is that improvement limited to the sampling phase, or does it extend to faster convergence during training as well? \n- With current trends of one-pass models such as consistency models, distillation methods, shortcut models, and flow models, what are the immediate benefit of proximal diffusion models? Are the main benefits faster inference? If so, those aforementioned methods already \"straighten\" the gradient trajectory—whether along the denoising path or the velocity field, allowing for a single one pass inference rendering a proximal operator unnecessary.\n- Overall, the paper is well written and engaging, but unless you can demonstrate potential innovation for future diffusion works or significant benefits beyond sampling speed, which many recent methods already address, I remain a little skeptical, though I am open to further discussion.", "limitations": "Yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths\n- The paper is well-structured and easy to follow.\n- The derivations are comprehensive and, to the best of my knowledge, mathematically sound.\n- The proofs are rigorous, and the theoretical convergence guarantees are clearly presented.\n- The experiments compellingly demonstrate superior sample quality with fewer sampling steps.\n\nWeaknesses\n- Beyond sampling efficiency, no clear advantage over standard diffusion models.\n- Experiments rely solely on FID versus sampling steps, offering a limited view of performance.\n- On larger datasets (e.g., CIFAR-10), traditional diffusion at full timesteps still achieves better FID than the proximal approach.\n- Comparative analyses of training efficiency and convergence behavior are absent.", "was_revised": true, "final_justification": "Written in response.", "rebuttal": "- **On the Significance of Our Work:**\n  - **”Beyond sampling efficiency, no clear advantage over standard diffusion models.”**\n  - **”With current trends of one-pass models such as consistency models, distillation methods, shortcut models, and flow models, what are the immediate benefit of proximal diffusion models?...”**\n\n\n**Response:** Thank you for allowing us to elaborate on this critical point. All diffusion models (both vanilla and accelerated versions) rely on *forward discretization* of continuous processes. It is true that in the last ~5 years, a lot of research has gone into developing accelerated methods based on deterministic processes (probability flows), shortcuts, distillation models, and more. Our goal was not just to provide an approach to accelerate inference using the well-established score-based paradigm, but rather to develop a new kind of diffusion model based on novel discretization schemes (which have stronger theoretical properties). A key contribution of our work is to show, for the first time, how backward discretization schemes can lead to a new type of diffusion models (based on proximal operators instead of scores). In doing so, we show that sampling based on backward discretization leads (immediately, and somewhat surprisingly) to better sampling efficiency than vanilla SDE score models. Naturally, one can envision developing corresponding accelerated models based on Prox-schemes (on ODES, shortcut methods, and more). These are only conceivable given the formalization of ProxDM.\n\nFurthermore, our work can be seen as a complementary and orthogonal line of work to these recent advancements. Our approach does not rely on any distillation technique and is trained purely from data. Moreover, the idea of backward discretization can be extended to other processes, such as the “straighter” ones mentioned by the reviewer. We believe it opens up promising directions for future research. As reviewer ZawX writes, \"I cannot wait to see more adoption of this method,\" and reviewer ot8V, “This method in this paper could yield new, improved algorithms for solvers for diffusion models”.\n\nFinally, we believe the findings from ProxDM may lead to a novel understanding of current methods. For example, recent few-step generative models have found the benefit of MAP-promoting losses [1,2], and our findings may provide new insights into the effectiveness of these new losses and inspire new improvements in methodology.\n\nReferences:\n\n[1] Yang Song and Prafulla Dhariwal.  Improved techniques for training consistency models.  In International Conference on Learning Representations (ICLR), 2024.\n\n[2] Geng, Zhengyang, Mingyang Deng, Xingjian Bai, J. Zico Kolter, and Kaiming He. \"Mean flows for one-step generative modeling.\" \n\n\n\n- **”Experiments rely solely on FID versus sampling steps, offering a limited view of performance.”**\n\n**Response:** Thank the reviewer for raising this point. Although FID is the most commonly used metric for unconditional image generation, we agree that additional metrics can provide a better view of the performance - we compute Precision and Recall and present the results in the tables below.\n\n**MNIST:**\n\nPrecision:\n| Method \\ NFE   |    5 |   10 |   20 |   50 |   100 |   1000 |\n|----------------|------|------|------|------|-------|--------|\n| Score SDE      | 0.00 | 0.00 | 0.00 | 0.00 |  0.00 |   0.75 |\n| Score ODE      | 0.00 | 0.00 | 0.35 | 0.48 |  0.55 |   0.69 |\n| Prox Backward  |      |      | 0.77 | 0.76 |  0.76 |   0.76 |\n| Prox Hybrid    | 0.24 | 0.70 | 0.74 | 0.74 |  0.75 |   0.75 |\n\nRecall:\n| Method \\ NFE   |    5 |   10 |   20 |   50 |   100 |   1000 |\n|----------------|------|------|------|------|-------|--------|\n| Score SDE      | 0.00 | 0.00 | 0.00 | 0.00 |  0.00 |   0.80 |\n| Score ODE      | 0.00 | 0.00 | 0.19 | 0.40 |  0.49 |   0.67 |\n| Prox Backward  |      |      | 0.55 | 0.73 |  0.75 |   0.72 |\n| Prox Hybrid    | 0.01 | 0.53 | 0.69 | 0.75 |  0.76 |   0.74 |\n\n**CIFAR-10:**\n\nPrecision:\n| Method \\ NFE   |    5 |   10 |   20 |   50 |   100 |   1000 |\n|----------------|------|------|------|------|-------|--------|\n| Score SDE      | 0.00 | 0.00 | 0.00 | 0.01 |  0.14 |   0.54 |\n| Score ODE      | 0.00 | 0.02 | 0.38 | 0.55 |  0.57 |   0.57 |\n| Prox Hybrid    | 0.00 | 0.24 | 0.37 | 0.47 |  0.50 |   0.53 |\n\nRecall:\n| Method \\ NFE   |    5 |   10 |   20 |   50 |   100 |   1000 |\n|----------------|------|------|------|------|-------|--------|\n| Score SDE      | 0.00 | 0.16 | 0.26 | 0.37 |  0.57 |   0.74 |\n| Score ODE      | 0.12 | 0.43 | 0.58 | 0.67 |  0.69 |   0.70 |\n| Prox Hybrid    | 0.43 | 0.80 | 0.79 | 0.76 |  0.74 |   0.72 |\n\nOn MNIST, both variants of ProxDM consistently outperform score-based samplers in both precision and recall across most step counts from 5 to 100. \nOn CIFAR-10, ProxDM shows clear advantages over Score SDE across nearly all step numbers and metrics, except for NFE=1000. Compared to the Score ODE sampler, ProxDM provides notable improvements at small NFEs (5, 10 and 20), while maintaining comparable performance at higher step numbers.\n\n\n- **”On larger datasets (e.g., CIFAR-10), traditional diffusion at full timesteps still achieves better FID than the proximal approach.”**\n\n**Response:** This is correct, and we argue that this is not a significant limitation: 1) Practically speaking, using 1000 time steps is usually computationally expensive and impractical for many applications. Our method is designed to be more efficient while still achieving competitive results. 2) The difference between the FID in this regime does not lead to significant differences in the visual quality of the samples. In contrast, our method can generate more visually appealing samples with fewer steps. 3) Improving performance in the fine-grained discretization regime is not the primary focus of this work. Our emphasis is on high quality under low step budgets. 4) Additionally, although our algorithm requires a proximal operator at each step, the mapping learned by our U-Net is not explicitly constrained to be a proximal (but only asymptotically). In the future, we believe that improved network parameterization (such as that from [1]) and training strategies may close the gap at high step counts.\n\n[1] Fang et al. \"What's in a prior? learned proximal networks for inverse problems.\" ICLR 2024.\n\n\n- **On Connection to Maximum Likelihood Training: ”Do you have any theoretical guarantees that minimizing the proximal matching loss will also improve—or even converge to—the true data log‐likelihood?”**\n\n**Response:** Thank you for this insightful question.\n\n1) Since our model is mainly built upon the continuous-time SDE framework in [1], it does not naturally come with an MLE interpretation as originated from the discrete framework in DDPM [2]. It remains an open question whether such an equivalence exists between proximal matching and MLE training, and establishing such an equivalence can be non-trivial. For instance, in the derivation in [2], the KL term between Gaussian distributions reduces to an MSE between means—this is critical for linking denoising score matching to ELBO. In contrast, our proximal matching loss is not an MSE, and the transition kernel in ProxDM is non-Gaussian. Nonetheless, we believe this is a promising direction and plan to explore it in future work.\n\n2) Although our method does not have a likelihood-based guarantee, it is supported by rigorous theoretical guarantees from a sample distribution perspective. In Theorem 1, we showed that if true proximal operators are known, the sample distribution generated by our algorithm converges to the true data distribution in KL divergence. Furthermore, prior work [3] establishes theoretical guarantees for learning true proximal operators via proximal matching. Together, these results provide strong theoretical support for our approach.\n\n3) Additionally, as Gribonval [4] showed, MAP and MMSE denoisers can be closely related under a change of distributions. This suggests that our training paradigm may correspond to MLE training for a related distribution, which may be worth investigating in the future.\n\n\nReferences:\n\n\n[1] Song, et al. \"Score-based generative modeling through stochastic differential equations.\" ICLR 2021.\n\n[2] Ho, et al. \"Denoising diffusion probabilistic models.\" NeurIPS 2020.\n\n[3] Fang, et al. \"What's in a prior? learned proximal networks for inverse problems.\" ICLR 2024.\n\n[4] Gribonval R. “Should Penalized Least Squares Regression be Interpreted as Maximum A Posteriori Estimation?” IEEE Trans. Signal Process. 2011.\n\n\n\n- **On Training Efficiency:**\n  - **”The paper reports faster convergence, but is that improvement limited to the sampling phase, or does it extend to faster convergence during training as well?”**\n  - **”Comparative analyses of training efficiency and convergence behavior are absent.”**\n\n**Response:** We will be happy to provide training curves (e.g., FID vs. steps) in the supplementary material. Based on our experiments, we don't observe a clear advantage in training convergence speed compared to standard diffusion models (neither do we have reasons to believe this should be the case, unlike in sampling). This can be because proximal diffusion models are conditioned on two parameters - time and lambda, whereas score models are conditioned on time only. This higher dimension requires proximal diffusion models to learn more complex functions. In the future, better training strategies may help accelerate training, such as better loss weighting or sampling for time and lambda, as commented in Section 3.3."}, {"review_id": "qYRGQY3iLG", "reviewer": "Reviewer_ot8V", "summary": "This paper proposes a new method for improving the discretization for backwards processes for diffusion models. Inspired by proximal techniques in optimization, they derive a proximal method-based solver for the backwards process of the OU evolution, as well as a hybrid method, loosely based on forward-backwards methods in optimization theory. Their main theoretical result is a set of error bounds for these methods, given L2-accurate proximal operators. Under some reasonable assumptions on the distribution, they show that the proximal method achieves error eps with O(d / eps) steps, and the hybrid method requires O(d / sqrt(eps)).", "questions": "One feedback I have is that it would be nice to have a more complete comparison to the previous non-asymptotic rates for how many timesteps are necessary, given L2 accurate score functions. I believe the best rates are currently given by [1]; notably, their dependence on their eps is 1/eps^2, which maybe worth noting (although again these results are not quite comparable). \n\n[1] Nearly d-Linear Convergence Bounds for Diffusion Models via Stochastic Localization. Benton et al., ICLR 2024.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Quality: The technical content and ideas presented in the paper are very nice. The approach taken by the paper is very natural in retrospect given the success of proximal methods for optimization, and the bounds they obtain are relatively clean, and assuming learning accurate proximal estimators is about as difficult as learning score functions, achieve better error rates in terms of eps, although of course this is not directly comparable.\n\nClarity: The paper is well-written and very easy to read.\n\nSignificance: This method in this paper could yield new, improved algorithms for solvers for diffusion models, which are a very important area of study. \n\nOriginality: The idea of using such proximal methods, while clearly inspired by optimization, is very nice in this setting.", "was_revised": true, "final_justification": "The authors resolved my concerns to my satisfaction. I still believe the paper is quite strong and I will maintain my original score.", "rebuttal": "- **“One feedback I have is that it would be nice to have a more complete comparison to the previous non-asymptotic rates for how many timesteps are necessary, given L2 accurate score functions. I believe the best rates are currently given by [1]; notably, their dependence on their eps is 1/eps^2, which maybe worth noting (although again these results are not quite comparable).”**\n\n[1] Nearly d-Linear Convergence Bounds for Diffusion Models via Stochastic Localization. Benton et al., ICLR 2024.\n\n**Response:** Thank you for this comment. Yes, we are happy to add this comparison and a new paragraph on more comprehensive comparison to previous works. We agree with the reviewer that their results are not directly comparable to ours: they assumed non-smooth data distributions and L2-accurate score estimates, whereas we assumed L-smooth distributions and exact proximal operators. We are working on extending our analysis to learned proximal operators and non-smooth distributions."}, {"review_id": "DrphUBQI3r", "reviewer": "Reviewer_JaE7", "summary": "The paper proposes an alternative approach to training diffusion models - based on *proximal matching* instead of score matching. The key idea behind their approach is to use a *backward discretization* of the reverse stochastic process  for the sampling (in place of the standard forward discretization used by score matching methods).\n\nThey show that the backward discretization update can be interpreted as a proximal step w.r.t the log-density, and provide an algorithm \"ProxDM\" (with two variants for two kinds of discretizations) to generate samples using this proximal operator (analogous to the sampling algorithm used by score-based models like DDPMs).\n\nThey show how to train a neural network to learn this proximal step (\"proximal-matching\"). Their method involves sampling different times and regularizer strengths (t, $\\lambda$ pairs) in order to to handle multiple step sizes at sampling time, and they also propose appropriate parametrization and loss weighting to balance the training loss contribution of different times / regularizer strengths.\n\nThey also prove (under assumption that training algorithm converges to oracle proximal operator) that both variants of their ProxDM sampling algorithm converge to the true distribution, at a faster rate (fewer sampling steps) than score-matching methods.\n\nFinally, they provide experimental verification using toy datasets (MNIST and CIFAR10) that ProxDM generates better samples than score-matching methods for small number of steps.", "questions": "(nit) The text says that you include pseudo-code for proximal-matching training in the appendix but I couldn't find it, could you please add it (or else remove this statement)?\n\nMy main questions are about comparison with score-based methods:\n\n1. [Sampling of (t, $\\lambda$) pairs during training]: since the choice of $\\lambda$ depends on the number of sampling steps N that we wish to use during inference, ProxDM requires training the network with, say, 10x more ( t, $\\lambda$  ) pairs in order to support 10 different Ns compared to a single fixed N. Is this an inherent limitation of ProxDM compared to score-based methods like DDPM (which as I understand requires sampling only t during training, and then any number of sampling steps can be used at inference)?\n\n2. For score-matching, the training objective is actually equivalent to minimizing the variational bound on the negative log likelihood of the training data. Does the ProxDM training objective also have such an equivalent formulation?", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths:\n1. Excellent presentation, nice theory, very easy and enjoyable read especially for a mathematically dense paper.\n2. The idea of using proximal operators to do sampling in diffusion models is certainly novel, and the motivation behind it (faster convergence) is very compelling.\n3. The authors clearly explain the salient differences between their method and the score-based method, and also provide convenient mental models for comparing them (e.g. interpretation as MAP denoiser vs MMSE denoiser).\n4. The paper does a great job closing the loop from theory to practically useful methods (evidenced by their detailed analysis/experiments which revealed that the hybrid discretization method is better in practice even though the fully backward discretization was better in theory). They didn't shy away from discussing practical considerations (sampling of t and lambda, loss weighting) and limitations (minimum steps requirement for the fully backward method).\n\nWeaknesses:\nMore experimental results could have been provided (e.g. using large image datasets like CelebA) to better justify the practical utility of these methods for \"very high quality\" image generation. In this regime (large image + high quality), I would have liked to know if either of the two ProxDM variants give a significant reduction in number of steps compared to score-based methods. It would also be interesting to see if the theoretical minimum on number of steps for the fully backward discretization is even a concern in this regime.", "was_revised": true, "final_justification": "The authors have done an excellent job with this paper. They also acknowledges my concerns satisfactorily and therefore I have kept my \"Accept\" rating.", "rebuttal": "- **”More experimental results could have been provided (e.g. using large image datasets like CelebA) to better justify the practical utility of these methods for \"very high quality\" image generation. In this regime (large image + high quality), I would have liked to know if either of the two ProxDM variants give a significant reduction in number of steps compared to score-based methods. It would also be interesting to see if the theoretical minimum on number of steps for the fully backward discretization is even a concern in this regime.”**\n\n**Response:** Thank you for this suggestion. Yes, the empirical component of our work could have been more comprehensive (we focused mainly on the foundational components of the work, and experimentally we only got to perform experiments on CIFAR). We are implementing a latent space version of our model for generation of high-dimensional images, as suggested.\n\n\n- **”The text says that you include pseudo-code for proximal-matching training in the appendix but I couldn't find it, could you please add it (or else remove this statement)?“**\n\n**Response:** Certainly. We're sorry about that. We will certainly add it.\n\n\n- **Question 1: [Sampling of (t, lambda) pairs during training]: since the choice of depends on the number of sampling steps N that we wish to use during inference, ProxDM requires training the network with, say, 10x more (t, lambda) pairs in order to support 10 different Ns compared to a single fixed N. Is this an inherent limitation of ProxDM compared to score-based methods like DDPM (which as I understand requires sampling only t during training, and then any number of sampling steps can be used at inference)?”**\n\n**Response:** The reviewer is correct. ProxDM requires training with different sets of (t, lambda) pairs to support different sampling steps, because the proximal operator invoked at each step depends on the step size. Score models natively support arbitrary sampling steps, as the score function invoked at each step is independent of step size. We plan to explore continuous sampling strategies for (t, lambda) in the future to train any-step ProxDM.\n\n\n- **Question 2: “For score-matching, the training objective is actually equivalent to minimizing the variational bound on the negative log likelihood of the training data. Does the ProxDM training objective also have such an equivalent formulation?”**\n\n**Response:**\nThank you for this insightful question.\n\n1) Since our model is mainly built upon the continuous-time SDE framework in [1], it does not naturally come with an MLE interpretation as originated from the discrete framework in DDPM [2]. It remains an open question whether such an equivalence exists between proximal matching and MLE training, and establishing such an equivalence can be non-trivial. For instance, in the derivation in [2], the KL term between Gaussian distributions reduces to an MSE between means—this is critical for linking denoising score matching to ELBO. In contrast, our proximal matching loss is not an MSE, and the transition kernel in ProxDM is non-Gaussian. Nonetheless, we believe this is a promising direction and plan to explore it in future work.\n\n2) Although our method does not have a likelihood-based guarantee, it is supported by rigorous theoretical guarantees from a sample distribution perspective. In Theorem 1, we showed that if true proximal operators are known, the sample distribution generated by our algorithm converges to the true data distribution in KL divergence. Furthermore, prior work [3] establishes theoretical guarantees for learning true proximal operators via proximal matching. Together, these results provide strong theoretical support for our approach.\n\n3) Additionally, as Gribonval [4] showed, MAP and MMSE denoisers can be closely related under a change of distributions. This suggests that our training paradigm may correspond to MLE training for a related distribution, which may be worth investigating in the future.\n\n\nReferences:\n\n\n[1] Song, et al. \"Score-based generative modeling through stochastic differential equations.\" ICLR 2021.\n\n[2] Ho, et al. \"Denoising diffusion probabilistic models.\" NeurIPS 2020.\n\n[3] Fang, et al. \"What's in a prior? learned proximal networks for inverse problems.\" ICLR 2024.\n\n[4] Gribonval R. “Should Penalized Least Squares Regression be Interpreted as Maximum A Posteriori Estimation?” IEEE Trans. Signal Process. 2011."}, {"review_id": "0YNYO987d0", "reviewer": "Reviewer_ZawX", "summary": "This paper proposes to use proximal optimization techniques to speed up sampling in diffusion models. Two variants of the algorithm is proposed, one uses full backward discretization, and one uses a hybrid of current and previous point. Motivated by the algorithm, the authors also proposed proximal loss functions to train the \"proximal score\", giving the data-dependant proximal operator used during sampling. The authors analyzed the convergence properties, discussed relations to the usual forward discretization, and verified performance on synthetic and real image datasets.", "questions": "1. Why can't you let $\\zeta$ go to zero? Is the schedule used in the CIFAR-10 experiments (which I appreciate the honesty) from previous work? How important is this parameter? \n\n2. Why does the synthetic dataset composed of very well separated points? What if you use a distribuion where the outline of the picture is more connected, i.e. almost on a low-d manifold? This would provide a case where the score field becomes low-d, and hence help show advantage to MMSE learning? \n\n3. Why can't the FID at high function evaluation beat the normal forward diffusion? Please discuss on this, as people still care about the quality. \n\n4. Why is the backward method missing in the CIFAR-10 experiment?", "limitations": "Limitations are well-discussed. Although addressing the performance saturation at high NFE would be helpful.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n1. The method is well-motivated by research in proximal-optimization, and is well-developed into two plausible versions. \n1. The sampling method then motivates a new loss function, leading to the a new proximal score matching technique that is novel to my knowledge. \n2. Writing of the paper is of exceptionally high quality. The main text is accessible to a wide audience, has the perfect amount of technicality, with connections to more conventional methods clearly discussed. I did not fully checked the proofs in Appendix (yet).\n3. The experimental results show promising performance at small number of function evaluations. I cannot wait to see more adoption of this method. \n4. I also appreciate that the authors spent enough passage discussing implementation details, including the crucial reparameterization of the model. I find these non-trivial details very useful and educational, highlighting the practical focus of the authors, in addition to their analytical rigor.\n\nWeaknesses:\nThere is hardly any weaknesses of large concerns (except the proofs in Appendix which I haven't checked). I point out a few minor suggestions.\n1. The authors discussed the benefits of the MAP estimate from proximal points vs the MMSE denoiser (which is provides a lot of nice insights). However, the experiments did not support the claims. I would have thought that the synthetic dataset has a sharp score field at low noise, and hence the authors would show that the good performance comes from using the MAP instead of MMSE. It would be great to use this experiment to further develop this point. \n2. I am a little confused by the definition of $f_\\theta(\\cdot;t,\\lambda)$, I think the reader can figure it out almost by themselves, but it would be nice to provide a quick derivation (e.g. in the extra page if accepted). \n3. Would be nice to provide some intuition on the loss of Fang et al., if there is any. I think this is a very cool reference to read, but for self-contained purposes do give wide audience some high-level ideas. \n4. It would be nice to visualize the score field of the proximal model. I think this would help us understand the differences and connections between the two frameworks, and give readers more visuals to help understanding and interpreting the method.\n\nIf I have time I will read more into the Appendix, and provide more feedback. I put my confidence low for now (sorry for the limited review). If the authors could address the points above and provide more results (the suggestions do not require any large-scale experiments on new datasets), I am willing to further increase the score.", "was_revised": true, "final_justification": "It's a good paper. Sound techniques and theory, and full of useful details. The clarifications from this discussions should improve the paper. The only downside that prevents a 6 from me is empirical performance being not super impressive. But other than that, this paper should get in.", "rebuttal": "- **”There is hardly any weaknesses of large concerns (except the proofs in Appendix which I haven't checked). I point out a few minor suggestions.”**\n\n**Response:** We are glad the reviewer enjoyed our work!\n\n\n- **”The authors discussed the benefits of the MAP estimate from proximal points vs the MMSE denoiser (which is provides a lot of nice insights). However, the experiments did not support the claims. I would have thought that the synthetic dataset has a sharp score field at low noise, and hence the authors would show that the good performance comes from using the MAP instead of MMSE. It would be great to use this experiment to further develop this point.”**\n\n**Response:** We appreciate the reviewer’s interest in the MAP vs. MMSE comparison. Could the reviewer please clarify what aspect of the experiment they found unconvincing? To clarify, as shown in Figure 1, our ProxDM method - based on MAP denoisers - outperforms the score samplers based on MMSE denoiser, in terms of coverage of the target distribution and Wasserstein distance. Please note that no training is involved in the synthetic experiment - we use the exact score and proximal operators that were analytically derived from the ground-truth distribution.\n\nPrior work has shown that MAP denoisers produce more perceptually appealing outputs than MMSE denoisers [1], which conceptually supports our use of proximal operators in generative models.\n\nWe are happy to further develop these experiments as the reviewer suggests.\n\n[1] Blau, Yochai, and Tomer Michaeli. \"The perception-distortion tradeoff.\" In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 6228-6237. 2018.\n\n\n\n- **”I am a little confused by the definition of f_theta(;t, lambda), I think the reader can figure it out almost by themselves, but it would be nice to provide a quick derivation (e.g. in the extra page if accepted).”**\n\n**Response:** Certainly, we are happy to add more clarification on this. We use $f_\\theta( ; t, \\lambda)$ to denote the learned proximal operator at time $t$ and with regularization parameter $\\lambda$, i.e., $prox_{-\\lambda \\ln p_t}$. In implementation, we parameterize $f$ by $f_\\theta(x; t, \\lambda) = x - \\epsilon_\\theta(x; t, \\lambda)$, where $\\epsilon_\\theta$ is a neural network. We found this leads to better convergence of our model at large step numbers in early experiments. This parameterization is inspired by a similar $\\epsilon$ parameterization in DDPM (please see equation 11 in [1]).\n\n[1] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. \"Denoising diffusion probabilistic models.\" NeurIPS 2020.\n\n\n- **”Would be nice to provide some intuition on the loss of Fang et al., if there is any. I think this is a very cool reference to read, but for self-contained purposes do give wide audience some high-level ideas.”**\n\n**Response:** Certainly, we can add more intuition in the text, with a figure illustrating the loss function. Useful intuition comes from trying to recover the mode of a distribution in 1D. In that case, an L0 pseudo-norm (if one could minimize such a non-differentiable loss function), recovers the maximum of a distribution. The proximal matching loss is a generalization of this idea to R^d and differentiable functions.\n\n\n- **“It would be nice to visualize the score field of the proximal model. I think this would help us understand the differences and connections between the two frameworks, and give readers more visuals to help understanding and interpreting the method.”**\n\n\n**Response:** Thank the reviewer for this suggestion. Can we ask what the reviewer meant by the “score field of the proximal model”? If it refers to the displacement field of prox (i.e., prox(x) - x), we are happy to add such a comparison in 2D space. However, please note that we cannot include images in the rebuttal.\n\n\n\n- **”Why can't you let \\zeta go to zero? Is the schedule used in the CIFAR-10 experiments (which I appreciate the honesty) from previous work? How important is this parameter?”**\n\n**Response:** We thank the reviewer for this question. $\\zeta$ is the parameter in the proximal matching loss, and taking it to zero is the correct approach from a theoretical standpoint. In practice, however, when $\\zeta$ equals zero, the loss can blow up to infinity. In experiments, we use a schedule to set $\\zeta$, which can be tuned for different datasets and tasks. Previous work does not provide a schedule for image generation on CIFAR-10.  In our early experiments, we tuned hyperparameters for $\\zeta$, e.g., extreme values and decay rate. Although theoretically speaking, a smaller $\\zeta$ brings the learned mapping closer to the true proximal operator, empirically, a smaller $\\zeta$ does not always lead to improved generation performance. In general, we found that $\\zeta$ between 0.5 and 1 leads to similar performance in the tasks we considered. However, we also observed that the optimal $\\zeta$ could vary based on the dataset and even the number of sampling steps—a more detailed study of the relation between $\\zeta$ and generation performance is an interesting direction for future work.\n\n\n- **”Why does the synthetic dataset composed of very well separated points? What if you use a distribution where the outline of the picture is more connected, i.e. almost on a low-d manifold? This would provide a case where the score field becomes low-d, and hence help show advantage to MMSE learning?”**\n\n**Response:** Thank the reviewer for this suggestion. The synthetic dataset we use comes directly from the DataSauraus dataset [1]—we did not perform any explicit subsampling. We are working on adding an experiment on a low-dimensional manifold comparing MAP and MMSE learning.\n\n[1] Matejka, Justin, and George Fitzmaurice. \"Same stats, different graphs: generating datasets with varied appearance and identical statistics through simulated annealing.\" In Proceedings of the 2017 CHI conference on human factors in computing systems, pp. 1290-1294. 2017.\n\n\n\n- **”Why can't the FID at high function evaluation beat the normal forward diffusion? Please discuss on this, as people still care about the quality.”**\n\n**Response:** Thank the reviewer for pointing this out. We argue that this is not a significant limitation: 1) Practically speaking, using 1000 time steps is usually computationally expensive and impractical for many applications. Our method is designed to be more efficient while still achieving competitive results. 2) The difference between the FID in this regime does not lead to significant differences in the visual quality of the samples. In contrast, our method can generate more visually appealing samples with fewer steps. 3) Improving performance in the fine-grained discretization regime is not the primary focus of this work. Our emphasis is on high quality under low step budgets. 4) Additionally, although our algorithm requires a proximal operator at each step, the mapping learned by our U-Net is not explicitly constrained to be a proximal (but only asymptotically). In the future, we believe that improved network parameterization (such as that from [1]) and training strategies may close the gap at high step counts.\n\n[1] Fang et al. \"What's in a prior? learned proximal networks for inverse problems.\" ICLR 2024.\n\n\n\n\n- **”Why is the backward method missing in the CIFAR-10 experiment?”**\n\n\n**Response:** Thank the reviewer for raising this point. The backward method has a practical limitation on the maximum step size it can take. As shown in Eq. (PDA) and Algorithm 1, each backward update involves the proximal operator $prox_{-\\lambda \\ln p_{ t_{k-1} } }$, where the regularization weight is $\\lambda = \\frac{2\\gamma_k}{2 - \\gamma_k}$. To ensure $\\lambda > 0$, we require $\\gamma_k < 2$, and because $\\gamma_k$ increases with step size, this imposes a hard upper bound on the step size. In practice, as the step size increases, $\\gamma_k$ approaches 2, and the update becomes increasingly unstable, which degrades sampling quality at few steps. On MNIST, we found performance degrades noticeably at around 20 steps. On CIFAR-10, we observed the instability issue become more pronounced, which led to poor performance at few steps. For this reason, we chose not to include the backward method in the CIFAR-10 experiment. That said, we plan to explore ways to address this limitation in future work."}]

Author rebuttal:

Meta-review: The inference of diffusion model is typically enabled by a time discretization of its backward process (such as a backward SDE), and this paper considers discretizing the backward SDE using backward Euler numerical scheme instead of the commonly used explicit schemes. This is equivalent to a proximal map formulation. However, backward Euler, being an implicit scheme, is computationally expensive, and the key idea of this paper is to learn the proximal map instead of the score function, so that the necessity of implicit equation solves needed by backward Euler can be circumvented. The efficacy of the proposed method is then demonstrated both theoretically and empirically. Both reviewers and I were impressed and convinced by the clarity and the extent of demonstration. While some reviewers also pointed out (and I agree) that a comparison to more recent theoretical results, such as [1-3], should be discussed, all of us agree that this work has enough merits for an acceptance, which I'm delighted to recommend. I encourage the authors to take the discussions into consideration when preparing for a revised/final version.

[1] Nearly d-linear convergence bounds for diffusion models via stochastic localization
[2] KL convergence guarantees for score diffusion models under minimal data assumptions
[3] Evaluating the design space of diffusion-based generative models
(note these are not based on advanced discretization, unlike the compared `accelerated strategy')
Decision: Accept (poster)