PaperlayerView exact revision on arXiv

Source-linked AI summary

Steepest Guidance: A Practical and Principled Approach to Inference-Time Alignment of Flow and Diffusion-based Models

Shokichi Takakura, Akifumi Wachi, Rei Higuchi, Kohei Miyaguchi, Taiji Suzuki

arXiv:2609.39091v1stat.MLcs.LG

TL;DR

Inference-time alignment is difficult because optimal guidance can be challenging to estimate and existing approaches have limitations for nonlinear rewards. Steepest Guidance treats alignment as sequential optimization in probability-measure space, and experiments report consistent outperformance and handling of nonlinear reward functionals.

  • Problem

    Inference-time alignment methods face challenges with nonlinear reward functionals, and finite-sample estimates of optimal guidance can be biased.

    Source: §1 · p. 2

  • Method

    Steepest Guidance treats inference-time alignment as sequential optimization in probability-measure space, using local improvement of the reward functional.

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

  • Results

    Experiments show that Steepest Guidance consistently outperforms existing approaches and handles nonlinear reward functionals.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The training-free approach directly optimizes reward locally, simplifying guidance estimation and accommodating general reward functionals.

    Source: §6 · p. 9

    The KL-regularized improvement guarantee requires guidance strength λ to satisfy 0 < λ ≤ 2η, so λ cannot be chosen independently of η.

    Source: §4 · p. 7

Abstract

from arXiv · show

Inference-time alignment of flow and diffusion-based models is critical for achieving flexible generative modeling. Theoretically, Doob's $h$-transform provides an elegant solution to this problem, and most existing methods are based on this principle. However, in practice, estimating the optimal guidance derived from Doob's $h$-transform at inference time is challenging. To deal with this issue, we regard inference-time alignment as a sequential optimization problem in the space of probability measures and propose a novel framework called *Steepest Guidance*, based on the principle of maximizing local improvement in the objective. We provide a theoretical analysis of the proposed method and demonstrate its effectiveness through extensive experiments.

1 INTRODUCTION

The introduction frames inference-time alignment as reward optimization over generated distributions, but finite-sample guidance estimation and nonlinear objectives remain challenging. Steepest Guidance treats alignment as sequential optimization in probability-measure space and reports theoretical and experimental support for the approach.

1 INTRODUCTION

  • Reward-guided generation optimizes a functional over the sample distribution; examples beyond expected reward include diversity-promoting and risk-sensitive objectives.The paper describes Doob’s h-transform as a theoretical optimal-control solution and cites guidance-model approaches [Domingo-Enrich et al., 2024; Uehara et al., 2024].Source: §1 · p. 1
  • Existing inference-time methods face limited support for nonlinear rewards, while finite-sample plug-in and REINFORCE guidance estimates can be biased.The paper notes that the plug-in estimator also requires differentiability; it reports large REINFORCE bias in a later toy problem.Source: §1 · p. 2
  • Steepest Guidance optimizes local reward improvement sequentially in probability-measure space, with unbiased estimation for linear rewards and particle-based application to nonlinear rewards.The contributions also include an entropy-regularized variant, improvement guarantees, global convergence under suitable assumptions, and experiments reporting consistent outperformance and nonlinear-reward handling.Source: §1 · p. 2

1.1 RELATED WORK

  • Related work spans selection-based and guidance-based inference-time methods, as well as trained probability-measure optimization and Langevin approaches.Guidance methods may require costly backpropagation or suffer bias; the cited Langevin approaches face convergence difficulty on complex multimodal real-world distributions, while SALD cannot handle nonlinear reward functionals.Source: §1.1 · p. 2

2 PRELIMINARIES

The preliminaries formulate flow and diffusion generation as stochastic processes and alignment as reward-functional optimization with guidance. They review Doob’s h-transform and explain why its optimal guidance is difficult to estimate in practice.

2.1 FLOW AND DIFFUSION-BASED MODELS

  • The paper models generation with flow and diffusion processes, including an SDE initialized from πε; theoretical analysis assumes ε > 0 to avoid a noise-schedule blow-up at t = 0.The model passage states that the SDE’s marginal distributions match those of the flow-matching ODE for the specified construction.Source: §2 · p. 2, §2.1 · p. 3

2.2 INFERENCE-TIME ALIGNMENT OF FLOW MODELS

  • Inference-time alignment adds a guidance term to the generative process to maximize a reward functional over generated samples.The objective can include a KL regularizer, and examples of nonlinear functionals include CVaR and Rao’s quadratic entropy.Source: §2.2 · p. 3

