Greed is Good: A Unifying Perspective on Guided Generation

Training-free guided generation is a widely used and powerful technique that allows the end user to exert further control over the generative process of flow/diffusion models. Generally speaking, two families of techniques have emerged for solving this problem for gradient-based guidance: namely, posterior guidance (i.e., guidance by projecting the current sample to the target distribution via the target prediction model) and end-to-end guidance (i.e., guidance by performing backpropagation throughout the entire ODE solve). In this work, we show that these two seemingly separate families can actually be unified by looking at the posterior guidance as a greedy strategy of end-to-end guidance. We explore the theoretical connections between these two families and provide an in-depth theoretical understanding of these two techniques relative to the continuous ideal gradients. Motivated by this analysis, we then show a method for interpolating between these two families enabling a trade-off between compute and accuracy of the guidance gradients. We then validate this work on several inverse image problems and property-guided molecular generation. scheme (cf . Equation ( 2)) is part of the computation graph of the model reverse-mode automatic differentiation (Linnainmaa 1976) is applied, i.e., vanilla backpropagation. The memory cost of such techniques, however, is O(n), prompting researchers to explore the second method known as optimize-then-discretize (OTD) which instead solves another ODE in reverse-time which models the continuous-time dynamics of reverse-mode differentiation, this is called the continuous adjoint method (R. T. Chen et al. 2018; cf . Kidger 2022, Section 5.1.2).

Given a flow model u θ ∈ C 1,1 ([0, 1] × R d ; R d ) that is Lipschitz continuous in its second argument and the solution x : [0, 1] → R d , x t → x(t), let a x := ∂L/∂x t denote the adjoint state. Then a x (t) can be found by solving the continuous adjoint equation:

N.B., this technique was first proposed by Pontryagin et al. (1963) and popularized for neural differential equations by R. T. Chen et al. (2018). This approach has a constant memory cost O(1); however, this comes with the cost of several drawbacks related to the numerical scheme. While these issues are not particularly relevant to our theoretical analyses, we note them in Appendix E for the ML practitioner.

Now returning back to our problem statement from Equation ( 3), the end-to-end guidance techniques amount to optimizing the initial condition x 0 in light of the entire solution trajectory admitted by the numerical scheme. A natural question we consider for problems of this form is that rather than finding the full sequence x n , can we make use of local information instead? I.e.,

Rather than solving the full ODE from x t , what if we greedily took a locally optimal step at each x t instead?