Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Tilt Matching for Scalable Sampling and Fine-Tuning

Peter Potaptchik, Cheuk-Kit Lee, Michael S. Albergo

TL;DR

Tilt Matching provides a scalable framework for adapting transport velocities to reward-tilted targets without backpropagating through trajectories or requiring reward gradients. By deriving a covariance-based evolution (and its implicit-all-orders variant) and connecting it to Doob's h-transform, the method achieves lower-variance updates and practical training stability. Empirical results show state-of-the-art sampling for Lennard-Jones systems and competitive fine-tuning of Stable Diffusion without reward multipliers, underscoring the approach's robustness and broad applicability to tilting flow-based models. The work advances a principled, regression-based alternative to trajectory-based SOC methods for reward-tilted transport in both sampling and generative refinement tasks.

Abstract

We propose a simple, scalable algorithm for using stochastic interpolants to sample from unnormalized densities and for fine-tuning generative models. The approach, Tilt Matching, arises from a dynamical equation relating the flow matching velocity to one targeting the same distribution tilted by a reward, implicitly solving a stochastic optimal control problem. The new velocity inherits the regularity of stochastic interpolant transports while also being the minimizer of an objective with strictly lower variance than flow matching itself. The update to the velocity field can be interpreted as the sum of all joint cumulants of the stochastic interpolant and copies of the reward, and to first order is their covariance. The algorithms do not require any access to gradients of the reward or backpropagating through trajectories of the flow or diffusion. We empirically verify that the approach is efficient and highly scalable, providing state-of-the-art results on sampling under Lennard-Jones potentials and is competitive on fine-tuning Stable Diffusion, without requiring reward multipliers. It can also be straightforwardly applied to tilting few-step flow map models.

Tilt Matching for Scalable Sampling and Fine-Tuning

TL;DR

Tilt Matching provides a scalable framework for adapting transport velocities to reward-tilted targets without backpropagating through trajectories or requiring reward gradients. By deriving a covariance-based evolution (and its implicit-all-orders variant) and connecting it to Doob's h-transform, the method achieves lower-variance updates and practical training stability. Empirical results show state-of-the-art sampling for Lennard-Jones systems and competitive fine-tuning of Stable Diffusion without reward multipliers, underscoring the approach's robustness and broad applicability to tilting flow-based models. The work advances a principled, regression-based alternative to trajectory-based SOC methods for reward-tilted transport in both sampling and generative refinement tasks.

Abstract

We propose a simple, scalable algorithm for using stochastic interpolants to sample from unnormalized densities and for fine-tuning generative models. The approach, Tilt Matching, arises from a dynamical equation relating the flow matching velocity to one targeting the same distribution tilted by a reward, implicitly solving a stochastic optimal control problem. The new velocity inherits the regularity of stochastic interpolant transports while also being the minimizer of an objective with strictly lower variance than flow matching itself. The update to the velocity field can be interpreted as the sum of all joint cumulants of the stochastic interpolant and copies of the reward, and to first order is their covariance. The algorithms do not require any access to gradients of the reward or backpropagating through trajectories of the flow or diffusion. We empirically verify that the approach is efficient and highly scalable, providing state-of-the-art results on sampling under Lennard-Jones potentials and is competitive on fine-tuning Stable Diffusion, without requiring reward multipliers. It can also be straightforwardly applied to tilting few-step flow map models.
Paper Structure (34 sections, 14 theorems, 97 equations, 6 figures, 2 tables, 2 algorithms)

This paper contains 34 sections, 14 theorems, 97 equations, 6 figures, 2 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Preliminaries
  3. Dynamical Transport and Stochastic Interpolants
  4. Flow Maps.
  5. Fine-tuning and Sampling as Tilting
  6. Deriving Tilt Matching
  7. Explicit Tilt Matching
  8. Regularity of transport.
  9. Approximation error due to incomplete minimization of the objective.
  10. Discretization error in $a$.
  11. PINN objective.
  12. Implicit Tilt Matching
  13. Higher-order expansions.
  14. Expanding to all orders.
  15. Variance reduction via control variates.
  16. ...and 19 more sections

Key Result

Proposition 1

(Esscher Transform.) Let $I_t^a = \alpha_t x_0 + \beta_t x_1^a$ be the interpolant constructed from samples $x_1^a \sim \rho_{1,a}$. Then the augmented velocity field $b_{t,a}(x)$ is given by: Furthermore, for any shift $h$ from $a$ to $a+h$, the updated velocity $b_{t,a+h}(x)$ satisfies:

Figures (6)

  • Figure 1: For each prompt, we display three paired samples coming from the Base model vs our Tilt Matching method.
  • Figure 2: Schematic overview of the proposed method. When flow matching with a stochastic interpolant $I_t^{a=0}$ is used to learn a generative model $b_{t,a=0}$ that samples $\rho_{t,a=0}$ and in particular has terminal samples $\rho_{1,a=0}$ (gray curve), then the evolution of that velocity field in $a$ in order to sample $\rho_{1,a>0} = \tfrac{1}{Z_a}\rho_{1,0}e^{a r(x)}$, where $r$ is a reward function, has closed form given by the conditional covariance of the dynamics of the interpolant at $(t,a)$ and the reward. The velocity field, denoted as up or down arrows showing direction of motion in $x$, changes from negative to positive in the above toy example.
  • Figure 3: Pictorial additivity of higher-order corrections to $b_{t,a+h}$. The first-order term is the covariance, while higher-order terms are cumulants $\kappa^n$.
  • Figure 4: Example improvements to Stable Diffusion 1.5 using Tilt Matching and ImageReward. Key parts of each prompt are highlighted in color.
  • Figure 5: Sampling performance and ablation results for LJ13.
  • ...and 1 more figures

Theorems & Definitions (24)

  • Proposition 1
  • Proposition 2
  • Proposition 3
  • Proposition 4
  • Proposition 5
  • Remark 1
  • Proposition 6
  • Proposition 7
  • Proposition 7
  • proof
  • ...and 14 more