PaperlayerView exact revision on arXiv

Source-linked AI summary

DOHF: Online Diffusion Fine-tuning with Doob's $h$-transform Guidance

Zhengyi Guo, Jiayuan Sheng, Wenpin Tang, David D. Yao

arXiv:2609.31882v2cs.LGcs.AI

TL;DR

Reward-based diffusion fine-tuning can struggle when high-reward regions are poorly represented and correction estimation relies on sparse trajectories. DOHF re-estimates and distills an h-transform correction under the evolving rollout policy; across three settings, it reports faster adaptation and strong reward alignment in such regions.

  • Problem

    Fixed-reference h-guidance may rely on trajectories with little information in high-reward regions that are rare under the original proposal.

    Source: §1 · p. 1

  • Method

    DOHF assigns optimality weights to current-policy samples, estimates the corresponding local h-correction, and distills it into the generative model online.

    Source: §1 · p. 1

  • Results

    Across rare-mode adaptation, constrained contingency-table generation, and text-to-image fine-tuning, DOHF reports faster adaptation and strong reward alignment when high-reward regions are poorly represented.

    Source: §1 · p. 2, §5 · p. 9

  • Takeaways & Limitations

    The framework offers an online way to adapt the rollout distribution together with the model, particularly when reward-relevant regions are difficult to sample from the base model.

    Source: §5 · p. 9

    For contingency-table generation, marginals may be noisy or partially observed, so they are treated as soft structural targets.

    Source: §4.2 · p. 7

Abstract

from arXiv · show

Reward-based diffusion fine-tuning faces practical challenges when desirable outcomes are rare or conditioning corrections are costly to estimate. In this work, we propose Diffusion Online $h$-guidance Fine-tuning (DOHF), which turns Doob's $h$-transform into a practical online training algorithm. DOHF assigns optimality weights to generated samples, estimates the normalized local correction $\nabla\log h$ under the current rollout policy, and distills it directly into the generative model. Theoretically, we characterize the population-optimal DiffusionNFT update as well as the various classfier free guidance methods through a unified $h$-transform perspective. Methodologically, our framework accommodates black-box and non-differentiable rewards without additional network evaluations. We further show improved alignments under three empirical scenarios. Our work demonstrates how adapting probabilistic conditioning through inexpensive estimation and iterative distillation can improve generative learning across statistical sampling and visual generation.

1 Introduction

DOHF addresses the difficulty of estimating useful conditioning corrections when high-reward outcomes are rare by re-estimating the Doob h-correction under an evolving rollout policy and distilling it into the model. Across three settings, it reports faster adaptation and strong reward alignment, while also characterizing how DiffusionNFT relates to the exact h-transform update.

1 Introduction

  • When high-reward regions are rare under the original proposal, fixed-reference guidance may estimate corrections from trajectories with little information where adaptation is most needed.This motivates estimating the h-function under the evolving rollout policy.Source: §1 · p. 1
  • DOHF assigns generated samples optimality weights, estimates the resulting local correction under the current rollout policy, and distills it into the generative model.The estimator reuses the velocity already computed at the current step, requires no additional network evaluations, and supports non-differentiable weights.Source: §1 · pp. 1–2
  • Across rare-mode adaptation, constrained contingency-table generation, and SD3.5-Medium text-to-image fine-tuning, DOHF reports faster adaptation and strong reward alignment when high-reward regions are poorly represented.The reported benefit comes from jointly adapting the rollout distribution and the h-correction.Source: §1 · p. 2
  • A single optimality weight represents exponential tilting, hard constraints, smooth relaxations, and relative-optimality objectives as targets proportional to w(x) p_old(x).This retains conditioning structure, including indicator-type hard constraints.Source: §1 · p. 2
  • DiffusionNFT's population-optimal update matches the exact h-transform only with the state-dependent guidance strength β = 2h(x_t, t, c), which DOHF estimates explicitly.The paper reports compatibility with classifier-free guidance for DOHF, unlike forward-matching baselines.Source: §1 · p. 2

2 Preliminaries

The preliminaries set up an online rollout policy and express fine-tuning as conditioning on an optimality event. They define the intermediate h-function and show how its score correction changes the flow-matching velocity target.

2 Preliminaries

  • The setup distinguishes the pre-trained base model from the current rollout policy used to generate samples online, with c denoting optional conditioning context.Examples of c include a text prompt, market covariate, or no context.Source: §2 · p. 3

2.1 A Unified Optimality Weight

  • An optimality weight w assigns generated samples a probability of being optimal, and conditioning on that event yields a target proportional to w(x_0) p_old(x_0 | c).Multiplying w by a positive constant does not change the resulting target.Source: §2.1 · p. 3

2.2 The h-Function

  • The h-function is the probability of optimality given the intermediate state, time, and context; its log-gradient supplies the score correction for fine-tuning.At time zero, h equals the terminal optimality weight.Source: §2.2 · p. 3

