PaperlayerView exact revision on arXiv

Source-linked AI summary

Routing Policy Optimization for Online Diffusion Reinforcement Learning

Qiuhong Shen, Yiding Yang, Yuxin Zhang, Xinwei Huang, Minxuan Lin, Bo Liu, Hao Kang, Haotian Yang, Yizhi Wang, Chongyang Ma, Xinchao Wang

arXiv:2610.22282v1cs.CV

TL;DR

Online reinforcement learning for visual generators lacks an explicit, tractable policy that preserves the original deterministic sampler. RPO samples attention routes during training and restores dense attention at inference, improving reward-aligned generation across image and video settings while reducing training cost relative to sampler-space baselines.

  • Problem

    Existing score-function methods perturb reverse denoising trajectories, while the paper identifies a need for an explicit, tractable policy that preserves the deterministic sampler.

    Source: §1 · p. 1

  • Method

    RPO samples attention routes from native attention mass at selected denoising levels and scores routed transitions with a Gaussian velocity likelihood optimized through a controlled two-rollout objective.

    Source: §1 · p. 2, §4.3 · p. 6

  • Results

    RPO improves reward-aligned quality across text-to-image and text-to-video; on Wan2.2, it achieves the strongest overall VBench score using 24.1 GPU-hours versus 75.4 for Flow-GRPO and 102.3 for DanceGRPO.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Because the likelihood updates the dense velocity, RPO removes routing at inference and retains dense attention with the original deterministic sampler.

    Source: §1 · pp. 2–3

    RPO requires an attention-based backbone, and its block-averaged route proposal and Gaussian velocity likelihood do not fully represent token-level structure or discrete route choices.

    Source: §6 · p. 10

Abstract

from arXiv · show

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.

1 INTRODUCTION

Online reinforcement learning for visual generators lacks an explicit, tractable policy that preserves the original deterministic sampler. RPO explores through attention routes during training, then removes them for dense inference; experiments report improved image and video generation quality at lower training cost.

1 INTRODUCTION

  • Prior methods either perturb reverse trajectories for score-function learning or avoid explicit sampled actions; the introduction identifies a missing tractable policy that preserves the deterministic sampler.The cited approaches include stochastic SDE conversions and forward-matching objectives.Source: §1 · p. 1
  • Two rollouts share the prompt, initial latent, routed levels, and anchor layers, differing only in sampled routes; their reward difference supplies an advantage for the routed trajectories.The reward difference reflects the aggregate effect of the two recorded trajectories.Source: Figure 1 · p. 2, §1 · p. 2
  • RPO samples attention routes at selected denoising levels, scores their velocity perturbations with a Gaussian likelihood around the dense velocity, and optimizes that likelihood with a score-function objective.Routes are sampled from the generator’s native attention mass, without an auxiliary policy head or sampler modification.Source: §1 · pp. 2–3
  • Routing is a training interface: the likelihood updates the dense velocity, and inference removes masks to restore dense attention and the original deterministic sampler.The paper tests whether reward learned through temporary routing transfers to dense image and video generation.Source: §1 · p. 2
  • RPO improves reward-aligned image and video generation; on Wan2.2, it uses 24.1 GPU-hours versus 75.4 for Flow-GRPO and 102.3 for DanceGRPO.On Stable Diffusion 3.5, it exceeds Flow-GRPO’s final dense validation reward within 25 of 50 training steps and ends 1.1 points higher.Source: §1 · p. 2

2 RELATED WORK

Related work differs in where its learning signal enters the generator: sampler trajectories, objectives without sampled sampler actions, or attention computation. RPO places a score-function action in attention and removes routing before inference.

2 RELATED WORK

  • Sampler-space methods use stochastic denoising transitions as actions; Flow-GRPO and DanceGRPO convert deterministic flow samplers into SDEs, while later work reduces trajectory cost or refines credit assignment.Examples include temporal windows, shared-prefix branches, one-step updates, and designated branch points (Li et al., 2025a;b; He et al., 2025; He et al., 2026).Source: §2 · p. 3
  • Other objectives avoid sampler actions through reward backpropagation, forward-process or preference objectives, native regression targets, hybrid training, or offline preferences.These objectives leave inference untouched but give up an explicit sampled action.Source: §2 · p. 3
  • Prior discrete policies target corruption or unmasking, while attention sparsity methods primarily pursue efficiency; KVPO explores historical-key routing in autoregressive video.The section also situates RPO alongside reinforcement learning for sparse-attention levels and attention-allocation methods.Source: §2 · p. 3
  • RPO combines block routes sampled from native attention mass in a bidirectional diffusion transformer with a Gaussian velocity likelihood that updates the dense generator; routes disappear at inference.The related-work section distinguishes this combination from prior attention-routing and paired-estimator approaches.Source: §2 · p. 4

