Constrained Diffusion Models via Dual Training

Summary

In this paper, the authors introduce a novel approach termed `Dual Training` for training constrained diffusion models, particularly focusing on scenarios involving `biased data generation`. Initially, the authors adeptly derive the learning objectives for diffusion models within a constrained optimization framework. Subsequently, they develop corresponding learning methodologies based on the principles of Lagrangian duality and propose an innovative training procedure. The effectiveness of their approach is rigorously validated through experiments on two distinct generative tasks.

Strengths

1. **Robust Derivation Procedure**: The authors present a thorough and well-founded derivation procedure throughout the manuscript, which stands out as particularly compelling. This rigorous approach significantly strengthens the theoretical foundation of their work. 2. **Relevant Application Scenario**: The authors explore a critical yet underrepresented field in their research. This focus not only highlights the relevance of their study but also underscores its potential impact in areas that have previously received limited attention.

Weaknesses

**Major Issues:** 1. The authors state, "Compared with the loss re-weighting method [9], our constrained formulation provides an optimal trade-off between matching given data distribution and following reference distribution." It would be helpful if the authors could clarify whether the re-weighting-based approach can be considered a special case of the proposed method in this manuscript. 2. The discussion on diffusion models appears to focus predominantly on the DDPM model. Could the authors explore whether this proposed approach can be extended to other diffusion processes, such as the Variance Exploding Stochastic Diffusion? 3. There seems to be an unclear transition from (U-KL) to (U-LOSS). Can the authors directly justify this conversion? It would be beneficial to include a derivation or a more detailed explanation, especially how it relates to equations (6) and (U-MIX). 4. The manuscript lacks comparative results with baseline models. Given that the authors suggest this approach provides a more general framework for generative models, could comparisons be made with other methods mentioned in reference [1]? 5. Could the authors discuss whether this approach can be interpreted from the perspective of partial optimal transport? The constraints discussed appear to share similarities with those in unbalanced optimal transport approaches, such as those outlined in equation 3 of reference [2]. 6. The paper discusses a dual training procedure involving primal-dual optimization. In the context of augmented direction multiplier methods used in directed acyclic graph learning, batch size is known to significantly influence model efficacy. Could the authors provide a sensitivity analysis regarding batch size effects? **Minor Issue:** 1. Under Assumption 1, should the range of $\zeta \in (0, b_i)$, considering that KL divergence cannot be negative? 2. Regarding the term (U-LOSS), should the subject be revised to $\nabla{\log{q^i(x_t)}}$? Additionally, it would be helpful if the authors could provide a detailed explanation or derivation of how the constraint term is converted into the (U-LOSS) term. --- **References** [1] Choi, Kristy, et al. "Fair generative modeling via weak supervision." International Conference on Machine Learning. PMLR, 2020. [2] Duque, Andrés F., Guy Wolf, and Kevin R. Moon. "Diffusion transport alignment." International Symposium on Intelligent Data Analysis. Cham: Springer Nature Switzerland, 2023.

Questions

Please see weaknesses.

Rating

5

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

1. **Organization:** The manuscript lacks a dedicated section for related work, which is essential for contextualizing the study within the existing literature. The absence of this section diminishes the foundational motivation of the manuscript, potentially limiting the reader's understanding of its contributions relative to prior research. 2. **Baseline Comparison:** The authors did not include experiments comparing their methods against established baselines. Such comparisons are crucial for demonstrating the efficacy and advancements of the proposed approach over existing techniques. 3. **Hyper-parameter Sensitivity Analysis:** The manuscript does not address hyper-parameter sensitivity. Including such analysis is important to evaluate the robustness and reliability of the model across various parameter settings.

Summary

This paper proposes Dual Training, a method designed to constrain the distributions that denoising diffusion models can learn. The authors propose an extension to the standard diffusion model training objective that minimizes the Kullback-Leibler (KL) divergence of the learned distribution with respect to two components: (1) the original data distribution and (2) a set of auxiliary distributions representing relevant constraints. The authors derive a tractable algorithmic approximation of their approach and apply it to two scenarios: the fair generation of underrepresented image classes and the fine-tuning of pre-trained models on new data while preserving image classes from the pre-training dataset.

Strengths

* The paper proposes an original approach to addressing two critical challenges in diffusion-based image generation: fairness and the ability to integrate new data with pre-trained diffusion models. * The method is well-motivated and the theoretical analysis of both the constrained optimization problem and the proposed solution embeds it in a rigorous mathematical framework. * Experimental results demonstrate the efficacy of the proposed approach on two relevant tasks (fair generation and fine-tuning), showcasing quantitative and qualitative improvements over unconstrained and pre-trained baselines.

