💻 computer science

Scaling Reinforcement Learning for Diffusion Models via Velocity Matching

This paper introduces Reward-based Velocity Matching (RVM), a simple and computationally efficient trajectory-free framework that directly optimizes the velocity field of diffusion models for reward fine-tuning, outperforming existing likelihood-based policy-gradient methods while offering a unified perspective on recent approaches and enabling improved dynamic rewards for video generation.

Original authors: Jaemoo Choi, Wei Guo, Yuchen Zhu, Arash Vahdat, Molei Tao, Julius Berner, Yongxin Chen

Published 2026-08-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Jaemoo Choi, Wei Guo, Yuchen Zhu, Arash Vahdat, Molei Tao, Julius Berner, Yongxin Chen

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

In the rapidly evolving world of artificial intelligence, a specific class of computer programs has emerged that can create stunning images and realistic videos from simple text descriptions. These systems, known as diffusion models, work by learning to reverse a process of gradual corruption. Imagine taking a clear photograph and slowly adding static noise until it becomes unrecognizable; these models learn the mathematical path to go backward, starting from pure noise and carefully removing it step by step to reveal a coherent picture. While these tools have become incredibly powerful, teaching them to follow human preferences—such as making a video move more naturally or a picture look more beautiful—has proven difficult. Traditionally, researchers have tried to adapt methods used for language models, which rely on calculating the probability of every possible next step in a sequence. However, because image and video generation involves millions of pixels changing simultaneously, calculating these probabilities is computationally expensive and often impossible to do with perfect accuracy.

A team of researchers from Georgia Tech and NVIDIA has found a simpler, more direct way to teach these models what humans prefer. Instead of trying to calculate complex probabilities for every tiny step of the generation process, they proposed a method that focuses directly on the "velocity" of the image as it forms. In the language of these models, the system predicts how the image should change from one moment to the next to reach the final result. The researchers discovered that they could simply nudge this prediction toward directions that lead to high-quality results and push it away from directions that lead to poor ones, using a reward signal as a guide. This approach, which they call reward-based velocity matching, bypasses the need for the complicated probability calculations that have slowed down previous attempts.

The researchers tested this idea on large-scale video generation models, which are particularly expensive to train because creating a single video requires significant computing power. They compared their new method against existing techniques that rely on tracking the entire path of the video's creation. The results were striking: their simpler method produced videos that were not only better in quality but also required a fraction of the computing time. In one specific test using a model called Wan2.1, their approach achieved the highest overall quality scores while using only about one-twelfth of the training time required by the previous best methods. This suggests that the complexity of the old methods was unnecessary; the key to improvement lay not in refining the mathematical machinery of probability, but in how the reward signals were designed and applied.

A critical part of their discovery involved understanding what the models were actually optimizing. The researchers found that standard rewards, which often focus on visual clarity or how well the video matches a text description, could inadvertently encourage the model to produce static, unmoving images that looked clean but were boring. To fix this, they introduced a new type of reward that specifically tracked motion, ensuring the video contained meaningful movement. When they added this motion-focused reward to their system, the videos became significantly more dynamic without losing their visual quality. This finding highlights that for these powerful models, the most important factor is not the complexity of the training algorithm, but the clarity and specificity of the goals set for the machine.

The study also revealed that the specific mathematical formula used to update the model matters less than previously thought. The researchers showed that their new method was mathematically equivalent to several other recent approaches, meaning that the differences in performance between those methods were likely due to how they set up their rewards rather than the core algorithm itself. By stripping away the unnecessary layers of probability estimation, they demonstrated that a direct, velocity-based update is sufficient to guide the model toward better results. This simplification opens the door for more efficient training, allowing researchers to iterate faster and potentially scale these technologies to even more complex tasks without being bottlenecked by computational costs.

Ultimately, this work suggests a shift in how we think about training generative AI. Rather than treating the problem as a complex puzzle of probability estimation, the researchers showed it is better viewed as a direct alignment of the model's internal direction with human preferences. The success of their method, which achieved superior results with dramatically reduced resources, indicates that the future of efficient AI training may lie in simpler, more direct feedback loops. As these models continue to grow in size and capability, the ability to fine-tune them quickly and effectively will be essential, and this new approach offers a clear, practical path forward for making artificial intelligence more responsive to human needs.

Technical Summary: Scaling Reinforcement Learning for Diffusion Models via Velocity Matching

Problem Statement

Reward fine-tuning has become a critical mechanism for aligning diffusion and flow models with human preferences and specific task objectives (e.g., semantic alignment, visual quality, motion). However, existing methods largely inherit policy-gradient machinery from autoregressive Large Language Models (LLMs), such as PPO and GRPO. This transfer is problematic because, unlike LLMs, diffusion models do not provide tractable likelihoods for generated samples.

Consequently, current approaches attempt to circumvent this intractability through two complex routes:

Trajectory-based methods: These construct likelihoods from stochastic one-step denoising transitions along a sampled trajectory. This requires expensive stochastic rollouts, trajectory storage, and often yields high-variance updates because the intermediate denoising states are stochastic and not necessarily closer to the clean endpoint than the current state.
ELBO-based methods: These approximate the marginal likelihood of the final sample using an Evidence Lower Bound (ELBO). While this decouples training from the full rollout sampler, these methods still originate from a likelihood-based policy-gradient formulation, retaining estimated policy ratios and the computational overhead of likelihood approximation.

