Routing Policy Optimization for Online Diffusion Reinforcement Learning

arXiv:2610.22282 · cs.CV · Submitted 2026-10-15 · Read on arXiv

Listen

Transcript

Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.

Tom: I'm Tom.

Jane: And I'm Jane, here with Lu, Meng and our in-house large language model, Lalam.

Tom: We're looking at the paper: "Routing Policy Optimization for Online Diffusion Reinforcement Learning".

Jane: The gist The Routing Policy Optimization (RPO) method introduces a sampler-preserving technique for online reinforcement learning in diffusion models by turning attention routing into a stochastic exploration action to improve…

Tom: First, who's behind it and why it matters.

Title and authors: Tom: So, Jane, this paper is about something called "Routing Policy Optimization for Online Diffusion Reinforcement Learning." It sounds a bit dense, but the core idea seems to be that they want to align visual generators with reward signals without messing up how the original sampler works.

Jane: That’s right. The title itself is pretty telling, because it promises a way to do this online, and crucially, it claims it preserves the original sampler. Tom, can you tell us what they are actually proposing in plain English?

Tom: Well, they're tackling the problem where existing score-function methods introduce randomness by messing with the actual denoising trajectory. Instead of that, RPO turns attention routing into a stochastic exploration action to improve generation quality while keeping the original sampler intact.

Lu: What's interesting about that mechanism is that it doesn't add a new policy head or modify the existing sampler itself; it just samples routes from the generator’s native attention mass at a few specific denoising levels.

Jane: So, this routing action then perturbs the velocity of that transition, and they score this perturbation using a Gaussian likelihood around the dense velocity scores as a surrogate objective. That sounds like they're essentially guiding the generation process based on how well it fits the existing dense velocity field.

Meng: From an engineering standpoint, scoring transitions with a likelihood around the dense velocity is clever because it means that this update acts directly on that same velocity field that inference uses, which is good for deployment.

Tom: Exactly. And they use a controlled two-rollout estimator where two rollouts share the prompt, initial latent, routed denoising levels and anchor layers but differ only in their sampled attention routes. The difference in reward between those two rollouts gives them an advantage signal for the routed transitions.

Jane: That paired advantage signal is key to how they update things. The loss function L RPO(theta) is defined as one-half sigma over i of the velocity likelihood Si, where Si is that velocity likelihood from Equation five > It sounds like they are pulling the dense velocity toward the recorded routed transitions of the higher-reward member and pushing it away from those of the lower-reward member <ref:2610.22282#pg2>.

Title and authors: Lu: That mechanism directly addresses how to optimize without needing a massive rollout group, which is what sampler-space methods usually require > It moves stochastic exploration from being in the sampler space into being inside the attention mechanism itself.

Tom: The results are pretty compelling. They tested RPO on Stable Diffusion three point five and Wan2 point 2 and found that it improves generation quality while keeping dense inference on both models > On Stable Diffusion three point five, RPO exceeds the final dense validation reward of Flow-GRPO within twenty-five of fifty training steps, and ends one point higher > <ref:2610.22282#pg3,RPO on Stable Diffusion 3.5 and Wan2.2>

Jane: And for Wan2 point 2, RPO achieved the strongest overall VBench score among the evaluated checkpoints while only using twenty-four point one GPU hours, which is significantly lower than seventy-five point four for Flow-GRPO > That translates to three point one and four point two times lower training cost compared to those other methods <ref:2610.22282#pg1>.

Meng: From a production perspective, that efficiency gain is huge because it means we can train better models faster without needing massive compute clusters for these kinds of online adjustments > And since it retains the original sampler, we don't have to worry about breaking the deterministic generation process after training.

Lalam: I see this paper focuses on introducing a sampler-preserving formulation for online alignment of continuous visual generators where the exploration action is a block route sampled from native attention mass at a few selected denoising levels with no auxiliary policy head or SDE conversion > That sounds like they're keeping the core generation mechanism very clean.

Tom: So, to wrap up, this paper introduces Routing Policy Optimization for Online Diffusion Reinforcement Learning by turning attention routing into a stochastic exploration action to improve generation quality with nearly four times higher training efficiency than Flow-GRPO-style methods > It’s about using the generator's own attention mass for steering without changing how the final image or video is actually rendered.

Jane: It’s clear they are making a strong case for integrating exploration directly into the attention process rather than relying on external policy methods > The main implication here is that we can get reward alignment with much less computational overhead while maintaining the integrity of the original generation path.

