[Paper] Value Gradient Guidance for Flow Matching Alignment

Source: arXiv - 2512.05116v1

Overview

The paper introduces VGG‑Flow, a new way to fine‑tune large flow‑matching generative models (e.g., Stable Diffusion 3) so they obey human‑defined preferences without sacrificing the model’s original knowledge. By framing alignment as an optimal‑control problem, the authors achieve fast, compute‑efficient adaptation while keeping the pretrained “prior” intact.

Key Contributions

• Value‑gradient guidance: Shows that the optimal adjustment to a pretrained velocity field can be expressed as the gradient of a learned value function.
• Gradient‑matching finetuning: Proposes a simple, first‑order loss that aligns the model’s velocity field with the value‑gradient, avoiding expensive reinforcement‑learning loops.
• Heuristic value‑function initialization: Introduces a practical way to bootstrap the value function, dramatically speeding up convergence.
• Empirical validation on Stable Diffusion 3: Demonstrates that VGG‑Flow can align a state‑of‑the‑art text‑to‑image model under tight compute budgets while preserving visual quality and diversity.

Methodology

1. Background – Flow Matching

   • Flow‑matching models learn a velocity field (v_\theta(x,t)) that transports a simple noise distribution into the data distribution over continuous time.
   • Sampling is done by integrating this field (e.g., with an ODE solver).

2. Alignment Goal

   • We have a reward model (R(x)) that scores how well a generated sample matches a human preference (e.g., “more realistic” or “contains a cat”).
   • The ideal aligned model should generate samples that maximize expected reward and stay close to the original pretrained distribution.

3. Optimal‑Control Formulation

   • Treat the velocity adjustment (\Delta v(x,t)) as a control input.
   • The optimal control that maximizes expected reward while penalizing deviation from the pretrained field solves a Hamilton‑Jacobi‑Bellman (HJB) equation.

4. Value‑Gradient Guidance

   • The solution to the HJB yields a value function (V(x,t)) whose gradient (\nabla_x V) tells us the direction to modify the velocity field.
   • Instead of solving the full HJB, VGG‑Flow matches the pretrained velocity plus a learnable correction (\Delta v) to the gradient (\nabla_x V):

[ \min_{\phi}; \mathbb{E}{x,t}\bigl|,\Delta v\phi(x,t) - \nabla_x V_\psi(x,t),\bigr|^2 ]

• (\phi) are parameters of the correction network, (\psi) are parameters of the value network.

5. Heuristic Initialization

   • The value network is seeded with a simple proxy (e.g., the reward model’s logits passed through a shallow MLP) so that early gradients already point toward higher‑reward regions.

6. Training Loop

   • Sample trajectories from the pretrained model.
   • Compute rewards, update the value network via TD‑style regression, then update the correction network via the gradient‑matching loss.
   • No reinforcement‑learning roll‑outs or policy‑gradient estimators are required, keeping the compute cost low.

Results & Findings

• Alignment quality: VGG‑Flow consistently pushes samples toward higher reward regions, outperforming naïve fine‑tuning and RL‑based baselines.
• Prior preservation: Despite the shift in preferences, the Fréchet Inception Distance (FID) stays virtually unchanged, confirming that the original visual fidelity is retained.
• Efficiency: The gradient‑matching loss converges in a few hundred steps, thanks to the heuristic value initialization, making it feasible on a single workstation.

Practical Implications

• Fast, on‑device customization: Developers can adapt large text‑to‑image models to brand‑specific styles, safety filters, or user‑feedback loops without needing weeks of GPU time.
• Plug‑and‑play alignment: VGG‑Flow works as a thin wrapper around any pretrained flow‑matching model; you only need a reward model (often already available as a classifier or CLIP scorer).
• Reduced risk of “mode collapse”: Because the method penalizes deviation from the original velocity field, it avoids the common pitfall of RL‑based alignment that over‑optimizes for the reward at the cost of diversity.
• Potential for other modalities: The same optimal‑control view applies to audio, video, or 3‑D generative flows, opening a path for cross‑modal preference alignment.

Limitations & Future Work

• Reward model dependence: The quality of alignment hinges on the reliability of the external reward model; biased or noisy rewards will propagate to the generator.
• Heuristic value init: While effective, the current initialization is hand‑crafted; learning a more principled prior could further accelerate convergence.
• Scalability to extremely large models: Experiments were limited to Stable Diffusion 3 (≈ 3 B parameters). Extending VGG‑Flow to multi‑billion‑parameter diffusion pipelines may require additional memory‑efficient tricks.
• Theoretical guarantees: The paper provides an intuitive optimal‑control derivation but stops short of formal convergence proofs; future work could tighten the theoretical underpinnings.

Bottom line: VGG‑Flow offers a developer‑friendly recipe for aligning powerful flow‑matching generators with human preferences—fast, cheap, and with minimal impact on the model’s original capabilities. It’s a promising step toward making large generative models truly customizable in production environments.

Authors

• Zhen Liu
• Tim Z. Xiao
• Carles Domingo-Enrich
• Weiyang Liu
• Dinghuai Zhang

Paper Information

• arXiv ID: 2512.05116v1
• Categories: cs.LG, cs.CV
• Published: December 4, 2025
• PDF: Download PDF