2.3 OPTIMAL CONTROL VIA DOOB’S h-TRANSFORM

  • Doob’s h-transform provides the principled optimal guidance, but for general nonlinear rewards it requires the unknown optimal distribution and has no closed-form expression.The paper describes inference-time Monte Carlo estimation from the conditional distribution as one response to the lack of a closed form.Source: §2.3 · pp. 4–5
  • The plug-in estimator requires differentiable rewards and sampling, whereas REINFORCE only requires differentiability of the conditional log-likelihood.The conditional log-likelihood gradient can be computed from the marginal score function; the plug-in estimator also requires backpropagation through the generative process.Source: §2.3 · p. 4
  • For a toy problem with a one-dimensional Gaussian terminal distribution and linear reward r(y)=y, the plug-in estimator is unbiased, while REINFORCE can still exhibit large bias.Source: §2.3 · p. 5

3 PROPOSED METHOD: STEEPEST GUIDANCE

Steepest Guidance replaces difficult estimation of globally optimal guidance with a locally optimal direction of reward improvement. Its practical estimators differ by reward type: the linear-reward estimator is unbiased under conditional sampling, while finite-sample nonlinear estimation is generally biased.

3 PROPOSED METHOD: STEEPEST GUIDANCE

  • Steepest Guidance avoids estimating globally optimal guidance by choosing a local reward-improvement direction, extending the generator-based argument to general reward functionals.The functional derivative appears directly in the expectation without a nonlinear transformation, unlike Doob’s h-transform.Source: §3 · p. 5

3.1 LOCAL IMPROVEMENT OF THE REWARD FUNCTIONAL

  • For a small time interval, maximizing the quadratic local improvement yields the steepest guidance, characterized by local optimality of reward improvement.The guidance is applied over a short interval before the process resumes its unguided dynamics.Source: §3.1 · p. 5
  • Under admissibility, steepest guidance provably improves the reward functional while controlling KL divergence between the original and guided processes.The paper interprets this result as a functional version of the descent lemma; larger λ is expected to yield greater reward improvement.Source: §3.1 · p. 6

3.2 PRACTICAL IMPLEMENTATION OF STEEPEST GUIDANCE

  • For linear rewards, conditional sampling gives an unbiased estimator, whereas finite-sample nonlinear-reward estimation is generally biased.Approximate conditional sampling can bias the linear-reward estimator, though the paper reports it can still be effective in practice.Source: §3.2 · p. 6

4 REGULARIZED STEEPEST GUIDANCE

Regularized Steepest Guidance constructs guidance for a terminal KL-regularized objective using a lower-bound surrogate when direct variation requires an intractable density ratio. The paper proves objective improvement, including for λ > 2η, and establishes global convergence for concave rewards under stated assumptions.

4 REGULARIZED STEEPEST GUIDANCE

  • The earlier improvement guarantee applies only when 0 < λ ≤ 2η, tying guidance strength to the KL regularization parameter.The later theorem for regularized steepest guidance guarantees improvement even when λ > 2η.Source: §4 · p. 7
  • Directly differentiating the terminal KL-regularized objective requires a generally intractable density ratio, so the method uses a lower-bound surrogate to construct guidance.The surrogate is lower-bounded by the terminal objective via the data-processing inequality; an equivalent-marginal SDE avoids computing the density ratio.Source: §4 · p. 7
  • Regularized Steepest Guidance improves the KL-regularized objective even when λ > 2η.This extends beyond the earlier guarantee that required 0 < λ ≤ 2η.Source: §4 · p. 7

4.1 GLOBAL CONVERGENCE FOR CONCAVE REWARD FUNCTIONALS

  • For concave rewards, global convergence is established under a uniform log-Sobolev inequality and additional regularity conditions.The analysis handles a time-dependent, distribution-dependent objective, extending the convergence analysis discussed for Langevin dynamics.Source: §4.1 · pp. 7–8

5 NUMERICAL EXPERIMENTS

The experiments evaluate Steepest Guidance on synthetic and image-generation tasks. They report baseline comparisons, a reward–quality trade-off for regularized guidance, and results for nonlinear CVaR objectives.

5.2 IMAGE GENERATION

  • With α = 0.5, the proposed method achieves better CVaR when maximizing CVaR than when maximizing expected reward.Figure 4 compares the reward distributions; its horizontal lines indicate lower-tail CVaR.Source: Figure 4 · p. 9, §5.2 · p. 9
  • Regularized Steepest Guidance achieves better Blueness–ImageReward trade-offs than vanilla guidance as η controls regularization strength.ImageReward is used as a proxy for the otherwise unmeasurable terminal KL divergence.Source: Figure 3 · p. 9, §5.2 · p. 9

5.1 TOY EXPERIMENTS

  • In the 1D Gaussian-mixture task, empirical Doob sometimes misses the positive mode; steepest guidance selects it but misses the optimum, while regularized guidance approximates it.The task uses a step reward favoring the positive region.Source: §5.1 · pp. 8–9

