Direct Distributional Optimization for Provable Alignment of Diffusion Models

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Kawata, Ryotaro, Oko, Kazusato, Nitanda, Atsushi, Suzuki, Taiji
Formato: Preprint
Publicado: 2025
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866910860963741696
author Kawata, Ryotaro
Oko, Kazusato
Nitanda, Atsushi
Suzuki, Taiji
author_facet Kawata, Ryotaro
Oko, Kazusato
Nitanda, Atsushi
Suzuki, Taiji
contents We introduce a novel alignment method for diffusion models from distribution optimization perspectives while providing rigorous convergence guarantees. We first formulate the problem as a generic regularized loss minimization over probability distributions and directly optimize the distribution using the Dual Averaging method. Next, we enable sampling from the learned distribution by approximating its score function via Doob's $h$-transform technique. The proposed framework is supported by rigorous convergence guarantees and an end-to-end bound on the sampling error, which imply that when the original distribution's score is known accurately, the complexity of sampling from shifted distributions is independent of isoperimetric conditions. This framework is broadly applicable to general distribution optimization problems, including alignment tasks in Reinforcement Learning with Human Feedback (RLHF), Direct Preference Optimization (DPO), and Kahneman-Tversky Optimization (KTO). We empirically validate its performance on synthetic and image datasets using the DPO objective.
format Preprint
id arxiv_https___arxiv_org_abs_2502_02954
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Direct Distributional Optimization for Provable Alignment of Diffusion Models
Kawata, Ryotaro
Oko, Kazusato
Nitanda, Atsushi
Suzuki, Taiji
Machine Learning
We introduce a novel alignment method for diffusion models from distribution optimization perspectives while providing rigorous convergence guarantees. We first formulate the problem as a generic regularized loss minimization over probability distributions and directly optimize the distribution using the Dual Averaging method. Next, we enable sampling from the learned distribution by approximating its score function via Doob's $h$-transform technique. The proposed framework is supported by rigorous convergence guarantees and an end-to-end bound on the sampling error, which imply that when the original distribution's score is known accurately, the complexity of sampling from shifted distributions is independent of isoperimetric conditions. This framework is broadly applicable to general distribution optimization problems, including alignment tasks in Reinforcement Learning with Human Feedback (RLHF), Direct Preference Optimization (DPO), and Kahneman-Tversky Optimization (KTO). We empirically validate its performance on synthetic and image datasets using the DPO objective.
title Direct Distributional Optimization for Provable Alignment of Diffusion Models
topic Machine Learning
url https://arxiv.org/abs/2502.02954