Weaknesses

* The experimental evaluation in the main text is limited, focusing primarily on the MNIST and CelebA datasets. While some results on CIFAR-10 are provided in the appendix, they are not referenced in the main text. A more comprehensive empirical evaluation using more challenging datasets would strengthen the experimental section and better demonstrate the method's utility. * The empirical evaluation lacks relevant baselines. For instance, the fine-tuning experiments do not provide quantitative results for models fine-tuned without constraints. More generally, it would be beneficial to see how the model compares with standard conditioning approaches capable of enforcing similar constraints, for example, training a model with classifier-free guidance on datasets with underrepresented classes and providing balanced conditioning information at inference time. * The empirical results do not include uncertainty estimates or measures of statistical significance, which would enhance the robustness of the findings.

Questions

* Can you provide more details on the computational overhead of dual training compared to standard diffusion model training? * How does the method compare to other approaches that explicitly aim to fit mixture distributions, e.g. [1]? * How sensitive are the results to the choice of hyperparameters (number of dual iterations, primal/dual batch sizes, primal/dual learning rates, etc.)? Is there a principled way to set these values? --- [1] Du, Yilun, et al. "Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC." International conference on machine learning. PMLR, 2023.

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

3

Limitations

The authors discuss some limitations, including the need to experiment with more datasets/attributes and explore converting other types of constraints into the KL divergence formulation. It would be good to also discuss and compare the computational overhead of the proposed formulation as this would be beneficial to evaluate the method's practical implications and trade-offs

Summary

This paper studies the constrained diffusion models, motivated by customizing the generation by specific tasks. The idea is to formulate KL divergence-constrained optimization problem (U-KL), which is shown to have zero duality gap. The convergence (rate) of the sampling process is given by assuming a mixture initial distribution. Several experiments have been conducted, e.g., on MNIST and on the fairness.

Strengths

The paper provides a thorough study on the KL-constrained diffusion models: both optimization and sampling aspects are considered. The paper is overall interesting (methodologically). I also checked most proofs, and they are correct.

Weaknesses

There are several weaknesses (and questions): (1) The problem (U-KL) is still abstract. It is not clear how to choose $b_i$ to represent desired properties, or what are the right interpretations of the KL constraints (e.g., Section 3.2 on the fairness). (2) Theorem 7: the authors may emphasize that Theorem 7 is valid under the assumption that the initial distribution is a mixture. Also I think DDPM is used. The authors may also compare Theorem 7 with the results in [28]. (3) Experiments: the authors consider MNIST, CIFAR 10 and Celeb-A. However, the FID obtained seems to be large compared to the existing results. For instance, for MNIST the FID is typically less than 1 while the paper reports ~20-40; for CIFAR 10 the FID is less than 3 (the best is around 1.8) while the paper reports ~50. The authors may explain why this happens, otherwise the experiment results are not so convincing. (4) Literature and comments: There have also been a line of work on the SDE-based diffusion models (see ref. a), which show good empirical results. Related to (3), ref b. reports the FID 1.8 for CIFAR 10, and ref c. reports the FID 0.8 for MNIST (both without constraints). The authors may refer to these results on the continuous-time (SDE-based) models, and explain why there is so large difference in the empirical results. Another suggestion is that the authors may consider fine-tuning in the continuous framework, as in [45] and ref d. a. Score-Based Generative Modeling through Stochastic Differential Equations, Song et al., arXiv:2011.13456. b. Elucidating the Design Space of Diffusion-Based Generative Models, Karras, Aittala, Aila and Laine, arXiv:2206.00364. c. Contractive Diffusion Probabilistic Models, Tang and Zhao, arXiv:2401.13115. d. Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond, Tang, arXiv:2403.06279. I would be happy to raise the score if some (or all) of the above concerns are addressed.

Questions

See weakness

Rating

5

Confidence

5

Soundness

3

Presentation

3

Contribution

3

Limitations

NA

Summary

This paper aims to address the issue of generating biased data based on the training dataset for diffusion models. The authors introduce a constrained diffusion model by imposing the diffusion constraints based on desired diffusions that are informed by requirements/constraints. They propose a dual training algorithm to train the model, and characterize the convergence of the trained model. Two constrained generation tasks are explored: ensuring fairness to underrepresented classes and adapting a pretrained model to new data.