Title and authors: Lu: The limitation they mention is that this framework is confined to attention, so you need a backbone that computes attention for routing > Plus, they note that since the velocity likelihood treats the routed deviation as Gaussian although it's a discrete route, and because routing perturbs the generator less during training causing the variance s squared t to shrink which produces larger gradients > That suggests there are still some nuances in how well this approximation holds up.

Meng: And from my side, I think that limitation on the backbone is practical. It means if we have a text-to-video model where the attention structure isn't what they expect, this specific RPO formulation might not be applicable without modifications > We need to see how robust this is across different architectures.

Lalam: The velocity likelihood being a surrogate rather than the exact likelihood of the sampled route means it omits pathwise gradient information through later denoising transitions and through the frozen reward model > That’s a known trade-off when using surrogates in these types of optimization setups.

Tom: So, to wrap up, Routing Policy Optimization for Online Diffusion Reinforcement Learning provides an efficient policy interface for reward-aligned visual generation without altering the inference process by moving stochastic exploration from the sampler into attention > It shows we can achieve quality improvements across text-to-image and text-to-video generation with substantially less training compute than sampler space baselines.

Jane: It’s a solid piece of work because it manages to introduce this exploration mechanism while keeping the core deterministic sampling process intact for the final output > We should keep an eye on how they address those limitations regarding the velocity likelihood approximation in future work.

Lu: I think if they can figure out how to incorporate the probability of the sampled route into that velocity term, it would credit the route choice itself rather than just its effect on the dense velocity > That would make it even more principled from a modeling standpoint.

Meng: For deployment, we need that stability. If we can get better guarantees on those gradients when s squared t is shrinking, then this becomes much more viable for real-world training pipelines > Right now, it's promising but not completely plug-and-play yet.

Lalam: I'm fairly confident that the core insight—moving the exploration into attention while preserving the sampler—is a valuable direction for future work in diffusion RL > It’s a principled way to interface policy learning with generative models.

The paper's summary: Tom: So, to recap, this paper introduces Routing Policy Optimization for Online Diffusion Reinforcement Learning by turning attention routing into a stochastic exploration action to improve generation quality with nearly four times higher training efficiency than Flow-GRPO-style methods >

Jane: Right. It’s about using the generator's own attention mass for steering without changing how the final image or video is actually rendered, and it does this way faster computationally >

Lu: What they are doing there is turning a deterministic process into something stochastic during training, but keeping the inference process totally intact so you don't have to worry about breaking your sampler later >

Meng: I mean, that efficiency gain of nearly four times over Flow-GRPO is significant because it cuts down on GPU hours and training time, which is what matters when we're trying to ship models to production >

Tom: Exactly. And the way they handle the scoring—using a Gaussian likelihood around the dense velocity—that’s how they create an advantage signal for those routed transitions without needing a whole new policy head >

Jane: It sounds like they are using the generator's own internal state, that dense velocity field, as a proxy to judge how good a route is, which is really smart because it keeps everything tied to the generation process itself >

Lu: They use this controlled two-rollout setup where only the attention routes change between those two samples; the difference in reward tells them what worked >

Meng: That structure lets them get that advantage signal without needing a massive rollout group of samples, which makes sense for scaling up training on real traffic data >

Tom: And after training, they just remove all those masks and run the dense attention with its original deterministic sampler, so the output quality should be what we expect from a standard diffusion model >

Jane: So it boils down to this: they found a way to introduce exploration during training that isn't messing with the final generation step or requiring a separate policy network >

Lu: It’s an interesting formulation for online alignment, specifically because the exploration action is sampled directly from the native attention mass without needing some complex SDE conversion >

Meng: The main thing I see here is that this method makes reward alignment much more accessible to practitioners who can't afford massive training runs >

Tom: Yeah, and it shows we can get better quality across text-to-image and text-to-video generation while using substantially less compute than the established sampler space baselines >

Jane: So the big picture is that we’re moving away from external policy methods for reinforcement learning in diffusion models toward something integrated directly into the attention mechanism itself >

Lu: The limitation they mention is that this framework is confined to attention, so you need a backbone that computes attention for routing, which isn't always the case >

Meng: That’s a practical constraint. If we have a video model where the attention structure doesn't fit their routing assumption, this specific RPO formulation might not work without some heavy modification >

Tom: The velocity likelihood being a surrogate—meaning it’s an approximation instead of the exact likelihood—that omits pathwise gradient information through later denoising transitions >

