Direct Distributional Optimization for Provable Alignment of Diffusion Models

2025-02-05

AI Summary

This work addresses the alignment problem for diffusion models from a distributional optimization perspective. Methodologically, it formulates alignment as a minimization problem with a distributional regularization term, directly optimizes the target distribution via dual averaging, and approximates the score function using Doob’s h-transform to ensure sampleability. Key contributions include: (i) the first rigorous convergence guarantee and end-to-end sampling error bound for diffusion model alignment; (ii) a sampling complexity independent of isoperimetric or Poincaré constants—overcoming fundamental limitations of prior analyses; and (iii) a theoretically established upper bound on the sampling error under the shifted distribution when the exact score function is available. Experiments on synthetic and image data validate alignment efficacy across canonical scenarios including RLHF, DPO, and KTO.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search
• Machine Learning: Distributed Machine Learning & Federated Learning

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

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.

Problem

• Optimizes diffusion model alignment with distribution methods
• Uses Dual Averaging for direct distribution optimization
• Applies to RLHF, DPO, and KTO alignment tasks

Innovation

• Direct Distributional Optimization method
• Dual Averaging for optimization
• Doob's h-transform for sampling