D-Flow: Differentiating through Flows for Controlled Generation

1. Introduction

• Controlled-generation approaches include task-specific conditional training, training-free guidance, and variational optimization using differentiable costs.The introduction notes that conditional training may be prohibitive and that guidance approaches often rely on assumptions that limit their use.Source: §1 · p. 1
• D-Flow formulates controlled generation as optimizing the source noise point x0 against an arbitrary control cost through the generation ODE.The generated sample is constrained to be the ODE solution, while x0 is the optimized quantity.Source: §1 · pp. 1–2
• The framework applies differentiation through generation to general Diffusion and Flow-Matching models and a wider class of problems defined by general cost functions.The paper reports theoretical and practical evidence for this broader applicability.Source: §1 · p. 2
• D-Flow achieves state-of-the-art performance across image inverse problems, conditional molecule generation, and audio inpainting and super-resolution without careful tuning across domains.The reported applications use image, molecule, and audio generative priors.Source: §1 · p. 2
• Generation is relatively slow compared with some baselines, usually taking 5–15 minutes on ImageNet-128 using an NVIDIA V100 GPU, though the authors believe there is room for speed improvement.Source: §1 · p. 2
• For models trained with Gaussian probability paths, differentiating with respect to the source point projects the cost gradient onto major data-manifold directions.The paper describes this as an implicit prior in the optimization.Source: §1 · pp. 1–2

2. Preliminaries

• A flow model first samples an initial point x(0) from a source distribution p0(x0).The paper describes the source as a noise distribution.Source: §2 · p. 2
• The generated endpoint x(1) follows the distribution p1 induced by drawing x(0) from p0 and solving the ODE.The paper denotes p1 as the resulting distribution and density function.Source: §2 · p. 2
• The model then solves an ODE from t = 0 to t = 1 using a predetermined velocity field u.The velocity field maps time and state to the state’s evolution direction.Source: §2 · p. 2

3. Controlled Generation via Source Point Optimization

• D-Flow optimizes only the initial source point x0, constraining the output sample to follow the flow ODE while minimizing the desired cost.The framework seeks samples that have low cost and are likely under the flow model's distribution.Source: §3 · p. 2

3.1. Cost Functions

• Cost functions cover target matching, inverse problems such as inpainting, deblurring, and super-resolution, and conditional sampling guided by a classifier or energy function.For inverse problems, the loss compares the corrupted output H(x) with the observation; target matching seeks an ODE trajectory that reaches y.Source: §3.1 · pp. 2–3

3.2. Initialization

• When observations strongly constrain the desired sample, initializing x0 with a blend of source noise and the observation's backward ODE solution can speed convergence.The alternative natural initialization is a sample from the source distribution.Source: §3.2 · p. 3

3.3. Regularizations

• Target negative log likelihood can be added as a regularizer, but its high-dimensional suitability is unclear and it requires an ODE divergence term.The paper notes that a zero-masked ImageNet-128 test image can be more likely under the flow model than the original image.Source: §3.3 · p. 3
• Optimizing through x0 projects the output-cost gradient onto high-variance data directions, helping optimization stay close to the data distribution.The paper illustrates this with optimization steps that pass through different bird subspecies toward a target image.Source: §3.3 · p. 4

3.4. Practical Implementation

• The implementation uses differentiable ODE solving with gradient checkpointing and LBFGS with line search; typical optimization takes 5–15 minutes per sample.Most experiments use the midpoint solver with six function evaluations; large text-to-image and text-to-audio models can take 30–40 minutes.Source: §3.4 · p. 4

• The theory derives source-point gradients under AGPP and interprets the resulting sample velocity as pointing in the direction of the data distribution.It uses adjoint dynamics to obtain the gradient and its asymptotic effect on the generated sample.Source: §4 · p. 4
• AGPP interpolates Gaussian source noise and the data distribution, with scheduler conditions ensuring the endpoints correspond to source and target distributions.The source is p0 and the target is p1; the schedules satisfy α0 = 0, σ1 ≈ 0, α1 = 1, and σ0 = 1.Source: §4 · p. 4
• Under AGPP, the denoiser’s gradient with respect to its input is proportional to the conditional variance of the target given that input.Source: §4 · p. 5
• In continuous time, the differential of the generated sample with respect to its source is a time-ordered exponential of a bounded time-dependent matrix.Its matrix-vector product is an infinite sum involving powers of γtVar1|t(x(t)).Source: §4 · p. 5
• For an Euler solver, the discrete differential applies covariance matrices at successive times, projecting vectors onto major axes of the data conditioned on the current sample.This offers an intuition for why differentiating through a solver with few steps can still work.Source: §4 · p. 5
• The resulting change in x(1) follows the loss gradient projected by the flow differential, which iteratively applies covariance-based projections at different times.Source: §4 · p. 5

