Inference-Time Attribute Distribution Alignment for Unconditional Diffusion

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.