Value Gradient Guidance for Flow Matching Alignment

TL;DR

This work tackles efficient, prior-preserving alignment of flow matching models with human preferences by embedding the problem in an optimal-control framework. It introduces VGG-Flow, which uses value-gradient guidance from the Hamilton-Jacobi-Bellman equation to regularize the residual velocity between a finetuned model and its base, via a forward-looking value gradient model and a set of consistency, boundary, and matching losses. Empirically, VGG-Flow achieves faster reward convergence with better diversity and stronger prior preservation than baselines on Stable Diffusion 3 across multiple reward models and under constrained compute. The approach offers a principled, memory-efficient route to principled alignment and highlights avenues for further architectural and optimization refinements.

Abstract

While methods exist for aligning flow matching models--a popular and effective class of generative models--with human preferences, existing approaches fail to achieve both adaptation efficiency and probabilistically sound prior preservation. In this work, we leverage the theory of optimal control and propose VGG-Flow, a gradient-matching-based method for finetuning pretrained flow matching models. The key idea behind this algorithm is that the optimal difference between the finetuned velocity field and the pretrained one should be matched with the gradient field of a value function. This method not only incorporates first-order information from the reward model but also benefits from heuristic initialization of the value function to enable fast adaptation. Empirically, we show on a popular text-to-image flow matching model, Stable Diffusion 3, that our method can finetune flow matching models under limited computational budgets while achieving effective and prior-preserving alignment.

Figures (20)

• Figure 1: Left: Illustration of our proposed VGG-Flow algorithm. The velocity field is trained to match the value gradient field obtained from the optimal control problem. The value gradient field is parametrized as the reward gradient field of the one-step Euler prediction plus a learnable residual field. Right: Evolution of samples (with fixed seeds and prompts) during the course of finetuning on the reward model of Aesthetic Score.
• Figure 2: Comparison on samples generated by models finetuned with different methods. All models are finetuned with a maximum of $400$ update steps and for fair qualitative comparison we pick the model checkpoints that yield the best rewards without significant collapsing in image semantics (as ReFL and DRaFT are more prone to overfitting). For each set of images produced by each method, we display their average reward on the left.
• Figure 3: Qualitative results on HPSv2.
• Figure 4: Qualitative results on PickScore.
• Figure 5: Convergence curves of different metrics for different methods throughout the finetuning process on Aesthetic Score. Finetuning with our proposed VGG-Flow converges faster than the non-gradient-informed methods and with better diversity- and prior-preserving capability.

Theorems & Definitions (9)

• Remark 1: Connection to Equation \ref{['eqn:optimal-ctrl-objective']}
• Proposition 2: Basic adjoint matching for deterministic control
• proof
• Proposition 3: Adjoint matching for deterministic control
• proof
• Proposition 4: $W_2$ Bound via Grönwall's Inequality
• proof
• Proposition 5: KL Divergence Identity for ODEs
• Remark 6: Justification for the $L_2$ Proxy Objective