• Inverse-problem methods modify diffusion sampling through variational solvers, observation-consistency guidance, or derivatives through denoising steps.Related work also extends guidance to latent diffusion and uses reconstruction loss with score-matching regularization.Source: §5 · pp. 5–6
• Conditional sampling methods include task-specific condition training and plug-and-play optimization using pre-trained unconditional generative models.Prior optimization approaches include reconstruction objectives, learned controls, and methods for GANs and normalizing flows.Source: §5 · p. 6
• D-Flow’s analysis concerns gradients with respect to initial values in diffusion and flow models with affine Gaussian paths, including behavior with few solver steps.The passage contrasts this analysis with prior work on simplified linear models, manifold preservation, and discrete-diffusion differentiation.Source: §5 · p. 6

6. Experiments

• For inverse problems with informative observations, blend initialization can speed convergence and often improve performance; explicit regularization helps only in noisy linear inverse problems and molecule generation.The experiments use a blend of a source-distribution sample and a backward solution for image inverse problems.Source: §6 · p. 6

6.1. Linear Inverse Problems on Images

• For image inpainting and super-resolution, D-Flow improves on state-of-the-art methods in most metrics and performs well in noisy cases.The authors attribute higher fidelity to source-point optimization correcting the sampling trajectory toward the observed signal.Source: Table 1 · p. 7, §6.1 · p. 7

6.2.1. IMAGE INPAINTING

• For free-form image completion, RED-Diff performs better on structural metrics, while D-Flow produces more semantically plausible completions on perceptual metrics.The authors also report that RED-Diff often produces artifacts on this task.Source: Table 2 · p. 7, §6.2.1 · p. 7

6.2.2. AUDIO INPAINTING AND SUPER-RESOLUTION

• Across audio inpainting and super-resolution experiments, D-Flow achieves the lowest FAD, with slightly lower PSNR than RED-Diff for inpainting.The evaluation includes signal crops of 10% and 20%, and upscaling by factors of two, four, and eight.Source: Table 3 · p. 8, §6.2.2 · p. 8

6.3. Conditional Molecule Generation on QM9

• Molecule-generation baselines use separately trained property-specific conditional models, whereas D-Flow uses one unconditional model.The QM9 setup uses molecules with up to 29 atoms and compares against conditional EDM, EQUIFM, and GEOLDM.Source: §6.3 · pp. 8–9
• D-Flow significantly outperforms conditional-generation baselines in molecule-generation quality, while its generated molecules have lower stability than conditional EDM’s.The authors attribute the stability gap to the unconditional flow model’s lower stability and optimization with the property predictor.Source: Table 4 · p. 8, Table 5 · p. 8, §6.3 · p. 9
• D-Flow’s molecule-generation stability is below conditional EDM’s 82.1% average: the unconditional flow prior reaches 72.2% stability.The authors also report comparable MAE improvements for stable and non-stable generated molecules.Source: Table 5 · p. 8, §6.3 · p. 9

7. Discussion, Limitations and Future Work

• D-Flow’s main limitation is relatively long runtime because it back-propagates through multiple compositions of the velocity field.The authors identify reducing computational overhead while retaining the implicit bias as a future direction.Source: §7 · p. 9

Impact Statement

• The authors state that they see no societal consequences requiring specific emphasis.They describe the work’s goal as advancing machine learning.Source: Impact Statement · p. 9

A.1. Proof of Proposition 4.1

• For AGPP, the denoiser gradient with respect to the input is proportional to the variance of the conditional distribution over data points.The proof introduces affine Gaussian paths and their conditional and marginal velocity fields.Source: §A.1 · p. 12