Jane: That means they aren't getting the full, detailed gradient path through every single step, which is a trade-off for keeping the optimization tractable >

Lu: The authors suggest that if you added the probability of the sampled route to that velocity term, it would credit the route choice itself more directly >

Meng: If they could figure out how to make those gradients more stable when s squared t shrinks during training, then this becomes much more viable for real-world deployment >

Tom: So RPO provides a very efficient policy interface for reward alignment without altering inference, and it’s showing strong performance across different models >

Jane: It’s a solid piece of work because it manages to introduce this exploration mechanism while keeping the core deterministic sampling process intact for the final output >

The paper's improvements: Tom: So, they’re looking at how to make this RPO method even better by addressing some of those little technical gaps they mentioned in the paper >

Jane: Right, like that issue where you have a discrete route but you treat it like it’s a smooth Gaussian deviation; they want to fix that modeling mismatch >

Lu: They suggest adding the actual probability of the sampled route into that velocity term, which would credit the choice itself rather than just its effect on the dense velocity >

Meng: If they incorporate that probability, it means their objective function becomes more principled because it’s accounting for how likely a certain path was in the first place >

Tom: And another point is about those variance terms; they noticed that since routing perturbs the generator less during training, the variance s squared t shrinks which causes larger gradients >

Jane: So, if they can get better guarantees on those gradients when things are shrinking, then this becomes much more stable for real-world training pipelines >

Lu: That’s a big engineering hurdle. If we can stabilize that part of the loss function, the method becomes much more robust across different training runs >

Meng: I need to see that stability before I can tell my team it's production ready, because right now it sounds like a bit of a research artifact >

Tom: The limitation they flag is that the velocity likelihood is just a surrogate—it skips pathwise gradients through later denoising steps and through the frozen reward model >

Jane: That means even with these suggested improvements, you’re still relying on an approximation instead of having the full gradient information from every single step >

Lu: It's that trade-off between tractability and fidelity in these types of optimization setups; they are acknowledging what’s missing from the current formulation >

Meng: So while they’re pushing for better modeling, I still see that this framework is confined to attention, meaning it only works on models where routing isn't just an afterthought >

Tom: Exactly. It requires a backbone that actually computes attention for the routing action to be valid in the first place >

Jane: So, even with these suggested improvements, the core limitation remains that this method doesn't apply to every kind of generative model out there >

Conclusion: Tom: So to wrap up on "Routing Policy Optimization for Online Diffusion Reinforcement Learning," we’ve seen how this method turns attention routing into a stochastic action to guide learning without messing up the original sampler >

Jane: It really shows that you can get better alignment with rewards by steering the attention process itself, rather than adding a whole new policy network on top of it >

Lu: The core idea is keeping the sampling deterministic while using those sampled routes to create an advantage signal through a controlled two-rollout estimator >

Meng: From my side, the main thing that stands out is how much training compute this saves; nearly four times better than Flow-GRPO methods means we can iterate on models much faster >

Tom: Yeah, and the results on Stable Diffusion three point five and Wan2 point two show that this quality improvement comes without sacrificing dense inference performance on either model > <ref:2610.22282#pg1>

Jane: It’s a lot of practical benefit for anyone building visual generation tools because you get higher quality images or videos with less training time >

Lu: The main caveat, though, is that the framework is tied to attention; you need a backbone that computes attention for the routing action to be valid >

Meng: And as an engineer, I’m worried about those constraints; if we have a video model where the attention structure isn't what they expect, this specific RPO formulation might not work without some heavy modification >

Tom: Yeah, and they also noted that because the velocity likelihood is a surrogate—it skips pathwise gradients through later denoising steps—that’s an approximation we have to live with >

Jane: So while it’s a solid method for online alignment, you still aren't getting the full gradient information from every single step of the generation process >

Lu: The authors suggest adding the probability of the sampled route into that velocity term, which would credit the route choice itself more directly >

Meng: If they can stabilize those gradients when s squared t shrinks during training, then this becomes much more viable for real-world deployment >

Tom: So to summarize "Routing Policy Optimization for Online Diffusion Reinforcement Learning," it’s an efficient way to interface policy learning with generative models without altering the inference process >

Jane: It proves that moving stochastic exploration from the sampler into attention is a principled way to align visual generation quality with reward signals >

Lu: I think if they can incorporate the probability of the sampled route, it would make it even more principled from a modeling standpoint >

Meng: For production, we need those guarantees on gradient stability before we can say this is totally plug-and-play >

