Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control

2024/09/13 by Carles Domingo-Enrich, Michal Drozdzal, Domingo-Enrich, Carles +5 · 60 citations
Mathematics · Economics, Econometrics and Finance · #Markov Chains and Monte Carlo Methods #Stochastic processes and financial applications

paper · pdf · doi:10.48550/arxiv.2409.08861

Abstract

Dynamical generative models that produce samples through an iterative process, such as Flow Matching and denoising diffusion models, have seen widespread use, but there have not been many theoretically-sound methods for improving these models with reward fine-tuning. In this work, we cast reward fine-tuning as stochastic optimal control (SOC). Critically, we prove that a very specific memoryless noise schedule must be enforced during fine-tuning, in order to account for the dependency between the noise variable and the generated samples. We also propose a new algorithm named Adjoint Matching which outperforms existing SOC algorithms, by casting SOC problems as a regression problem. We find that our approach significantly improves over existing methods for reward fine-tuning, achieving better consistency, realism, and generalization to unseen human preference reward models, while retaining sample diversity.

Cited by

• Temporal Difference Learning for Diffusion Models
• Expand and Prune: Maximizing Trajectory Diversity for Effective GRPO in Generative Models
• Anchoring Values in Temporal and Group Dimensions for Flow Matching Model Alignment
• MolGuidance: Advanced Guidance Strategies for Conditional Molecular Generation with Flow Matching
• TreeGRPO: Tree-Advantage GRPO for Online RL Post-Training of Diffusion Models
• Value Gradient Guidance for Flow Matching Alignment
• Control Consistency Losses for Diffusion Bridges
• Diffusion Fine-Tuning via Reparameterized Policy Gradient of the Soft Q-Function
• Iterative Tilting for Diffusion Fine-Tuning
• Diffusion-Augmented Markov Decision Processes for Maximum Entropy Reinforcement Learning
• High-dimensional Mean-Field Games by Particle-based Flow Matching
• Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-Tuning
• FlexiFlow: decomposable flow matching for generation of flexible molecular ensemble
• PID-controlled Langevin Dynamics for Faster Sampling of Generative Models
• Sublinear iterations can suffice even for DDPMs
• Learning Paths for Dynamic Measure Transport: A Control Perspective
• Flow Matching for Measure Transport and Feedback Stabilization of Control-Affine Systems
• Adaptive Divergence Regularized Policy Optimization for Fine-tuning Generative Models
• Fine-tuning Flow Matching Generative Models with Intermediate Feedback
• Continuous Q-Score Matching: Diffusion Guided Reinforcement Learning for Continuous-Time Control
• ImagerySearch: Adaptive Test-Time Search for Video Generation Beyond Semantic Dependency Constraints
• Reinforced sequential Monte Carlo for amortised sampling
• Understanding Sampler Stochasticity in Training Diffusion Models for RLHF
• Calibrating Generative Models
• Wavefunction Flows: Efficient Quantum Simulation of Continuous Flow Models
• Enhancing Reasoning for Diffusion LLMs via Distribution Matching Policy Optimization
• Control-Augmented Autoregressive Diffusion for Data Assimilation
• FOCUS: Optimal Control for Multi-Entity World Modeling in Text-to-Image Generation
• Proximal Diffusion Neural Sampler
• Fine-Tuning Diffusion Models via Intermediate Distribution Shaping
• Plug-and-Play Prompt Refinement via Latent Feedback for Diffusion Model Alignment
• Diffusion Alignment as Variational Expectation-Maximization
• Training-Free Reward-Guided Image Editing via Trajectory Optimal Control
• TR2-D2: Tree Search Guided Trajectory-Aware Fine-Tuning for Discrete Diffusion
• GLASS Flows: Transition Sampling for Alignment of Flow and Diffusion Models
• Advantage Weighted Matching: Aligning RL with Pretraining in Diffusion Models
• Diffusion Bridge or Flow Matching? A Unifying Framework and Comparative Analysis
• Sample-Efficient Optimisation over the Outputs of Generative Models
• Landing with the Score: Riemannian Optimization through Denoising
• CREPE: Controlling Diffusion with Replica Exchange
• DriftLite: Lightweight Drift Control for Inference-Time Scaling of Diffusion Models
• Directly Aligning the Full Diffusion Trajectory with Fine-Grained Human Preference
• Feynman-Kac-Flow: Inference Steering of Conditional Flow Matching to an Energy-Tilted Posterior
• 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
• MDNS: Masked Diffusion Neural Sampler via Stochastic Optimal Control
• Noise Hypernetworks: Amortizing Test-Time Compute in Diffusion Models
• FlowBack-Adjoint: Physics-Aware and Energy-Guided Conditional Flow-Matching for All-Atom Protein Backmapping
• Physics-Constrained Fine-Tuning of Flow-Matching Models for Generation and Inverse Problems
• Reinforcement Learning for Flow-Matching Policies
• Neuro-Inspired Inverse Learning for Planning and Control
• HOTA: Hamiltonian framework for Optimal Transport Advection
• RODS: Robust Optimization Inspired Diffusion Sampling for Detecting and Reducing Hallucination in Generative Models
• Flows and Diffusions on the Neural Manifold
• Inversion-DPO: Precise and Efficient Post-Training for Diffusion Models
• Stein Diffusion Guidance: Training-Free Posterior Correction for Sampling Beyond High-Density Regions
• Posterior Inference in Latent Space for Scalable Constrained Black-box Optimization
• Adjoint Schrödinger Bridge Sampler
• Diffusion Tree Sampling: Scalable inference-time alignment of diffusion models

Related

• A Taxonomy of Loss Functions for Stochastic Optimal Control
• Score Matching Diffusion Based Feedback Control and Planning of Nonlinear Systems
• Score as Action: Fine-Tuning Diffusion Generative Models by Continuous-time Reinforcement Learning
• Scores as Actions: a framework of fine-tuning diffusion models by continuous-time reinforcement learning
• FOCUS: Optimal Control for Multi-Entity World Modeling in Text-to-Image Generation