# Inference-Time Attribute Distribution Alignment for Unconditional Diffusion

arXiv: https://arxiv.org/abs/2605.07456 · HTML: https://mlsanity.com/paper/2605.07456/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion foundations and sampling, one of 14 directions in Diffusion models.
- Builds on 67 of its 81 references that are in the map.
- First posted on arXiv 8 May 2026.

## Placed in

- Diffusion models / [Diffusion foundations and sampling](https://mlsanity.com/diffusion/diffusion-foundations-and-sampling/)

## Builds on (most cited first, 30 of 67)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Denoising Diffusion Implicit Models](https://mlsanity.com/paper/2010.02502/) (2020)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Elucidating the Design Space of Diffusion-Based Generative Models](https://mlsanity.com/paper/2206.00364/) (2022)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [Video Diffusion Models](https://mlsanity.com/paper/2204.03458/) (2022)
- [Diffusion Posterior Sampling for General Noisy Inverse Problems](https://mlsanity.com/paper/2209.14687/) (2022)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Denoising Diffusion Restoration Models](https://mlsanity.com/paper/2201.11793/) (2022)
- [Diffusion Model Alignment Using Direct Preference Optimization](https://mlsanity.com/paper/2311.12908/) (2023)
- [Solving Inverse Problems in Medical Imaging with Score-Based Generative Models](https://mlsanity.com/paper/2111.08005/) (2021)
- [Lumiere: A Space-Time Diffusion Model for Video Generation](https://mlsanity.com/paper/2401.12945/) (2024)
- [Understanding Diffusion Models: A Unified Perspective](https://mlsanity.com/paper/2208.11970/) (2022)
- [Perception Prioritized Training of Diffusion Models](https://mlsanity.com/paper/2204.00227/) (2022)
- [Permutation Invariant Graph Generation via Score-Based Generative Modeling](https://mlsanity.com/paper/2003.00638/) (2020)
- [Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control](https://mlsanity.com/paper/2409.08861/) (2024)
- [MM-Diffusion: Learning Multi-Modal Diffusion Models for Joint Audio and Video Generation](https://mlsanity.com/paper/2212.09478/) (2022)
- [Flow Matching in Latent Space](https://mlsanity.com/paper/2307.08698/) (2023)
- [An optimal control perspective on diffusion-based generative modeling](https://mlsanity.com/paper/2211.01364/) (2022)
- [Fair Generative Modeling via Weak Supervision](https://mlsanity.com/paper/1910.12008/) (2019)

## Built on by (most cited first, 2 of 2)

- [A Mean-Field Framework for Inference-Time Distributional Control of Diffusion Models](https://mlsanity.com/paper/2608.08770/) (2026)
- [Debias Anything: Fairness with Diversity without Supervision in Diffusion Models](https://mlsanity.com/paper/2610.01815/) (2026)

## Abstract (arXiv)

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.