VGG-Flow: Value Gradient Guidance for Flow

• The paper introduces a novel gradient-matching technique that aligns pretrained flow models with human preferences via an optimal control framework.
• It leverages heuristic initialization and finite-difference approximations to enable fast adaptation while maintaining sample diversity and prior fidelity.
• Empirical results on Stable Diffusion 3 demonstrate improved reward metrics and preserved diversity compared to baseline methods.

VGG-Flow (Value Gradient Guidance for Flow Matching Alignment) is a gradient-matching-based optimal control method designed to align flow matching generative models with human preferences while preserving adaptation efficiency and prior fidelity, as articulated in (Liu et al., 4 Dec 2025). The central innovation of VGG-Flow is leveraging the optimal difference between a finetuned and the base velocity field, matching it directly with the gradient of a learned value function. Empirical validations showcase its efficacy for rapid, robust finetuning in text-to-image generation, specifically on Stable Diffusion 3, under restricted computational budgets and diverse reward models.

1. Flow-Matching Models and the Alignment Challenge

Flow matching generative models are built around a time-indexed velocity field $v(x, t)$ that determines the trajectory from an initial noise sample $x_0 \sim \mathcal{N}(0, I)$ to a final sample $x_1$ in data space by integrating the ODE

$\frac{dx}{dt} = v(x, t).$

Training employs an $L_2$ objective in matching $v_\theta$ to a reference velocity $u(x_t|x_1) = (1-t)x_t + t x_1$ over $t \in [0,1]$ and $x_1 \sim \mathcal{D}$. Distinct from SDE-based diffusion models, flow matching uses deterministic ODE sampling, thereby necessitating novel alignment approaches. The alignment task seeks to finetune pretrained flows so model outputs maximize a learned human-preference reward $r(x)$, while simultaneously preserving sample diversity and base-model priors.

2. Optimal Control Formulation

Alignment in VGG-Flow is formulated as a deterministic optimal control problem. Given $v_0(x, t)$ (pretrained velocity field), the target is to optimize a new field $v_1(x, t)$ balancing proximity to $v_0$ and maximizing terminal reward $r(x_1)$. The residual control field is defined as

$\tilde{v}(x, t) := v_1(x, t) - v_0(x, t),$

with the expected cost

$$J[\tilde{v}] = \mathbb{E}_{x_0 \sim \mathcal{N},\, \dot{x}_t = v_0 + \tilde{v}} \left[ \frac{\lambda}{2} \int_0^1 \|\tilde{v}(x_t, t)\|^2 dt - r(x_1) \right],$$

where $\lambda > 0$ regulates the heating (temperature) of the control. Introducing the value function,

$$V(x, t) = \min_{\tilde{v}(\cdot)} \mathbb{E} \left[ \frac{\lambda}{2} \int_t^1 \|\tilde{v}(x_s, s)\|^2 ds - r(x_1) \mid x_t = x \right ],$$

establishes the foundation for deriving the optimal control law.

3. Hamilton–Jacobi–Bellman Derivation and Gradient-Matching Objective

The core control problem leads to the Hamilton–Jacobi–Bellman PDE:

$$-\partial_t V = \min_{\tilde{v}} \left[ \frac{\lambda}{2}\|\tilde{v}\|^2 + \nabla_x V^\top (v_0 + \tilde{v}) \right].$$

Solving the inner minimization yields the optimal residual

$$\tilde{v}^*(x, t) = -\frac{1}{\lambda} \nabla_x V(x, t).$$

This interprets the optimal modification to the pretrained velocity as the (negative) value gradient. Practical implementation parameterizes $V$ via $\psi$, enforcing:

(a) Gradient-Matching Control Law

$$L_{\text{match}}(\theta, \psi) = \mathbb{E}_{x_0, t}\big[ \|(v_1(x_t, t) - v_0(x_t, t)) + \beta \cdot g_\psi(x_t, t)\|^2 \big],$$

where $g_\psi(x, t) \approx \nabla_x V(x, t)$ and $\beta = 1/\lambda$.

(b) Value-Consistency (Bellman Residual)