5.2 IMAGE GENERATION

  • Across reward functions and models, Steepest Guidance outperforms empirical Doob’s transform and SVDD.Image experiments use Stable Diffusion v1.5 and FLUX.1, with Blueness, PickScore, ImageReward, and Compressibility rewards.Source: §5.2 · p. 9

6 CONCLUSION

Steepest Guidance is presented as a training-free approach that locally optimizes reward functionals, with theoretical guarantees and experiments in text-to-image generation. The supplied appendices compare related approaches and describe reward-functional examples, the algorithm, and supporting analyses.

6 CONCLUSION

  • Steepest Guidance directly optimizes reward functionals locally, simplifies guidance estimation, handles general functionals, and has provable improvement guarantees for ideal dynamics plus global convergence under suitable assumptions.The conclusion reports experimental effectiveness in text-to-image generation tasks.Source: §6 · p. 9

A.1 SYSTEMATIC COMPARISON WITH PREVIOUS WORKS

  • Table 2 compares reward-guided methods on training-free use, derivative-free use, and support for non-linear reward functionals.The supplied caption identifies these comparison dimensions but does not provide cell values.Source: §A.1 · p. 12, Table 2 · p. 13

A.2 OTHER RELATED WORKS

  • Diversity-seeking methods use repulsive particle forces but do not cover general functionals; SALD targets tilted distributions, excludes general non-linear functionals, and may require slow annealing for its improvement guarantee.The passage says such slowdown can significantly affect generation efficiency.Source: §A.2 · pp. 12–13

B EXAMPLES OF REWARD FUNCTIONALS

  • The appendix presents reward-functional examples and notes that entropy’s first-order variation cannot be computed directly because it requires the density µ(x).It instead uses a relation that interprets entropy maximization as KL-regularized expected reward; when log π1(x) is unavailable, the score function may be used if sampling is differentiable.Source: B EXAMPLES OF REWARD FUNCTIONALS · p. 13

B.1 FIRST-ORDER VARIATION OF REWARD FUNCTIONALS

  • The appendix gives first-order variations for CVaR and Rao’s quadratic entropy, which are used in the experiments.The CVaR statement assumes a strictly positive density and a continuous reward density near its α-quantile; Rao’s functional uses a symmetric positive definite kernel.Source: §B.1 · pp. 13, 15

C DETAILED ALGORITHM OF STEEPEST GUIDANCE

  • The detailed algorithm chooses a time grid, initializes particles from the unguided base process, estimates guidance at each step, updates particles, and returns the final samples.The excerpt specifies M particles and k samples in its sampling steps.Source: C DETAILED ALGORITHM OF STEEPEST… · p. 15

D AUXILIARY LEMMAS

  • Appendix bounds the optimality gap under concavity and finite first-order-variation assumptions, treating linear-flow and diffusion interpolations separately.The supplied passages identify these settings but do not expose the full bound expressions.Source: D AUXILIARY LEMMAS · pp. 18–19

E PROOF OF LEMMA 2.2

  • For the score decomposition, the marginal score is estimated by a neural network for diffusion models and can be computed from the velocity field for flow models.The derivation says only the conditional score term then needs to be computed.Source: E PROOF OF LEMMA 2.2 · p. 20

F PROOF OF PROPOSITION 3.1

The proof combines a reward-improvement argument with a relative-entropy identity for the guided process.

F PROOF OF PROPOSITION 3.1

  • The proof applies the chain rule for functions of measures and integration by parts to the guided SDE’s Fokker–Planck equation.The supplied excerpt gives the proof step but not the resulting expression.Source: F PROOF OF PROPOSITION 3.1 · p. 21
  • Integrating from ε to 1 and using Kεπε = π1 completes the reward-improvement argument..Source: F PROOF OF PROPOSITION 3.1 · p. 21
  • Because the base and guided processes share initial marginal πε, Girsanov’s theorem and admissibility conditions yield the KL identity.Source: F PROOF OF PROPOSITION 3.1 · p. 21

G PROOF OF THEOREM 3.2

The supplied proof excerpt reports that its conclusion follows from the preceding result.

G PROOF OF THEOREM 3.2

H PROOF OF COROLLARY 4.1

The excerpts state that the corollary proof follows from Theorem 3.2 and that the proposition proof follows from Lemma D.1.

H PROOF OF COROLLARY 4.1

I PROOF OF PROPOSITION 4.2

J PROOF OF THEOREM 4.3

The proof analyzes the guided SDE using the Kolmogorov backward equation and the specified guidance term.