3 PRELIMINARIES

The preliminaries define the deterministic visual-generation trajectory and reward, then describe block attention and the selected units where RPO introduces routing. The setup limits stochastic routing to a few denoising levels while retaining each query block’s own block.

3 PRELIMINARIES

  • A rollout starts from Gaussian noise and follows a deterministic sampler through denoising transitions; its terminal latent is decoded into an image or video scored by black-box reward models.Training uses reward values without differentiating through the reward model.Source: §3 · p. 4
  • Dense attention lets each query attend to every key, while block-sparse attention masks key blocks; RPO samples routes as training actions and restores dense attention at inference.Tokens are grouped into contiguous blocks of 128 along a locality-preserving Gilbert curve, producing compact image or video patches.Source: §3 · p. 4
  • RPO routes at a shared set of early denoising levels, where every query block makes one decision applied across routed layers and heads; other levels remain dense.Randomly selected anchor layers remain dense at routed levels.Source: §3 · p. 4
  • Each route retains its query block’s own block as a safety set, while the policy chooses additional candidate blocks or dense attention.This keeps local context and prevents an empty attention set; prompt, sampler, and safety block remain deterministic.Source: §3 · p. 4

4 ROUTING POLICY OPTIMIZATION

RPO uses attention routing as a temporary stochastic action for online optimization, then updates the dense velocity field so inference can return to dense attention and the original deterministic sampler.

4 ROUTING POLICY OPTIMIZATION

  • Unlike sampler-space methods that inject noise into the reverse trajectory, RPO explores by routing attention at selected denoising levels.The route determines which key blocks each query block can see.Source: §4 · p. 4

4.1 ATTENTION-MASS ROUTING

  • At routed levels, native block-attention mass proposes candidate or dense routes; Gumbel ranking selects kept blocks while safety blocks remain active and anchor layers stay dense.The scoring pass carries no gradient, and no extra parameters define the action.Source: §4.1 · p. 5

4.2 VELOCITY LIKELIHOOD OF A ROUTED TRANSITION

  • Replay scores each routed transition with a Gaussian likelihood around dense velocity; the surrogate omits route probabilities despite non-Gaussian deviations, while updating the dense field.Each level starts from its recorded input, so later levels condition on earlier routes.Source: §4.2 · p. 6

4.3 CONTROLLED TWO-ROLLOUT OBJECTIVE

  • Matched rollouts share prompt, content seed, anchor layers, and non-routing randomness; reward differences weight updates toward higher-reward routes, and masks are removed for deterministic dense inference.The pair’s reward differences are processed and used as fixed replay weights for the velocity-likelihood objective.Source: §4.3 · p. 6
  • On DrawBench, RPO improves all five reported metrics across 1,000 images from 200 prompts and 5 seeds, generated without classifier-free guidance.Source: Figure 3 · p. 7

5 EXPERIMENTS

Experiments evaluate whether reward learned through temporary attention routing transfers to dense image and video generators. RPO improves image and video results, learns faster on SD3.5 dense validation, and uses less training time than the sampler-space baselines in the reported video comparison.

5.1 SETUP

  • The experiments test transfer to dense inference on SD3.5 images and Wan2.2 videos, using HPSv3 rewards and Flow-GRPO comparisons, plus DanceGRPO for video.RPO collects 64 matched pairs per training step on both generators; DanceGRPO uses groups of eight and updates the full model.Source: §5 · p. 7, §5.1 · p. 7

5.3 LEARNING SPEED AGAINST SAMPLER-SPACE RL

  • On dense validation, RPO leads from step 10, reaches 8.90 at step 25 versus Flow-GRPO’s final 8.81, and ends at 9.91 versus 8.81.Training-rollout curves are not directly comparable as dense-model performance: RPO scores routed rollouts while Flow-GRPO scores SDE rollouts.Source: Figure 4 · p. 8, §5.3 · p. 8

5.2 TEXT-TO-IMAGE GENERATION

  • On DrawBench, ImageReward rises from −0.54 to −0.18, UnifiedReward from 2.10 to 2.32, and DeQA-Score from 3.35 to 3.56 after RPO training.The reported preference gains are clearest; the five evaluated metrics are not the training reward, and paired samples share prompt, initial noise, and sampler.Source: §5.2 · pp. 7–8