With $g_\psi$ also required to satisfy the Bellman gradient PDE,

$$L_{\text{cons}}(\psi) = \mathbb{E}\big[ (\partial_t g_\psi + [\nabla g_\psi]^\top (-\beta g_\psi) - [\nabla v_0]^\top g_\psi)^2 \big],$$

plus terminal consistency

$$L_{\text{bdry}}(\psi) = \mathbb{E}\big[\|g_\psi(x_1, 1) + \nabla r(x_1)\|^2\big].$$

4. Heuristic Initialization and Fast Adaptation

Directly optimizing for $g_\psi$ under full value-consistency is inefficient. VGG-Flow introduces a heuristic by initializing $g_\psi$ using first-order reward gradients from a one-step Euler look-ahead:

$$\hat{x}_1(x_t, t) = x_t + (1-t)\,\text{stopgrad}[v_0(x_t, t)],$$

$$g_\psi(x_t, t) = -\eta_t\,\text{stopgrad}[\nabla_x r(\hat{x}_1)] + \nu_\psi(x_t, t),$$

with a scheduled coefficient $\eta_t$ (e.g., $\eta_t = t^2$) and small learnable correction $\nu_\psi$. This initialization utilizes actual reward gradients for rapid adaptation, letting $\nu_\psi$ refine $g_\psi$ to satisfy PDE consistency.

5. Algorithmic Workflow

An explicit sequence for VGG-Flow is as follows:

Algorithm VGG-Flow Input: pretrained flow v₀, reward r, initialize θ ← small LoRA on v₀, initialize gϕ via heuristic. repeat until convergence or budget exhausted: 1. Sample batch of trajectories {xₜ} via dx/dt = vθ(x,t) from t=0 to 1. 2. Update value-gradient net ϕ by minimizing L_cons(ϕ) + α·L_bdry(ϕ). 3. Update flow-model net θ by minimizing L_match(θ, ϕ). end

Significant efficiency is gained by avoiding backward adjoint ODE solves (unlike adjoint-matching), employing finite differences and Jacobian-vector products for PDE terms, and disabling all second-order autograd. Trajectory subsampling further reduces computational cost.

6. Empirical Performance, Prior Preservation, and Diversity

VGG-Flow was evaluated on Stable Diffusion 3 (20-step Euler sampler, LoRA rank 8) with Aesthetic Score, HPSv2, and PickScore as reward metrics, a budget of 400 update steps, batch size 32, and 3 random seeds. Key quantitative results for Aesthetic Score (400 steps) are shown below:

Method	Reward↑	Diversity↑ (10⁻²)	FID↓
Base	5.99	23.12	212
ReFL	10.00	5.59	1338
DRaFT	9.54	7.78	1518
Adjoint M.	6.87	22.34	465
VGG-Flow	8.24	22.12	375

VGG-Flow achieves higher reward than baseline flow methods, with diversity and FID comparable to the base model. Pareto analysis confirms that at fixed reward levels, VGG-Flow maintains significantly higher diversity and lower FID compared to alternatives. Qualitatively, the method preserves semantic fidelity without inducing mode collapse, unlike ReFL and DRaFT which display overfitting artifacts.

7. Limitations and Prospects

The objective strictly matches the KL-regularized optimum only when $\lambda$ is small; at larger $\lambda$, the approximation deteriorates. Finite-difference approximations for the PDE contribute bias and are sensitive to step-size choices. There is inherent tension between exploration and exploitation, with the risk of missing high-reward modes or sample collapse under constrained updates. The method does not explicitly manage the divergence term of full KL regularization, instead relying on neural network implicit bias. Future directions identified include more robust second-order or adjoint-free solvers for the HJB term, richer value function modeling (potentially via implicit neural representations), extension to stochastic flows in conjunction with SDEs, adaptive $\lambda$ scheduling, and generalization to other modalities such as video, 3D data, or hybrid architectures.

The complete technical specification and results for VGG-Flow are presented in "Value Gradient Guidance for Flow Matching Alignment" (Liu et al., 4 Dec 2025).