J PROOF OF THEOREM 4.3

  • The proof uses the Kolmogorov backward equation for V(t, µ) = R[K_tµ] and analyzes the guided SDE with g_t(y) = λσ²…The supplied passages do not include the completed guidance expression or the subsequent derivation.Source: J PROOF OF THEOREM 4.3 · p. 22

K PROOF OF THEOREM 4.5

The proof bounds the terminal objective through the guided process, while the experimental estimators use centering to reduce variance.

K PROOF OF THEOREM 4.5

  • For linear flow interpolation, Δε = O(ε²), whereas for diffusion interpolation, Δε = O(ε).The proof assumes π₁ has a finite second moment and the first-order variation of R is bounded.Source: K PROOF OF THEOREM 4.5 · p. 23
  • The proof applies the Fokker–Planck equation to the marginally equivalent ideal SDE and uses concavity of R to establish concavity of V(t, ·).The argument then applies Jensen’s inequality to a probability measure over the initial time and [ε, 1].Source: K PROOF OF THEOREM 4.5 · pp. 23–24
  • The proof combines an initialization-error term with an integrated guidance term, then uses K_1 as the identity to recover Jη.For linear flow interpolation, the initialization term is O(ε²).Source: K PROOF OF THEOREM 4.5 · p. 24

L.1 VARIANCE REDUCTION

  • Experiments center coefficients multiplying the conditional score; the leave-one-out baseline preserves unbiasedness for linear rewards, while nonlinear finite-particle estimates generally remain biased.For DOIT, subtracting the softmax coefficients’ empirical mean 1/k removes the constant component and reduces variance without changing the target expectation.Source: §L.1 · p. 25

L.2 EXPERIMENTAL SETUP FOR BIAS EVALUATION.

The appendix describes a Gaussian toy bias evaluation and text-to-image experiments, then reports additional reward, diversity, guidance-scaling, SALD, and particle-count comparisons.

L.2 EXPERIMENTAL SETUP FOR BIAS EVALUATION

  • In the Gaussian linear-reward setting, the plug-in estimator has no bias; the evaluation uses X₁ ∼ N(0, 1), r(x) = x, and λ = 5.The setup evaluates guidance at t = 0.5, with repeated REINFORCE estimates across sampled states.Source: §L.2 · pp. 25–26

L.3 REWARD FUNCTIONS

  • ImageReward and PickScore predict image scores from human preferences, while Compressibility is a non-differentiable JPEG-size reward.The experiments also define a Blueness reward from image pixel values.Source: §L.3 · p. 26

L.4 DETAILED SETUP FOR TEXT-TO-IMAGE EXPERIMENTS

  • Text-to-image experiments use 50 FLUX steps and 160 SD steps, with reward-specific guidance schedules and different approximate posterior sampling procedures.FLUX uses the Diamond Map; SD uses Tweedie’s formula for one-step denoising.Source: §L.4 · p. 26

L.5 ADDITIONAL RESULTS ON TEXT-TO-IMAGE GENERATION

  • The proposed method achieves better or comparable performance across two additional text-to-image prompts.The tables summarize rewards for the Tokyo-street and mountain-lake prompts over 32 generated samples each.Source: §L.5 · p. 26, Table 5 · p. 29, Table 6 · p. 30

L.6 EXAMPLES OF GENERATED IMAGES

  • The image examples show SD outputs for Blueness, ImageReward, PickScore, and Compressibility, ordered from top to bottom.The accompanying examples passage also includes FLUX images for these rewards.Source: §L.6 · p. 26, Figure 5 · p. 27, Figure 6 · p. 28

L.7 ADDITIONAL RESULTS ON NON-LINEAR REWARD FUNCTIONALS

  • For the “VAN GOGH CAFE TERASSE” prompt, regularized sampling produces visually varied images, unlike the similar scene compositions observed in unguided samples.The comparison uses Rao’s quadratic entropy and entropy regularization; diversity is measured by mean pairwise cosine distance between normalized DINOv2 embeddings.Source: §L.7 · p. 29, Table 7 · p. 30, Figure 7 · p. 31

L.8 EFFECT OF GUIDANCE SCALING

  • Guidance scaling improves empirical Doob performance to some extent but does not outperform steepest guidance; scaling may amplify guidance bias.The authors infer that insufficient guidance strength alone cannot explain the main-experiment performance gap.Source: §L.8 · p. 30, Figure 8 · p. 31

L.9 COMPARISON WITH SALD; L.10 ABLATION STUDY ON k

  • SALD does not match steepest guidance, particularly on ImageReward; the authors suggest Gaussian smoothing may remove fine-grained features relevant to that reward.As the number of lookahead particles increases, performance improves, while steepest guidance outperforms baselines at every tested k.Source: §L.9 · p. 30, §L.10 · pp. 30–31, Figure 10 · p. 32

Related papers

Loading 2609.39091v1…