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 · 2 voices · 77 citations
Computer Science · Economics, Econometrics and Finance · Mathematics · #Applied mathematics #Artificial intelligence #Computer science #Diffusion #Flow (mathematics) #Generative grammar #Generative model #Geometry #Markov Chains and Monte Carlo Methods #Matching (statistics) #Mathematical optimization #Mathematics #Optimal control #Physics #Statistical physics #Statistics #Stochastic control #Stochastic processes and financial applications #Thermodynamics #cs.LG #math.OC #stat.ML

paper · pdf · doi:10.48550/arxiv.2409.08861

published in arXiv (Cornell University) (Cornell University)

openalex publication_date 2024/09/13 · arxiv published 2024/09/13 · openalex created_date 2024/10/23 · arxiv updated 2025/01/07 · openalex updated_date 2026/07/28

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
• Non-equilibrium Annealed Adjoint Sampler
• Nabla-R2D3: Effective and Efficient 3D Diffusion Alignment with 2D Rewards
• Provable Maximum Entropy Manifold Exploration via Diffusion Models
• ReGuidance: A Simple Diffusion Wrapper for Boosting Sample Quality on Hard Inverse Problems
• RNE: plug-and-play diffusion inference-time control and energy-based training
• Flow-GRPO: Training Flow Matching Models via Online RL
• Psi-Sampler: Initial Particle Sampling for SMC-Based Inference-Time Reward Alignment in Score Models
• Efficient Controllable Diffusion via Optimal Classifier Guidance
• Adaptive Diffusion Guidance via Stochastic Optimal Control
• Test-Time Scaling of Diffusion Models via Noise Trajectory Search
• Scaling Image and Video Generation via Test-Time Evolutionary Search
• Deep Learning for Continuous-Time Stochastic Control with Jumps
• VARD: Efficient and Dense Fine-Tuning for Diffusion Models with Value-based RL
• Minimum-Excess-Work Guidance
• Conformal Policy Control
• Diamond Maps: Efficient Reward Alignment via Stochastic Flow Maps
• Minimal-Action Discrete Schrödinger Bridge Matching for Peptide Sequence Design

Discussions

• On Monday, Carles Domingo-Enrich will present his new "Adjoint Matching: Fine-tuning Flow and Diffusion Generative Models with Memoryless Stochastic Optimal Control" arxiv.org/abs/2409.08861 Quite fre [bsky, 1 points, 0 comments]
• > [...] 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 fo [bsky, 0 points, 1 comments]

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