Inference-Time Attribute Distribution Alignment for Unconditional Diffusion

Hao Luan, See-Kiong Ng, Chun Kai Ling

Keywords: distribution alignment, fairness, diffusion

TL;DR: A training-free approach enabling pretrained diffusion models to generate samples with attributes aligning to flexible test-time target distributions.

Abstract: Controllable generation is crucial to real-world applications of diffusion models. However, many applications require control beyond individual samples: the distribution of semantic attributes (e.g., proportions of styles, objects, or demographics) over generated samples should match user-specified targets that may change at test time. We formalize this setting as the inference-time attribute distributional alignment problem for pretrained diffusion models. To address this, we cast the attribute distributional alignment as an optimal control problem over the reverse diffusion process, viewing the process as the rollout of a dynamical system and augmenting it with an additive, time-dependent perturbation as control. We solve for the perturbations using an optimal-control-based algorithm to optimize a differentiable distribution-matching objective while penalizing control effort to preserve data fidelity. Experiment results in image generation demonstrate that our proposed plug-and-play approach can better align attribute distributions to diverse test-time targets compared to inference-time baselines, without retraining or fine-tuning the pretrained diffusion model.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 5