Beyond Scores: Proximal Diffusion Models

• ”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.”

Response: We are glad the reviewer enjoyed our work!

• ”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.”

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.

Prior 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.

We are happy to further develop these experiments as the reviewer suggests.

[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.

• ”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).”

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]).

[1] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." NeurIPS 2020.

• ”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.”

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.

• “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.”

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.

• ”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?”

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.

• ”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?”

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.

[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.

• ”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.”

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.

[1] Fang et al. "What's in a prior? learned proximal networks for inverse problems." ICLR 2024.

• ”Why is the backward method missing in the CIFAR-10 experiment?”

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.

• ”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.”

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.

• ”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)?“

Response: Certainly. We're sorry about that. We will certainly add it.

• 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)?”

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.

• 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?”

Response: Thank you for this insightful question.

1. 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.

2. 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.

3. 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.

References:

[1] Song, et al. "Score-based generative modeling through stochastic differential equations." ICLR 2021.

[2] Ho, et al. "Denoising diffusion probabilistic models." NeurIPS 2020.

[3] Fang, et al. "What's in a prior? learned proximal networks for inverse problems." ICLR 2024.

[4] Gribonval R. “Should Penalized Least Squares Regression be Interpreted as Maximum A Posteriori Estimation?” IEEE Trans. Signal Process. 2011.

Dear reviewer, If not already, please take a look at the authors' rebuttal, and discuss if necessary. Thanks, -AC

• “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).”

[1] Nearly d-Linear Convergence Bounds for Diffusion Models via Stochastic Localization. Benton et al., ICLR 2024.

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.

• On the Significance of Our Work:
  • ”Beyond sampling efficiency, no clear advantage over standard diffusion models.”
  • ”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?...”

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.

Furthermore, 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”.

Finally, 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.

References:

[1] Yang Song and Prafulla Dhariwal. Improved techniques for training consistency models. In International Conference on Learning Representations (ICLR), 2024.

[2] Geng, Zhengyang, Mingyang Deng, Xingjian Bai, J. Zico Kolter, and Kaiming He. "Mean flows for one-step generative modeling."

• ”Experiments rely solely on FID versus sampling steps, offering a limited view of performance.”

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.

MNIST:

Precision:

Recall:

CIFAR-10:

Precision:

Recall:

On MNIST, both variants of ProxDM consistently outperform score-based samplers in both precision and recall across most step counts from 5 to 100. On 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.

• ”On larger datasets (e.g., CIFAR-10), traditional diffusion at full timesteps still achieves better FID than the proximal approach.”

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.

[1] Fang et al. "What's in a prior? learned proximal networks for inverse problems." ICLR 2024.

• 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?”

Response: Thank you for this insightful question.

1. 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.

2. 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.

3. 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.

References:

[1] Song, et al. "Score-based generative modeling through stochastic differential equations." ICLR 2021.

[2] Ho, et al. "Denoising diffusion probabilistic models." NeurIPS 2020.

[3] Fang, et al. "What's in a prior? learned proximal networks for inverse problems." ICLR 2024.

[4] Gribonval R. “Should Penalized Least Squares Regression be Interpreted as Maximum A Posteriori Estimation?” IEEE Trans. Signal Process. 2011.

• On Training Efficiency:
  • ”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?”
  • ”Comparative analyses of training efficiency and convergence behavior are absent.”

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.

Dear reviewer, If not already, please take a look at the authors' rebuttal, and discuss if necessary. Thanks, -AC

Dear reviewer,

Thank you again for your constructive comments. Please kindly let us know if there are any remaining concerns or if our answers have addressed all your questions.

Thank you!