Constrained Diffusion Models via Dual Training

TL;DR

This work casts the training of diffusion models under requirements as a constrained distribution optimization problem that aims to reduce the distribution difference between original and generated data while obeying constraints on the distribution of generated data.

Abstract

Diffusion models have attained prominence for their ability to synthesize a probability distribution for a given dataset via a diffusion process, enabling the generation of new data points with high fidelity. However, diffusion processes are prone to generating samples that reflect biases in a training dataset. To address this issue, we develop constrained diffusion models by imposing diffusion constraints based on desired distributions that are informed by requirements. Specifically, we cast the training of diffusion models under requirements as a constrained distribution optimization problem that aims to reduce the distribution difference between original and generated data while obeying constraints on the distribution of generated data. We show that our constrained diffusion models generate new data from a mixture data distribution that achieves the optimal trade-off among objective and constraints. To train constrained diffusion models, we develop a dual training algorithm and characterize the optimality of the trained constrained diffusion model. We empirically demonstrate the effectiveness of our constrained models in two constrained generation tasks: (i) we consider a dataset with one or more underrepresented classes where we train the model with constraints to ensure fairly sampling from all classes during inference; (ii) we fine-tune a pre-trained diffusion model to sample from a new dataset while avoiding overfitting.

Figures (4)

• Figure 1: Generation performance comparison of constrained and unconstrained models that are trained on MNIST with three minorities: 4, 5, 7. ( Left ) Frequencies of ten digits that are generated by an unconstrained model () and our constrained model (); ( Middle ) Generated digits from unconstrained model ( FID 15.9 ); ( Right ) Generated digits from our constrained model ( FID 13.4 ).
• Figure 2: Generation performance comparison of constrained and unconstrained models that are trained on Celeb-A with male minority. ( Left ) Frequencies of two genders that are generated by an unconstrained model () and our constrained model (); ( Middle ) Generated faces from unconstrained model ( FID 19.6 ); ( Right ) Generated faces from our constrained model ( FID 11.6 ).
• Figure 3: Generation performance comparison of constrained and unconstrained models that are trained on Image-Net with minority classes: 'Cassette player' (2), 'French horn' (5), and 'Golf ball' (8). ( Left ) Frequencies of ten classes that are generated by an unconstrained model () and our constrained model (); ( Middle ) Generated images from unconstrained model ( FID 36.0 ); (Right) Generated images from our constrained model ( FID 27.3 ).
• Figure 4: Fine-tuning performance comparison of constrained and unconstrained models that are trained on MNIST. ( Left ) Frequencies of ten digits that are generated by a pre-trained model without digit 9 () and our fine-tuned constrained model (); ( Middle ) Generated digits from unconstrained model ( FID 45.9 ); ( Right ) Generated digits from our constrained model ( FID 25.2 ).

Theorems & Definitions (36)

• Lemma 1: Equivalent formulations
• Lemma 2: Strong duality
• Theorem 1: Optimal constrained model
• Corollary 1
• Theorem 2
• Lemma 3: Convergence of diffusion model
• Lemma 4
• Theorem 3: Optimality of constrained diffusion model
• Theorem 4: Optimality of approximate constrained diffusion model
• proof