Inference-Time Attribute Distribution Alignment for Unconditional Diffusion

TL;DR

This paper introduces a novel inference-time method for aligning attribute distributions in unconditional diffusion models, enabling controlled generation without retraining.

Contribution

It formulates attribute distribution alignment as an optimal control problem over the diffusion process, providing a plug-and-play solution that improves distribution matching.

Findings

Better attribute distribution alignment compared to baselines.

No need for retraining or finetuning the pretrained model.

Effective in diverse, flexible test-time scenarios.

Abstract

Inference-time controllable generation is essential for real-world applications of unconditional diffusion models. However, most existing techniques focus on individual samples, struggling in applications that require the sample population to follow specific attribute distributions (e.g., demographic balance or semantic proportions). We formalize this setting as the inference-time attribute distributional alignment problem for pretrained unconditional diffusion models. To address this, we cast inference-time 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 additive, time-dependent perturbations 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 and flexible test-time targets compared to baselines, without retraining or finetuning the pretrained diffusion model.