Qiuhong Shen, Yiding Yang, Yuxin Zhang, Xinwei Huang, Minxuan Lin, Bo Liu

National University of Singapore · ByteDance

cs.CV

Submitted: 2026-10-15

Updated: 2026-10-15

Project page: https://velox-world.com/rpo

License: http://creativecommons.org/licenses/by/4.0/

Importance score: 91/100

The gist: The gist The Routing Policy Optimization (RPO) method introduces a sampler-preserving technique for online reinforcement learning in diffusion models by turning attention routing into a stochastic

Key concepts

Sampling Attention Routes: RPO explores by selecting specific attention routes directly from the generator's existing native attention mass at several denoising steps. This allows for a stochastic action—the route choice—to be made during training without needing to add a separate policy network or change how the original diffusion sampler operates.
Scoring Transitions: Each sampled transition is evaluated by calculating a Gaussian likelihood around the dense velocity scores of the generator. This likelihood acts as a surrogate score function, providing an objective signal that guides the optimization process toward better generation quality based on how well it aligns with the dense velocity field.
Controlled Two-Rollout Estimator: RPO uses a specific setup where two model rollouts are compared. These rollouts share the same input (prompt, latent, and routed levels) but differ only in their sampled attention routes. The difference in their resulting rewards creates an advantage signal that tells the system which routed transition was better.
Dense Velocity Field Update: The training update modifies the core velocity field that is used during standard inference. Because the likelihood function used for optimization differentiates this dense velocity, the learned adjustments remain effective when all routing mechanisms are removed, ensuring a sampler-preserving outcome.

Terminology

Generator — 3× in this paper · explained in 2 episodes
Reinforcement Learning — 2× in this paper · explained in 4 episodes
Routing Policy Optimization — 2× in this paper
GRPO — 1× in this paper · explained in 6 episodes

Summary

The gist The Routing Policy Optimization (RPO) method introduces a sampler-preserving technique for online reinforcement learning in diffusion models by turning attention routing into a stochastic exploration action to improve generation quality with nearly 4× higher training efficiency than Flow-GRPO-style methods

How it works

RPO addresses the bottleneck of policy in online reinforcement learning by introducing an explicit, tractable policy that preserves the original deterministic sampler for both image and video generators The core mechanism involves sampling attention routes from the generator’s own native attention mass at a few selected denoising levels Each sampled route perturbs the velocity of its transition, and a Gaussian likelihood around the dense velocity scores this perturbation with a score-function objective

Key Components of RPO

Sampling Attention Routes: The method samples attention routes from the generator’s own native attention mass at a few denoising levels This sampling exposes a stochastic action without adding a policy head or modifying the sampler
Scoring Transitions: Each routed transition is scored with a Gaussian likelihood around the dense velocity This likelihood supplies a score-function surrogate for the dense generator itself
Controlled Two-Rollout Estimator: RPO uses a controlled two-rollout estimator where two rollouts share the same prompt, initial latent, routed denoising levels, and anchor layers They differ only in their sampled attention routes Their reward difference provides an advantage signal for the routed transitions

Optimization and Inference

The update acts on the same velocity field that dense inference evaluates because the likelihood differentiates the dense velocity itself The loss function is defined as L RPO(θ) = 1/2 Σ i L pg i(θ), where Si is the velocity likelihood from Equation 5 This update pulls the dense velocity toward the recorded routed transitions of the higher-reward member and away from those of the lower-reward member After training, RPO removes all masks and runs dense attention with its original deterministic sampler

Experimental Results

Experiments on Stable Diffusion 3.5 and Wan2.2 show that RPO improves generation quality while retaining dense inference On Stable Diffusion 3.5, RPO exceeds the final dense validation reward of Flow-GRPO within 25 of 50 training steps and ends 1.1 points higher On Wan2.2, RPO attains the strongest overall VBench score among the evaluated checkpoints while using 24.1 GPU-hours, compared with 75.4 for Flow-GRPO and 102.3 for DanceGRPO This corresponds to 3.1× and 4.2× lower training cost, respectively

Contributions and Limitations

The contributions include introducing a sampler-preserving formulation for online alignment of continuous visual generators where the exploration action is a training-time attention route sampled from native attention mass with no auxiliary policy head or SDE conversion The method scores each routed transition with a Gaussian velocity likelihood around the dense velocity and optimizes it with a controlled two-rollout surrogate that holds the non-routing generation context fixed A limitation is that the framework is confined to attention, requiring a backbone that computes attention for routing Furthermore, the velocity likelihood treats the routed deviation as Gaussian although a discrete route produces it The estimator scope is limited because the velocity likelihood is a surrogate rather than the exact likelihood of the sampled route

