Efficient Controllable Diffusion via Optimal Classifier Guidance

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Oertell, Owen, Sun, Shikun, Chen, Yiding, Zhou, Jin Peng, Wang, Zhiyong, Sun, Wen
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866918036813905920
author Oertell, Owen
Sun, Shikun
Chen, Yiding
Zhou, Jin Peng
Wang, Zhiyong
Sun, Wen
author_facet Oertell, Owen
Sun, Shikun
Chen, Yiding
Zhou, Jin Peng
Wang, Zhiyong
Sun, Wen
contents The controllable generation of diffusion models aims to steer the model to generate samples that optimize some given objective functions. It is desirable for a variety of applications including image generation, molecule generation, and DNA/sequence generation. Reinforcement Learning (RL) based fine-tuning of the base model is a popular approach but it can overfit the reward function while requiring significant resources. We frame controllable generation as a problem of finding a distribution that optimizes a KL-regularized objective function. We present SLCD -- Supervised Learning based Controllable Diffusion, which iteratively generates online data and trains a small classifier to guide the generation of the diffusion model. Similar to the standard classifier-guided diffusion, SLCD's key computation primitive is classification and does not involve any complex concepts from RL or control. Via a reduction to no-regret online learning analysis, we show that under KL divergence, the output from SLCD provably converges to the optimal solution of the KL-regularized objective. Further, we empirically demonstrate that SLCD can generate high quality samples with nearly the same inference time as the base model in both image generation with continuous diffusion and biological sequence generation with discrete diffusion. Our code is available at https://github.com/Owen-Oertell/slcd
format Preprint
id arxiv_https___arxiv_org_abs_2505_21666
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Efficient Controllable Diffusion via Optimal Classifier Guidance
Oertell, Owen
Sun, Shikun
Chen, Yiding
Zhou, Jin Peng
Wang, Zhiyong
Sun, Wen
Machine Learning
Artificial Intelligence
The controllable generation of diffusion models aims to steer the model to generate samples that optimize some given objective functions. It is desirable for a variety of applications including image generation, molecule generation, and DNA/sequence generation. Reinforcement Learning (RL) based fine-tuning of the base model is a popular approach but it can overfit the reward function while requiring significant resources. We frame controllable generation as a problem of finding a distribution that optimizes a KL-regularized objective function. We present SLCD -- Supervised Learning based Controllable Diffusion, which iteratively generates online data and trains a small classifier to guide the generation of the diffusion model. Similar to the standard classifier-guided diffusion, SLCD's key computation primitive is classification and does not involve any complex concepts from RL or control. Via a reduction to no-regret online learning analysis, we show that under KL divergence, the output from SLCD provably converges to the optimal solution of the KL-regularized objective. Further, we empirically demonstrate that SLCD can generate high quality samples with nearly the same inference time as the base model in both image generation with continuous diffusion and biological sequence generation with discrete diffusion. Our code is available at https://github.com/Owen-Oertell/slcd
title Efficient Controllable Diffusion via Optimal Classifier Guidance
topic Machine Learning
Artificial Intelligence
url https://arxiv.org/abs/2505.21666