2.3 Flow Matching

  • For Gaussian flow matching, the fine-tuned velocity target subtracts t/(1−t) times the h-function log-gradient from the old velocity.This target is presented as a learning target for the fine-tuned velocity field.Source: §2.3 · p. 4

3 Diffusion Online h-guidance Fine-tuning

The section formulates online estimation of the h-correction along the generative trajectory and describes its use in training and sampling. It also characterizes when DiffusionNFT's population optimum matches the exact h-transformed velocity.

3 Diffusion Online h-guidance Fine-tuning

  • For Flow Matching, the backward process can be written as an SDE with selectable λ_t; setting λ_t=0 recovers the original ODE form.The method uses backward Gaussian transitions over discrete reverse times.Source: §3 · p. 4

3.1 Simulation-Based Estimator of ∇x log h(x, t, c) [56]

  • Monte Carlo estimation supplies ∇x log h during online fine-tuning, and a one-step-lookahead surrogate reduces the cost of multiple rollouts.The estimator reuses the already-computed rollout velocity at zero extra NFE and does not require differentiable w.Source: §3 · p. 4, §3.1 · p. 4

3.2 Training and Sampling

  • Training initializes the model and rollout policy from the pretrained model, samples rollout states, estimates ∇log h, and updates the model; a slow-mix strategy smooths on-policy sampling with an exponential moving average.Text-to-image sampling can use classifier-free guidance by leaving some prompts empty during training; other tasks sample directly from the fine-tuned model.Source: §3.1 · p. 5, §3.2 · p. 5

3.3 Population-Level Relation to DiffusionNFT

  • DiffusionNFT's population-optimal velocity coincides with the exact h-transformed velocity, while fixed β follows the unnormalized ∇h direction and attenuates updates in low-h regions.DOHF instead estimates normalized ∇log h under the evolving rollout policy, without a state-dependent guidance hyperparameter.Source: §3.3 · p. 6

4 Experiment

Across synthetic, structured-sampling, and image settings, DOHF improves rare-mode concentration, margin matching, and image reward alignment through online fine-tuning.

4.1 Gaussian Mixture

  • 100% of samples reach the minority mode after 200 DOHF updates, versus 14.88% after 2,000 off-policy guidance updates.DOHF also better preserves the target distribution, with mean (−0.82, −0.88) and standard deviation (0.38, 0.38).Source: §4.1 · p. 6, Table 1 · p. 7

4.2 Contingency Table

  • After 300 updates, DOHF reduces average and maximum standardized margin errors to 0.210 and 0.531, versus 0.379 and 1.005 for off-policy guidance after 800 updates.None of 500,000 uniformly sampled pretraining tables exactly satisfied all target margins.Source: Table 2 · p. 8, §4.2 · p. 8
  • DOHF achieves the smallest Wasserstein distance to the Diaconis–Gangolli sampler among the three methods.The DG sampler generates approximately uniform contingency tables with fixed margins; the experiment uses perturbed margins.Source: §4.2 · p. 8

4.3 Image Fine-tuning

  • Without classifier-free guidance (CFG), DOHF surpasses FlowGRPO on most metrics using 2k iterations versus 9k, and performs slightly better than DiffusionNFT [54] overall after 2k iterations.Source: §4.3 · p. 8

SD3.5M w/ CFG DOHF w/o CFG DOHF w/ CFG FlowGRPO

  • With classifier-free guidance, all quality scores increase to some extent at little cost to diversity, while GenEval rises from 44.27% to 62.52%.GenEval measures out-of-domain, rule-based accuracy on compositional prompts.Source: SD3.5M w/ CFG DOHF w/o CFG DOHF w/ CFG… · p. 9

5 Conclusion

The paper presents DOHF as online diffusion fine-tuning grounded in Doob’s h-transform. Across three settings, it reports benefits from adapting the rollout distribution alongside the model, particularly when reward-relevant regions are difficult to sample from the base model.

5 Conclusion

  • DOHF turns reward-based adaptation into a local score or velocity correction estimated online and distilled directly into the generative model.The framework treats continuous rewards, hard and soft constraints, and relative optimality objectives, and clarifies its connection to DiffusionNFT.Source: §5 · p. 9
  • Across rare-mode adaptation, constrained contingency-table generation, and text-to-image fine-tuning, DOHF benefits from adapting the rollout distribution together with the model.The paper particularly reports this benefit when reward-relevant regions are difficult to sample from the base model.Source: §5 · p. 9
  • The authors present online h-transform fine-tuning as a promising alternative to inference-time guidance and trajectory-based reinforcement learning.They describe the framework as a flexible interface between probabilistic conditioning and generative-model optimization.Source: §5 · p. 9

A Relation between score and velocity

This section relates score and velocity through conditional expectations and describes an alternative estimator of the h-correction. It also places DiffusionNFT’s guidance direction and training objectives within the h-transform framework.