5.4 VBENCH TRANSFER AND TRAINING COST

  • RPO is the only trained checkpoint above the base model overall and exceeds both RL baselines on eight VBench dimensions, using 24.1 GPU-hours—3.1× less than Flow-GRPO and 4.2× less than DanceGRPO.VBench evaluates 16 dimensions using 6,220 13-frame videos; the Wan2.2 overall-score ranking is directional because the four scores span only 0.81 points on one seed set.Source: §5.4 · pp. 8–9, Table 1 · p. 9
  • RPO lowers color and appearance style versus the base model, although it leads the trained baselines on color; Flow-GRPO raises dynamic degree by 12.50 points but loses several quality dimensions.The paper characterizes RPO’s transfer as mainly semantic alignment and Flow-GRPO’s as trading visual quality for motion.Source: §5.4 · p. 9
  • On four held-out video prompts, RPO follows prompts more closely; the compared checkpoints use dense attention with the same prompt, initial noise, and sampler.Figure 5 shows the first and last of 13 generated frames for each model.Source: §5.4 · p. 9, Figure 5 · p. 9

6 DISCUSSION AND LIMITATIONS

RPO’s current policy action is confined to attention-based generators, and its routing proposal and likelihood simplify the underlying structure. The method currently adapts existing generators rather than shaping their architectures.

6 DISCUSSION AND LIMITATIONS

  • RPO requires a backbone that computes attention; the framework has not been generalized to architectures without attention mass for proposing routes.The policy action is a mask over key-value blocks.Source: §6 · p. 10
  • Route logits average whole-block attention mass across routed layers and heads, hiding token-level and layer-specific structure.The velocity likelihood also models routed deviations as Gaussian even though routes are discrete.Source: §6 · p. 10
  • RPO adapts existing generators rather than shaping them; architecture co-design is proposed as a direction for higher efficiency and larger RL gains.Source: §6 · p. 10

7 CONCLUSION

RPO moves stochastic exploration into attention routing while preserving the generator’s denoising trajectory and deterministic inference sampler. Experiments show improved reward-aligned generation with less training compute than sampler-space baselines.

7 CONCLUSION

  • RPO samples block-wise attention routes from native attention mass at selected denoising levels without perturbing the reverse denoising trajectory.Inference removes the masks and restores dense attention with the original deterministic sampler.Source: §7 · p. 10
  • A controlled pair fixes non-routing generation context, and its reward difference weights a Gaussian velocity likelihood over the recorded routed trajectories.The likelihood updates the dense velocity field.Source: §7 · p. 10
  • Experiments on Stable Diffusion 3.5 and Wan2.2 show improved reward-aligned quality across text-to-image and text-to-video generation with substantially less training compute than sampler-space baselines.The paper frames RPO as an efficient policy interface that does not alter inference.Source: §7 · p. 10

A GILBERT BLOCKS AND THE SAFETY SET

The appendix constructs block-wise routes using a frame-sliced Gilbert order, which makes contiguous token runs compact spatial patches. Each query block retains its own patch while a sampled route determines which additional context it can attend to.

A.1 FRAME-SLICED GILBERT ORDER

  • Raster-order runs can span distant content, whereas the frame-sliced Gilbert order makes contiguous runs compact patches and continues across frames with adjacent steps.The ordering is reused across routed layers, heads, noise levels, and rollouts; attention outputs are permuted back to raster order, leaving dense computation unchanged.Source: A GILBERT BLOCKS AND THE SAFETY SET · p. 14, §A.1 · pp. 15–16

A.2 BLOCKS, SAFETY SET, AND EXECUTED MASK

  • Blocks use 128 tokens, and each query block’s safety set is its own block while its candidates include every other block across all frames.Because attention is bidirectional, routes can add context from earlier or later frames.Source: §A.2 · p. 16
  • The executed mask combines the always-visible safety block with sampled candidates, while anchor layers remain dense and one decision serves routed layers and heads.Stable Diffusion 3.5 keeps text-involving attention pairs dense; Wan2.2 cross-attention is not routed.Source: §A.2 · p. 16

A.3 BLOCK COUNTS IN THE REFERENCE CONFIGURATIONS

Table 2 compares visual-token grids, block counts, and routing options for the reference configurations. Sparse routes retain the query’s own block and one candidate, while dense options and anchor layers preserve dense computation.

A.3 BLOCK COUNTS IN THE REFERENCE CONFIGURATIONS

  • Stable Diffusion 3.5 uses a 32 × 32 image grid; Wan2.2 uses a 4 × 22 × 40 video grid.Counts exclude text tokens, which are never routed.Source: §A.3 · p. 16, Table 2 · p. 17
  • A sparse Stable Diffusion 3.5 query block keeps its own block and one of seven candidates, seeing a quarter of the image.The configuration uses 128-token blocks.Source: §A.3 · p. 16, Table 2 · p. 17
  • A sparse Wan2.2 query block keeps its own block and one of 27 candidates, seeing about 7% of the video.The dense option and anchor layers keep the remaining computation dense.Source: §A.3 · p. 16, Table 2 · p. 17

B IMPLEMENTATION DETAILS