A.2. Proof of Theorem 4.2

• The proof computes the output’s differential with respect to the initial point using adjoint dynamics, whose linear ODE solution uses a time-ordered exponential.The cited passages describe solving the adjoint ODE backward from time 1 to time 0.Source: §A.2 · pp. 13–14

A.3. Discrete Time Analysis

• The discrete-time analysis differentiates the final Euler-solver state with respect to the initial point using the chain rule.For the CondOT path, the analysis specifies α_t = t and σ_t = 1 − t.Source: §A.3 · p. 15

A.4. On Flow-Matching, Denoisers and Noise Prediction

• The appendix describes affine conditional probability paths and expresses their marginal velocity field in terms of either a denoiser or a noise predictor.The path parameters α_t and σ_t can represent known diffusion and flow-matching paths.Source: §A.4 · p. 16

B.1. Linear Inverse Problems on Images

• ImageNet-128 linear inverse problems use χd source-point regularization for noisy tasks, with algorithmic choices listed in Table 6.The χd regularization corresponds to equation 10; λ denotes its coefficient.Source: §B.1 · p. 17, Table 6 · p. 17
• Noiseless inpainting, super-resolution, and Gaussian deblurring take 10, 12.5, and 15.5 minutes per image; noisy counterparts take 4, 2.5, and 3.5 minutes.Experiments ran on a 32GB NVIDIA V100 GPU; RED-Diff parameters were searched, retaining prior results when the search did not outperform them.Source: §B.1 · p. 17

B.2.1. IMAGE INPAINTING

• For latent-space RED-Diff, the corruption function is modeled as f = H(decode(µ)), with learning rate 0.25 and λ = 0.25.The decoder maps from the autoencoder latent space before applying the corruption operator H.Source: §B.2.1 · p. 17

B.2.2. AUDIO INPAINTING

• Audio inpainting uses 10 inner iterations, stops after 100 global iterations, and initializes the algorithm with a backward blend using α = 0.1.The RED-Diff baseline uses lr = 0.05 and λ = 0.5.Source: §B.2.2 · p. 17

B.3. Conditional Molecule Generation on QM9

• QM9 generation initializes Gaussian atom features, normalizes them to zero mean and unit standard deviation after each step, and uses 100 midpoint function evaluations.LBFGS uses five optimization steps with up to five inner iterations; QM9 contains molecules with up to nine heavy atoms.Source: §B.3 · p. 18
• D-Flow’s MAE remains state of the art on generated stable molecules, so its improved MAE is not attributable only to non-stable outputs.Table 8 reports MAE over 10K generated samples and separates stable from non-stable molecules because D-Flow produces a lower stability percentage.Source: §B.3 · p. 18, Table 8 · p. 18

C.1. Linear Inverse Problems on Images: Denoising

• The denoising experiment adds Gaussian noise with σ_y = 0.05 and reports evaluation in Table 9.The noisy-experiment hyperparameters follow Table 6.Source: §C.1 · p. 19, Table 9 · p. 19

C.2. Ablation: Regularization Coefficient

• The ablation varies λ, the χd regularization coefficient, for noisy super-resolution on ace-blurred ImageNet-128.Figure 7 plots evaluation metrics against λ.Source: §C.2 · p. 19, Figure 7 · p. 19
• The plotted metrics are FID, LPIPS, PSNR, and SSIM for noisy super-resolution.The experiment uses the other hyperparameters reported in Table 6.Source: §C.2 · p. 19, Figure 7 · p. 19
• The regularized source point is χd, corresponding to equation 10, and λ is its coefficient.The figure caption identifies the task as noisy super-resolution on ace-blurred ImageNet-128.Source: §C.2 · p. 19, Figure 7 · p. 19

D. Additional Qualitative Results

• Figure 8 presents qualitative results for noiseless linear inverse problems on ImageNet-128, with face-blurred validation images as ground truth.Source: Figure 8 · p. 20
• Figure 9 presents qualitative results for noisy linear inverse problems on ImageNet-128, with face-blurred validation images as ground truth.Source: Figure 9 · p. 21
• Figure 10 compares free-form inpainting on MS-COCO using a text-to-image latent flow-matching model, against MS-COCO validation ground truth.Source: Figure 10 · p. 22