A Relation between score and velocity

  • The score marginalization identity expresses the marginal score as a conditional expectation, while the noise conditional expectation is related to the velocity field.The section presents these relationships in the derivation connecting score and velocity.Source: A Relation between score and velocity · p. 15
  • The alternative h-correction estimator approximates h with the reward evaluated at a Tweedie point estimate; its finite-difference method requires differentiable rewards.The finite-difference method requires two evaluations of the old velocity and one backward pass through the reward, but no M rollouts.Source: A Relation between score and velocity · p. 15

C Relation to Diffusion Negative-aware Fine-Tuning (NFT)

C.1 Distribution split coefficient α(xt) is equvalent to h(xt, t, c)

The appendix connects DiffusionNFT’s distribution split coefficient to h and interprets its guidance direction and anti-guidance field through the h-transform.

C.1 Distribution split coefficient α(xt) is equvalent to h(xt, t, c)

  • The old-policy velocity is a mixture of the optimal and non-optimal velocities, weighted by α(xt) and 1 − α(xt); the derivation identifies α(xt) with h(xt, t, c).The optimal-branch posterior coincides with the fine-tuned distribution, and velocity is a linear functional of the posterior mean.Source: §C.1 · p. 16

C.2 The Guidance Direction ∆and v−as Anti-Guidance

  • With α = h, DiffusionNFT’s improvement direction Δ is the h-gradient ∇xh, rescaled by −t/(1−t), rather than the normalized ∇log h correction.Source: §C.2 · p. 16
  • The non-optimal velocity v− is the h-transform velocity for the complementary event O = 0, making it an anti-guidance field pointing away from the optimal region.Source: §C.2 · p. 16

C.3 Double-Branch Optimal Velocity

The double-branch objective yields an h-attenuated guidance update at fixed β, while the single-branch objective has an exact population optimum for any fixed β.

C.3 Double-Branch Optimal Velocity

  • For fixed β, the double-branch update is damped by local optimality h and vanishes as h approaches zero.The section derives this from the double-branch training objective.Source: §C.3 · pp. 17–18
  • Choosing β = 2h(xt, t, c) makes the guidance strength state-dependent, allowing more movement where optimality is likelier and less where it is not.DiffusionNFT cannot choose this state-dependent β without explicitly estimating h, so it leaves β as a hyperparameter.Source: §C.3 · p. 18

C.4 Single-Branch Optimal Velocity

  • For the single-branch objective, the pointwise L2 minimizer is the conditional mean regardless of the positive weight h.This optimum is exact for any fixed β, without h-dependent damping.Source: §C.4 · p. 18

C.5 Unifying the Three via h

The three guidance optima are unified through h: setting the double-branch strength to 2h cancels its damping and recovers the exact h-transformed update.

C.5 Unifying the Three via h

  • Setting β(xt, t, c) = 2h(xt, t, c) cancels double-branch damping and yields the exact h-transformed velocity, matching the single-branch optimum at β = 1.The resulting velocity is vft.Source: §C.5 · p. 18

D.1 Birthday-Deathday Data

This experiment uses a classical birthday–deathday contingency table for 82 descendants of Queen Victoria and prepares normalized tables for flow-model pretraining.

D.1 Birthday-Deathday Data

  • The birthday–deathday data are a 12 × 12 contingency table of 82 descendants of Queen Victoria.The original table is reproduced as Table 4.Source: §D.1 · p. 18, Table 4 · p. 19

D.2 Pretraining the Contingency-Table Flow Model

  • For pretraining, integer tables with total count N = 82 are normalized into probability vectors on the 144-dimensional simplex for a continuous flow model.At inference, the generated probability table is converted back to an integer table with total count exactly N using largest-remainder rounding.Source: §D.2 · p. 19

E Additional Details on Text-to-Image Task

The text-to-image setup specifies rollout, sampling, and optimization choices; additional qualitative comparisons show aesthetic-quality and semantic-alignment prompts across methods.

E Additional Details on Text-to-Image Task

  • Text-to-image rollouts use 40 time steps, Euler integration in training, and DPM-Solver++ in inference; online h-guidance uses CFG scale 4 by default.The SD3.5-Medium CFG scale is 3.5 by default.Source: E Additional Details on Text-to-Image… · p. 20
  • Online data collection starts on-policy and gradually approaches an equal mixture of the old and fine-tuned velocity fields after 500 iterations.The off-policy guidance baseline uses the base model for both rollout policy and regression target.Source: E Additional Details on Text-to-Image… · p. 20
  • The implementation uses seven Nvidia L40S GPUs and restricts training times to 0.05 ≤ t ≤ 0.9 because two Flow-equivalent SDE denominators approach zero as t approaches 1.It also sets λt = 1, learning rate 3e-5, gradient clipping 1.0, and LoRA regularization 1e-5.Source: E Additional Details on Text-to-Image… · p. 20

F Additional Qualitative Comparison

Related papers

Loading 2609.31882v2…