Variational Control for Guidance in Diffusion Models

2025-02-06

AI Summary

Existing diffusion guidance methods often rely on additional training or task-specific customization, limiting their generalizability and efficiency. This paper introduces Diffusion Trajectory Matching (DTM), the first framework that formulates guidance as a variational optimal control problem under terminal constraints—enabling universal, task-driven generation without fine-tuning pre-trained diffusion models. DTM unifies diverse guidance paradigms by integrating variational inference with stochastic optimal control, optimizing diffusion trajectories to satisfy terminal observational constraints. On ImageNet nonlinear blind deblurring, DTM achieves an FID of 34.31, improving upon the best training-free baseline by 43.76. Moreover, it generalizes effectively to various linear and blind nonlinear inverse problems, substantially expanding the zero-shot applicability frontier of pre-trained diffusion models in complex inverse settings.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Search and Optimization: Learning to Search
• Machine Learning: Deep Generative Models & Autoencoders

Application Category

• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Diffusion models exhibit excellent sample quality, but existing guidance methods often require additional model training or are limited to specific tasks. We revisit guidance in diffusion models from the perspective of variational inference and control, introducing Diffusion Trajectory Matching (DTM) that enables guiding pretrained diffusion trajectories to satisfy a terminal cost. DTM unifies a broad class of guidance methods and enables novel instantiations. We introduce a new method within this framework that achieves state-of-the-art results on several linear and (blind) non-linear inverse problems without requiring additional model training or modifications. For instance, in ImageNet non-linear deblurring, our model achieves an FID score of 34.31, significantly improving over the best pretrained-method baseline (FID 78.07). We will make the code available in a future update.

Problem

• Enhances diffusion model guidance
• Unifies diverse guidance methods
• Improves inverse problems performance

Innovation

• Unifies guidance via Diffusion Trajectory Matching
• Enables guiding pretrained diffusion trajectories
• Achieves state-of-the-art without model retraining