Balancing Act: Distribution-Guided Debiasing in Diffusion Models

TL;DR

This work tackles demographic bias in diffusion models by introducing Distribution Guidance, a non-retraining debiasing approach that steers generated attribute distributions toward a user-specified reference $\mathbf{p^a_{ref}}$. Central to the method is an Attribute Distribution Predictor (ADP) that operates in the diffusion model's h-space, predicting the batch-level attribute distribution $\hat{\mathbf{p}^a_{\theta}}$ and guiding the reverse process to minimize $\mathcal{L}(\hat{\mathbf{p}^a_{\theta}}, \mathbf{p^a_{ref}})$ through gradient updates on $\mathbf{h_t}$. The authors demonstrate strong, quantitative debiasing across single- and multi-attribute scenarios for both unconditional DMs and text-to-image models like Stable Diffusion, while maintaining image quality as measured by FD and FID. They also show practical downstream benefits, including improved class-balanced data for training attribute classifiers. Overall, the approach provides a data-efficient, scalable path to fair generation in large diffusion models without costly retraining.

Abstract

Diffusion Models (DMs) have emerged as powerful generative models with unprecedented image generation capability. These models are widely used for data augmentation and creative applications. However, DMs reflect the biases present in the training datasets. This is especially concerning in the context of faces, where the DM prefers one demographic subgroup vs others (eg. female vs male). In this work, we present a method for debiasing DMs without relying on additional data or model retraining. Specifically, we propose Distribution Guidance, which enforces the generated images to follow the prescribed attribute distribution. To realize this, we build on the key insight that the latent features of denoising UNet hold rich demographic semantics, and the same can be leveraged to guide debiased generation. We train Attribute Distribution Predictor (ADP) - a small mlp that maps the latent features to the distribution of attributes. ADP is trained with pseudo labels generated from existing attribute classifiers. The proposed Distribution Guidance with ADP enables us to do fair generation. Our method reduces bias across single/multiple attributes and outperforms the baseline by a significant margin for unconditional and text-conditional diffusion models. Further, we present a downstream task of training a fair attribute classifier by rebalancing the training set with our generated data.

Figures (12)

• Figure 1: a) Random sampling from a pretrained DM ldm generates images with gender imbalance. b) Proposed method takes a user-defined reference attribute $\mathbf{p^a_{ref}}$ and performs distribution guidance on a pretrained DM. c) Sampling with distribution guidance results in fair generation that follow user define $\mathbf{p^a_{ref}}$.
• Figure 2: Distribution of samples w.r.t. decision boundary
• Figure 3: Sample guidance vs Distribution guidance. a) After a few steps of denoising ($t=\tau$), the generated samples have learned some discriminative features for gender. Sample guidance randomly selects samples from the batch uniformly from all the quantiles for conversion enforcing samples with dominant female features to also change. However, distribution guidance majorly converts the samples close to the decision boundary (Q1/Q2), which is easy to convert. b) Visualization of the generated samples. Samples transformed with distribution guidance are natural looking without any distortion, whereas images with sample guidance suffer from distortion or unnatural appearance.
• Figure 4: Distribution guidance in the h-space. For a given batch $\mathbf{x_t^{[1:N]}}$, we extract the intermediate h-space representation $\mathbf{h_t^{[1:N]}}$ and pass it through ADP to obtain attribute distribution $\mathbf{\hat{p^a_\theta}}$. Guidance updates $\mathbf{h_t^{[1:N]}}$ by backpropagating the derivative of loss.
• Figure 5: Classification accuracy for linear h-space classifiers and ResNet-18 image space classifiers trained on $2K$ training examples. h-space classifiers are data efficient and achieve superior performance even with a linear layer