Fine-Tuning of Continuous-Time Diffusion Models as Entropy-Regularized Control

2024/02/23 by Masatoshi Uehara, Yulai Zhao, Uehara, Masatoshi +15 · 35 citations
Computer Science · Mathematics · Physics and Astronomy · #Advanced Mathematical Modeling in Engineering #Artificial Intelligence (cs.AI) #FOS: Computer and information sciences #Machine Learning (cs.LG) #Machine Learning (stat.ML) #Model Reduction and Neural Networks #Numerical methods in inverse problems

paper · pdf · doi:10.48550/arxiv.2402.15194

openalex publication_date 2024/02/23 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

Diffusion models excel at capturing complex data distributions, such as those of natural images and proteins. While diffusion models are trained to represent the distribution in the training dataset, we often are more concerned with other properties, such as the aesthetic quality of the generated images or the functional properties of generated proteins. Diffusion models can be finetuned in a goal-directed way by maximizing the value of some reward function (e.g., the aesthetic quality of an image). However, these approaches may lead to reduced sample diversity, significant deviations from the training data distribution, and even poor sample quality due to the exploitation of an imperfect reward function. The last issue often occurs when the reward function is a learned model meant to approximate a ground-truth "genuine" reward, as is the case in many practical applications. These challenges, collectively termed "reward collapse," pose a substantial obstacle. To address this reward collapse, we frame the finetuning problem as entropy-regularized control against the pretrained diffusion model, i.e., directly optimizing entropy-enhanced rewards with neural SDEs. We present theoretical and empirical evidence that demonstrates our framework is capable of efficiently generating diverse samples with high genuine rewards, mitigating the overoptimization of imperfect reward models.

Cited by

• Generalized Fine-Tuning of Diffusion Models via Stochastic Control and FBSDEs
• Value Gradient Guidance for Flow Matching Alignment
• Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function
• Iterative Tilting for Diffusion Fine-Tuning
• Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-Tuning
• Test-Time Alignment of Text-to-Image Diffusion Models via Null-Text Embedding Optimisation
• ProxT2I: Efficient Reward-Guided Text-to-Image Generation via Proximal Diffusion
• Coffee: Controllable Diffusion Fine-tuning
• Embodiment Transfer Learning for Vision-Language-Action Models
• Schrödinger bridge for generative AI: Soft-constrained formulation and convergence analysis
• Understanding Sampler Stochasticity in Training Diffusion Models for RLHF
• Calibrating Generative Models
• Fine-Tuning Diffusion Models via Intermediate Distribution Shaping
• Diffusion Alignment as Variational Expectation-Maximization
• TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion
• DriftLite: Lightweight Drift Control for Inference-Time Scaling of Diffusion Models
• Composition and Alignment of Diffusion Models using Constrained Learning
• Source-Guided Flow Matching
• Trust Region Constrained Measure Transport in Path Space for Stochastic Optimal Control and Inference
• Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models
• Inference-Time Scaling of Diffusion Language Models with Particle Gibbs Sampling
• Stein Diffusion Guidance: Training-Free Posterior Correction for Sampling Beyond High-Density Regions
• GeoAda: Efficiently Finetune Geometric Diffusion Models with Equivariant Adapters
• Posterior Inference in Latent Space for Scalable Constrained Black-box Optimization
• Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models
• Nabla-R2D3: Effective and Efficient 3D Diffusion Alignment with 2D Rewards
• Provable Maximum Entropy Manifold Exploration via Diffusion Models
• Psi-Sampler: Initial Particle Sampling for SMC-Based Inference-Time Reward Alignment in Score Models
• ChatVLA-2: Vision-Language-Action Model with Open-World Embodied Reasoning from Pretrained Knowledge
• Efficient Controllable Diffusion via Optimal Classifier Guidance
• Diffusion Blend: Inference-Time Multi-Preference Alignment for Diffusion Models
• Scaling Image and Video Generation via Test-Time Evolutionary Search
• VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RL
• Conformal Policy Control
• Diamond Maps: Efficient Reward Alignment via Stochastic Flow Maps

Related

• Fine-tuning of diffusion models via stochastic control: entropy regularization and beyond
• Feedback Efficient Online Fine-Tuning of Diffusion Models
• Directly Fine-Tuning Diffusion Models on Differentiable Rewards
• Understanding Reinforcement Learning-Based Fine-Tuning of Diffusion Models: A Tutorial and Review
• Towards Controllable Diffusion Models via Reward-Guided Exploration