The central question posed by this work is whether likelihood estimation is truly necessary for effective reward fine-tuning in diffusion models, or if the model update can be simplified to act directly on the native training representation.

Methodology: Reward-Based Velocity Matching (RVM)

The authors propose Reward-Based Velocity Matching (RVM), a simple, trajectory-free update mechanism that operates directly on the velocity field of the diffusion model.

Core Objective

Instead of constructing a policy-gradient objective and approximating intractable likelihood ratios, RVM treats reward fine-tuning as a regression problem on the velocity field. Given a group of generated endpoints $\{x_0^i\}$ , each sample is forward-noised to a single intermediate state $x_t^i$ with a velocity target $v_i = \epsilon_i - x_0^i$ . The model velocity $v_\theta(x_t, t)$ is updated using a loss function composed of two terms:

$\mathcal{L}_{\text{RVM}} = \frac{1}{2} \mathbb{E} \left[ r_i \| v_\theta(x_t^i, t) - v_i \|_2^2 + \beta(r_i) \| v_\theta(x_t^i, t) - v_{\text{anc}}(x_t^i, t) \|_2^2 \right]$

Where:

$r_i$ (Reward Weight): A signed preference signal derived from the reward $R(x_0^i)$ . Positive rewards reinforce the velocity direction associated with the sample (pulling $v_\theta$ toward $v_i$ ), while negative rewards suppress it (pushing $v_\theta$ away).
$v_{\text{anc}}$ (Anchor Velocity): An optional term controlling drift from a reference velocity field (e.g., a frozen reference model or an EMA of the current policy). This prevents excessive policy drift during fine-tuning.
Efficiency: The method requires only a single noised state per sample, avoiding the storage of full denoising trajectories and the computation of likelihood ratios.

Theoretical Unification

The paper demonstrates that RVM provides a general framework that recovers recent fine-tuning methods as special cases:

RAM (Reinforce Adjoint Matching): Is shown to be update-equivalent to RVM when the sampling policy is on-policy ( $\pi_{\text{old}} = \pi_\theta$ ), the anchor is the frozen reference velocity ( $v_{\text{anc}} = v_{\text{ref}}$ ), and the anchor strength is unity.
DiffusionNFT: Is shown to be a special case of RVM where the anchor is an EMA-updated velocity and the anchor strength is reward-dependent.
ELBO-based Policy Gradients: The paper shows that Exact Policy Gradient (EPG) methods, which rely on ELBO approximations of log-ratios, reduce exactly to RVM when the importance weights and time-dependent weighting factors are omitted.

This unification suggests that the specific algebraic form of the velocity loss is less critical than the design of the reward weights and the anchor mechanism.

Key Contributions

Introduction of RVM: A simple, trajectory-independent method that incorporates reward feedback directly into the native velocity-matching representation, eliminating the need for policy-ratio estimation and trajectory storage.
Theoretical Unification: Proof that RVM generalizes recent methods (RAM, DiffusionNFT) and reduces ELBO-based policy gradients to velocity regression, shifting the focus from loss algebra to reward and anchor design.
Empirical Efficiency and Performance: Demonstration that velocity-based methods match or outperform trajectory-based baselines (FlowGRPO, DanceGRPO) while substantially reducing training costs.
Reward Design Insights: Identification that for video generation, standard preference rewards often favor static outputs. The authors introduce a Dynamic-Tracking (DT) reward based on optical flow, which explicitly encourages meaningful motion, significantly improving dynamic degree and overall video quality.

Experimental Results

The authors evaluated RVM across text-to-image (Stable Diffusion 3.5), text-to-video (Wan2.1-T2V-1.3B), and image-to-video (SkyReels-I2V) models.

Performance: On Wan2.1-T2V-1.3B, RVM achieved the highest VBench Overall score (84.13), outperforming trajectory-based methods like FlowGRPO (75.91) and DanceGRPO (79.46). Notably, RVM achieved a Dynamic Degree of 75.00, compared to 55–58 for trajectory-based methods.
Efficiency: RVM reduced training costs significantly. On Wan2.1, RVM required 525 GPU-hours (8×H100), compared to 1,159 for FlowGRPO and 6,171 for DanceGRPO/TaRoS. This represents a 2.2× to 11.8× reduction in cost.
Ablation Studies:
- Loss Variants: RVM, RAM, and DiffusionNFT performed comparably, confirming that the specific velocity-loss variant matters less than the reward and anchor design.
- Reward Design: The addition of the Dynamic-Tracking (DT) reward was the primary driver for improving motion quality, raising the Dynamic Degree from ~5.5 to 75.0.
- Few-Step Generation: RVM fine-tuned models demonstrated superior robustness to aggressive step reduction (e.g., 4-step inference), maintaining high quality where base models degraded sharply.

Significance and Claims

The paper claims that scalable reward fine-tuning for diffusion models is better posed in the native velocity representation rather than as likelihood-based policy optimization. The authors argue that the complexity of trajectory-based methods and likelihood approximations is unnecessary.

Key takeaways emphasized by the authors include:

Simplicity: Effective fine-tuning can be achieved with a simple velocity-matching loss without policy ratios or trajectory storage.
Reward over Loss: The specific choice of velocity-loss functional is secondary; the critical factors are how rewards are transformed into update weights and how the model is anchored.
Motion in Video: Generic preference rewards can lead to static videos; explicit dynamic tracking is essential for high-quality video generation.

The work suggests that future scalability in diffusion model post-training should focus on reward specification and anchor design rather than increasingly elaborate policy-loss formulations.