Conclusion

RPO provides an efficient policy interface for reward-aligned visual generation without altering the inference process by moving stochastic exploration from the sampler into attention The method demonstrates improved reward-aligned quality across text-to-image and text-to-video generation while requiring substantially less training compute than sampler-space baselines By moving stochastic exploration from the sampler into attention, RPO provides an efficient policy interface for reward-aligned visual generation without altering the inference process The remaining failure modes involve a dense-gain fade and reward hacking against the frozen preference model The velocity likelihood is a surrogate that omits the pathwise gradient of the generator through later denoising transitions and through the frozen reward model The width s 2 t is re-estimated on every batch as routing perturbs the generator less during training, causing s 2 t to shrink and produce larger gradients The extension of adding the probability of the sampled route to the velocity term would credit the route choice itself rather than only its effect on the dense velocity The larger groups would give a leave-one-out baseline with more members, trading part of the pair’s variance cancellation for a more stable advantage scale The framework is confined to attention and requires a backbone that computes attention for routing

--- Page 1 ---

Key Components of RPO

Sampling Attention Routes: The method samples attention routes from the generator’s own native attention mass at a few selected denoising levels This sampling exposes a stochastic action without adding a policy head or modifying the sampler
Scoring Transitions: Each routed transition is scored with a Gaussian likelihood around the dense velocity This likelihood supplies a score-function surrogate for the dense generator itself
Controlled Two-Rollout Estimator: RPO uses a controlled two-rollout estimator where two rollouts share the same prompt, initial latent, routed denoising levels, and anchor layers They differ only in their sampled attention routes Their reward difference provides an advantage signal for the routed transitions

Optimization and Inference

Improvements for AI systems

Bold header: Sampler-Preserving Policy Optimization

This method introduces Routing Policy Optimization (RPO), a sampler-preserving formulation for online alignment of continuous visual generators, which allows online reinforcement learning to optimize objectives without perturbing the reverse denoising trajectories.

Bold header: Attention-Based Stochastic Exploration

RPO explores by routing attention, defining the action as a block route sampled from native attention mass at a few selected denoising levels and scoring it with a Gaussian likelihood around the dense velocity, thereby enabling stochastic exploration without adding a policy head or modifying the sampler.

Bold header: Controlled Reward Signal Generation

The system utilizes a controlled two-rollout estimator where rewards are generated by pairing rollouts that share the prompt, initial latent, routed denoising levels, and anchor layers, and using the reward difference provides an advantage signal for the routed transitions.

Bold header: Dense Velocity Field Update

The optimization update is applied directly to the core generator via a score-function surrogate; specifically, the likelihood differentiates the dense velocity itself, ensuring that the learned updates remain effective after routing is removed at inference.

Bold header: Enhanced Training Efficiency

RPO demonstrates significant computational savings, achieving nearly 4× higher training efficiency than Flow-GRPO-style methods and requiring substantially less training compute while improving reward-aligned quality across text-to-image and text-to-video generation.

Abstract

Reinforcement learning can align visual generators with reward signals beyond their pretraining objectives, but existing score-function methods typically introduce stochasticity by perturbing the reverse denoising trajectory, leading to costly rollout-based optimization. We instead ask whether stochastic exploration can be introduced inside the generator while preserving its original sampler. Routing Policy Optimization (RPO) answers this question by turning attention routing into a stochastic exploration action. At a few selected denoising levels, each query block keeps its own block and samples further key blocks in proportion to the generator's native attention mass, or keeps dense attention. The sampled route perturbs the velocity of that transition, and a Gaussian likelihood around the dense velocity scores the perturbation with a score-function objective. RPO further uses a controlled two-rollout estimator: two rollouts share the same prompt, initial latent, routed denoising levels, and dense anchor layers, differing only in their sampled attention routes. Their reward difference provides an advantage signal for the routed transitions. Because the likelihood differentiates the dense velocity itself, the learned updates remain effective after routing is removed. At inference, RPO restores dense attention and retains the original deterministic sampler. Experiments on text-to-image and text-to-video generation show that RPO improves generation quality while achieving nearly 4x higher training efficiency than Flow-GRPO-style methods. Please see our project page at https://velox-world.com/rpo for visual results.