A unified perspective on fine-tuning and sampling with diffusion and flow models
| Authors | Carles Domingo-Enrich et al. |
| Year | 2026 |
| Field | Statistics / ML |
| arXiv | 2605.00229 |
| Download | |
| Categories | stat.ML, cs.LG |
Abstract
We study the problem of training diffusion and flow generative models to sample from target distributions defined by an exponential tilting of a base density; a formulation that subsumes both sampling from unnormalized densities and reward fine-tuning of pre-trained models. This problem can be approached from a stochastic optimal control (SOC) perspective, using adjoint-based or score matching methods, or from a non-equilibrium thermodynamics perspective. We provide a unified framework encompassing these approaches and make three main contributions: (i) bias-variance decompositions revealing that Adjoint Matching/Sampling and Novel Score Matching have finite gradient variance, while Target and Conditional Score Matching do not; (ii) norm bounds on the lean adjoint ODE that theoretically support the effectiveness of adjoint-based methods; and (iii) adaptations of the CMCD and NETS loss functions, along with novel Crooks and Jarzynski identities, to the exponential tilting setting. We validate our analysis with reward fine-tuning experiments on Stable Diffusion 1.5 and 3.
Engineering Breakdown
Plain English
This paper provides a unified mathematical framework for training diffusion and flow-based generative models to sample from exponentially-tilted target distributions—a problem that covers both sampling from unnormalized densities and reward-based fine-tuning of pre-trained models. The authors bridge three different perspectives (stochastic optimal control, adjoint methods, and score matching) and prove that certain gradient estimators like Adjoint Matching and Novel Score Matching have finite variance, while others like Target and Conditional Score Matching have unbounded variance. They provide theoretical norm bounds on the lean adjoint ODE and derive practical algorithms. This work unifies previously disparate approaches and provides principled guidance on which methods to use based on variance properties and theoretical guarantees.
Core Technical Contribution
The core novelty is a unified theoretical framework that connects stochastic optimal control, adjoint-based methods, and score matching for exponentially-tilted sampling—three approaches that were previously treated separately in the literature. The authors prove formal bias-variance decompositions showing that Adjoint Matching and Novel Score Matching achieve finite gradient variance (critical for practical optimization), while Target and Conditional Score Matching suffer from unbounded variance that scales poorly with problem dimension. They derive norm bounds on the lean adjoint ODE trajectory, providing theoretical justification for when certain methods will be numerically stable. This bridges the gap between thermodynamics-inspired sampling and control-theoretic perspectives, enabling principled algorithm selection.
How It Works
The framework starts with a base density and an exponential tilting function (e.g., a reward function) that defines the target distribution—proportional to base_density × exp(reward). The authors formulate this as a stochastic optimal control problem where you must learn a drift function that transforms samples from the base distribution to the target. Three algorithmic approaches compete: (1) Adjoint Matching solves the optimal control problem by matching the learned model to gradients of a value function (computed via the adjoint method); (2) Score Matching fits the learned score function directly to target score differences; (3) Conditional Score Matching conditions on trajectory endpoints. The key insight is that Adjoint Matching and Novel Score Matching compute gradients through an ODE with well-controlled norms, yielding finite variance, while other variants suffer from gradient explosion. The framework unifies these by viewing them as different parameterizations of the same underlying stochastic control problem.
Production Impact
For teams building reward-fine-tuned generative models (e.g., image generation with aesthetic rewards, LLM decoding with constraint satisfaction), this framework provides concrete guidance: use Adjoint Matching or Novel Score Matching instead of naive score matching, because the variance properties directly impact gradient stability during training. In practice, this means faster convergence, fewer divergences, and better sample quality—especially critical when fine-tuning large pre-trained models where instability is expensive. The norm bounds on the lean adjoint ODE inform numerical solver choice (step size, integrator type), reducing trial-and-error tuning. Integration cost is moderate: you need to implement adjoint backprop through ODEs (standard in modern AD frameworks) and compute value function estimates. Latency is comparable to standard diffusion training since the core loop is identical; the benefit is statistical efficiency and theoretical guarantees rather than speed.
Limitations and When Not to Use This
The paper assumes you can evaluate or approximate the reward/tilting function efficiently—if reward computation is expensive or noisy, the framework's variance guarantees degrade. The theoretical results focus on gradient variance; they don't directly address sampling quality, mode coverage, or whether the learned model actually matches the target distribution's density (only that gradients are stable). The framework requires differentiability of the target distribution and base density, limiting applicability to discrete distributions without relaxations. The paper appears incomplete (abstract cuts off mid-sentence), so we cannot assess claims about 'norm bounds on the lean adjoint ODE' fully. Real-world deployment would need empirical validation that the variance improvements translate to wall-clock convergence gains at scale, and the method's behavior under high-dimensional rewards (e.g., multi-objective fine-tuning) remains unclear.
Research Context
This work synthesizes two active research directions: (1) score-based generative models and diffusion for sampling, which have dominated generative modeling since 2019; and (2) stochastic optimal control / reinforcement learning perspectives on diffusion (e.g., RL for fine-tuning, control-as-inference). It builds directly on prior work connecting diffusion to SOC (e.g., Guidance and the broader optimal transport / control literature) and score matching (Song et al., 2019). The exponential tilting formulation is motivated by recent work on reward fine-tuning pre-trained models (RLHF for LLMs, reward-guided image generation). The paper opens directions for: certifying other gradient estimators under the unified framework, extending to discrete/combinatorial targets, and combining with recent work on flow matching as an alternative to diffusion. It also informs the design of foundation model fine-tuning pipelines where theoretical guarantees on optimization stability are increasingly valued.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::