Strengths

- The paper is well written and easy to follow. - Its key idea is using the quadratic loss formulation of score matching, which is equivalent to the ELBO of diffusion models. The authors employ Lagrangian duality to update both the parameters of the diffusion model and dual variables of constraints. While Lagrangian duality is commonly used for constrained optimization, particularly for constrained generative models, see [1,2,3], this paper first introduces this approach for diffusion models. - Thanks to the quadratic loss minimization of diffusion models, the authors can establish the convergence of the proposed constrained models. I think this is the most novel part of the paper. [1] Liu et al., Sampling with trustworthy constraints: a variational gradient framework , NeurIPS 2021. [2] Danilo et al., Generalized ELBO with Constrained Optimization, GECO, Bayesian Deep Learning (NeurIPS 2018). [3] Ferdinando et al., Lagrangian Duality for Constrained Deep Learning, ECMLPKDD 2020.

Weaknesses

- This paper shares a close connection to [1] regarding its motivation: formulating fairness as a constrained distributional optimization. Both of them employs a methodology based on primal-dual optimization derived from Lagrangian duality to address the constraints. While [1] focuses on the constrained SVGD, this paper focuses on constrained diffusion models. - Some minor corrections: +) Algorithm 1, line 4: x_{\theta}(h) should be s_{\theta}(h) +) Section 5, line 314, q(x_0:T) should be q_{i}(x_0:T) (constraint part)

Questions

- How efficient the model is as the number of constraints increases? Each time the model updates the duality variables, it needs to train a diffusion model, which is slow. - how the proposed method differs from existing constrained generative models, for example [1].

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The paper can be improved by: - a discussion of the efficiency of the proposed constrained diffusion models should be included. - a discussion of how the proposed method differs from existing constrained generative models should be added.

Response to Rebuttal

I thank the authors for the detailed response. I have raised my score in response to the additional empirical results and clarifications and will outline my remaining concerns below. --- **Re Weakness 1**: ... more challenging datasets ... Thank you for providing these additional results. I believe that they are helpful and strengthen the empirical evaluation presented in the paper. **Re Weakness 2**: ... quantitative results for models fine-tuned without constraints ... baselines ... > We refer to Figure (2b) in the paper which in fact provides a quantitative baseline for models fine-tuned without constraints, both as generated samples and its FID score. Thank you for pointing out the FID scores in the figure captions. They do indeed provide the quantitative results I was referring to. Could you briefly touch on why they are so much higher than the results reported in other works (e.g. 3.17 in [1] for unconditional image generation on CIFAR-10)? > Regarding the second point, we believe there is no meaningful metric of comparing a constrained unconditional model to existing conditional models since the conditional model can sample from any class if it is conditioned to do so. The suggested approach of using a conditional model with balanced conditioning information at inference time is interesting. However, we were unable to find any existing baselines using the suggested approach. Hence, we believe implementing such a novel approach to use as a baseline is beyond the scope of this paper. I am not sure I understand this point. Conditioning diffusion models via classifier-free guidance [2] is a standard approach in image diffusion models, including the latent diffusion model [3] that I assume was used for the additional ImageNet results. Since this approach enables the specification of arbitrary class labels at inference time, it should be straightforward to sample them uniformly before sample generation. Since all class labels need to be known at training time to specify the KL divergence-constrained optimization problem (U-KL), it seems reasonable to compare the method to other approaches that use class labels to train conditional diffusion models that allow us to flexibly constrain the sample generation process. The reasoning behind this question is to better understand the potential practical impact the proposed approach, compared to simply adapting existing diffusion model techniques to the experimental settings investigated in the manuscript. **Re Questions 1-3**: I appreciate the detailed response and the primal/dual batch size sensitivity study. In addition to the conceptual considerations provided in response to Question 1 of Reviewer EijC, I think it would be helpful to add an actual empirical comparison of the computational cost of standard and dual training approaches in the paper. --- ### References [1] Ho, Jonathan, Ajay Jain, and Pieter Abbeel. "Denoising diffusion probabilistic models." Advances in neural information processing systems 33 (2020): 6840-6851. [2] Ho, Jonathan, and Tim Salimans. "Classifier-Free Diffusion Guidance." NeurIPS 2021 Workshop on Deep Generative Models and Downstream Applications. 2021. [3] Rombach, Robin, et al. "High-resolution image synthesis with latent diffusion models." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2022.