The appendix reports distinct routing and optimization settings for Stable Diffusion 3.5 and Wan2.2, alongside robust processing of paired reward differences. Both configurations use two prompt-group rollouts and train for 50 steps.

B IMPLEMENTATION DETAILS

  • Stable Diffusion 3.5 routes levels 2, 3, and 4 with weights 1, 0.76, and 0.5, using one candidate or dense attention per 128-token query block.It uses 16 anchor layers, centered routed steps, rank-32 LoRA adapters, and 50 training steps on eight H100 GPUs.Source: B IMPLEMENTATION DETAILS · p. 17
  • Both configurations use 64 prompt groups per wave at G = 2, split into two optimizer minibatches of 32 groups, with 5% winsorization and batch-IQR scaling.Neither configuration gates pairs.Source: B IMPLEMENTATION DETAILS · p. 17
  • Wan2.2 routes levels 2 and 4 with weights 1 and 0.5, using one of 27 candidates or dense attention per 128-token query block.It uses 20 anchor layers, classifier-free guidance at scale 5, routed-velocity steps without centering, and trains all but three specified output components for 50 steps on 16 H100 GPUs.Source: B IMPLEMENTATION DETAILS · p. 17
  • Three generic safeguards stabilize optimization but are not components of RPO: PPO ratio clipping, Wan2.2 second-minibatch log-ratio centering, and a global gradient-norm cap.The clipping width is 3 × 10^-4 and the gradient-norm cap is 1.Source: B IMPLEMENTATION DETAILS · pp. 17–18

C ROBUST ADVANTAGE PROCESSING

  • Advantage processing symmetrically winsorizes each reward dimension and scales by its interquartile range with a numerical floor, yielding a fixed replay weight.The shared symmetric processing preserves pair antisymmetry up to numerical precision; experiments use one HPSv3 reward dimension.Source: C ROBUST ADVANTAGE PROCESSING · p. 18

D EVALUATION PROTOCOLS

The evaluation protocols cover text-to-video benchmarks and text-to-image validation, with controlled comparisons accounting for shared settings and backend differences. Training compute excludes VBench evaluation.

DanceGRPO comparison

  • The DanceGRPO comparison differs from RPO in group size, sampler, rollout backend, and training schedule.DanceGRPO uses full-model GRPO with G = 8 and an SDE sampler, while RPO routes two transitions with G = 2; DanceGRPO also uses a longer 120-step schedule.Source: D EVALUATION PROTOCOLS · p. 18

VBench

  • VBench evaluates all 16 dimensions using 6,220 videos per Wan2.2 model, with shared prompts and initial noise across four Wan2.2 models.Videos use 13 frames at 704 × 1280 and 24 fps with a deterministic 50-step ODE sampler; GPU-hour accounting excludes VBench evaluation.Source: D EVALUATION PROTOCOLS · p. 18

Stable Diffusion 3.5 evaluation

  • DrawBench uses 1,000 images from 200 prompts and 5 seeds without classifier-free guidance.The Stable Diffusion 3.5 dense evaluation during training instead uses 64 held-out prompts and a 28-step sampler.Source: D EVALUATION PROTOCOLS · p. 18

Stable Diffusion 3.5 comparison

  • The Flow-GRPO comparison controls the Stable Diffusion 3.5 validation protocol, but rollout backends differ and affect training rewards.On base weights, the RPO backend scores 0.92 raw HPSv3 and the Flow-GRPO backend scores 1.38, so dense validation is the controlled comparison.Source: D EVALUATION PROTOCOLS · p. 18

E EXTENDED DISCUSSION

The discussion describes dense-inference gains that can diverge from routed-rollout rewards, identifies limitations of the velocity-likelihood surrogate, and proposes two compatible estimator extensions.

Known failure modes

  • In some runs, dense-inference gains peak and decay while routed-rollout reward keeps rising; centered routed steps reduce this split on Stable Diffusion 3.5, but the mechanism is not isolated.The discussion also identifies reward hacking against the frozen preference model as a general online-RL risk, not specific to RPO.Source: E EXTENDED DISCUSSION · pp. 18–19

Estimator scope

  • The velocity likelihood is a surrogate, modeling discrete route effects as isotropic Gaussian deviations around dense velocity and omitting later-transition and reward-model pathwise gradients.Its width s2_t is re-estimated each batch and shrinks as routing perturbs the generator less, increasing the gradient for the same deviation.Source: E EXTENDED DISCUSSION · p. 19

Extensions

  • Adding sampled-route probability to the velocity likelihood would credit the route choice itself, while larger groups would trade pairwise variance cancellation for a more stable advantage scale.Both extensions are described as compatible with the stated estimator.Source: E EXTENDED DISCUSSION · p. 19

Related papers

Loading 2610.22282v1…