Routing Policy Optimization for Online Diffusion Reinforcement Learning

AI summary

Overview

Research area: Reinforcement learning for visual generative models (text-to-image and text-to-video diffusion/flow transformers), specifically reward alignment of pretrained generators.

Technical level: Advanced. The paper assumes familiarity with diffusion/flow-matching samplers, score-function policy gradients, and transformer attention.

Scope: This paper proposes Routing Policy Optimization (RPO), a method that moves the stochastic exploration action of online reinforcement learning from the denoising sampler into the generator's attention routing, so that reward alignment can be trained without perturbing the reverse trajectory and without changing the sampler at inference.

What This Paper Is About

Online reinforcement learning can align image and video generators with reward models, but existing score-function methods (such as Flow-GRPO and DanceGRPO) create exploration by injecting noise into the reverse denoising trajectory, which forces many costly stochastic rollouts. The authors ask whether exploration can instead come from inside the generator itself, so the original deterministic sampler is preserved. RPO answers this by treating the choice of which attention key-blocks each query block can see — a "route" — as the stochastic action, and by scoring the resulting velocity perturbation with a Gaussian likelihood around the dense velocity.

Key Contributions

1. A sampler-preserving formulation for online alignment. Exploration is a training-time attention route sampled from the generator's own native attention mass, with no auxiliary policy head and no conversion of the ODE sampler into an SDE.

2. A Gaussian velocity likelihood for routed transitions. Each routed transition is scored by a Gaussian centered on the dense velocity, with a fixed width derived from the mean squared routing deviation at that level, so the update differentiates the dense velocity field directly.

3. A controlled two-rollout estimator. Two rollouts share the prompt, initial latent, routed denoising levels, and anchor layers, differing only in routing seeds; their signed reward difference serves as a leave-one-out advantage.

4. Routing is removed at inference. Because the likelihood updates the dense velocity, all masks are dropped after training and the model runs dense attention with its original deterministic sampler, validated on Stable Diffusion 3.5 Medium (text-to-image) and Wan2.2-TI2V-5B (text-to-video).

Main Findings

• Image quality improves on metrics that were never the training reward. On DrawBench (1,000 images from 200 prompts with 5 seeds, no classifier-free guidance), RPO raises ImageReward from −0.54 to −0.18, UnifiedReward from 2.10 to 2.32, and DeQA-Score from 3.35 to 3.56; the paper states all five reported metrics improve.

• Faster dense-model reward gain than Flow-GRPO on SD3.5 Medium. Both methods train the same LoRA adapters with HPSv3 reward and a 10-step training sampler for 50 steps from a base model that scores 3.94 and 3.96 on their shared validation. RPO reaches 8.90 at step 25, already above Flow-GRPO's step-50 score of 8.81, peaks at 10.20 at step 40, and ends at 9.91 — a gain of +5.97 over its step-zero score versus +4.86 for Flow-GRPO, ending 1.1 points higher.

• Training curves are reported separately for routed and dense rollouts. The paper notes that RPO's training-rollout reward (smoothed, above 9 by step 30, near 9.5) and Flow-GRPO's SDE-rollout reward (below 8) are not comparable, and treats the dense validation curve as the decisive comparison. Flow-GRPO is slightly faster at step 5 only.

• Best overall VBench score among evaluated checkpoints. On 6,220 generated 13-frame videos, RPO scores 82.00 overall versus 81.61 for the Wan2.2-TI2V-5B base model, 81.46 for DanceGRPO, and 81.19 for Flow-GRPO (with CogVideoX-5B at 81.30 and HunyuanVideo at 78.69 as external references). It is the only trained checkpoint above the base model overall.

• Gains are mainly semantic. Relative to the base model, RPO improves object class by +4.40, human action by +4.80, scene by +3.43, and multiple objects by +2.15, and also improves aesthetic quality, imaging quality, subject consistency, and motion smoothness. RPO lowers color and appearance style relative to the base model, though it still leads the trained baselines on color.

• Much lower training cost. RPO uses 24.1 GPU-hours on H100 GPUs versus 75.4 for Flow-GRPO and 102.3 for DanceGRPO — 3.1× and 4.2× lower, respectively. The authors state the runs are not compute-matched.

• Different trade-off from Flow-GRPO on video. Flow-GRPO raises dynamic degree by 12.50 points but loses consistency, smoothness, aesthetic quality, color, and object-class accuracy; RPO transfers reward mainly through semantic alignment.

• The authors flag the aggregate ranking as directional. The four Wan2.2 overall scores span only 0.81 points on one seed set, and the image learning-speed curves come from a single training seed.

Methodology in Plain English

The generator is a diffusion transformer that splits a latent into visual tokens (M tokens, one per spatial or spatiotemporal patch). In attention, every query token compares itself to every key token. The authors group tokens into contiguous blocks of B = 128 tokens — ordered along a frame-sliced Gilbert curve so each block is a compact patch rather than a thin horizontal strip — and let each query block choose which key blocks it is allowed to see.

At a small fixed set of early denoising levels (the "routed levels"), a gradient-free scoring pass measures how much attention each query block already places on each key block. Those attention masses become logits for a categorical draw over the other blocks plus a "stay dense" option, using a temperature T. Each query block always keeps its own block (a safety set, so the softmax is never empty), and keeps its top-K candidates by a keep ratio κ with Gumbel noise for exploration. A random subset of transformer blocks, the "anchor layers," stays dense at every routed level.

Because the sampled route changes the velocity at that transition, the authors model the routed step as a sample from a Gaussian policy centered on the dense velocity, with the variance set to the batch's mean squared routing deviation at that level so the log-likelihood stays of order one. The route-designing pass carries no gradient; the gradient flows only through the dense velocity, which is why the update still applies to the dense model used at inference.

For the learning signal, each prompt is rendered twice with identical content seed (same prompt, same initial latent, same anchor layers, same routed levels) and different routing seeds. The sampler is deterministic, so the two outputs differ only through routing. The reward difference between the pair gives each rollout an advantage with the partner as a leave-one-out baseline. Differences are winsorized, scaled by a robust interquartile scale, summed over reward dimensions, and frozen with a stop-gradient before replay. Replay reconstructs each routed velocity from the recorded latents, recomputes the dense velocity with gradient, and applies the policy-gradient loss.

At inference, all masks and reordering are removed and the stock dense sampler runs unchanged.

Why This Matters

Impact on research. The paper separates two things that prior online RL for visual generators bundled together: exploration and the sampler. By showing that a discrete attention-level action can carry a usable score-function signal while leaving inference untouched, it opens a route to RL alignment that avoids the SDE conversion that makes Flow-GRPO-style methods expensive. It also connects structured attention sparsity — previously studied mainly for efficiency — to reward alignment, and shows that a two-rollout estimator (echoing 2-GRPO in language models) can be enough when the two rollouts are tightly matched.

Real-world applications:

• Cheaper post-training of text-to-image models against human-preference or aesthetic reward models, since the training reward need not be the deployment metric.
• Video generation alignment, where rollout cost dominates: 24.1 versus 75.4 and 102.3 GPU-hours matters at production scale.
• Prompt-adherence fixes for commercial image and video generators (the paper's examples include correctly rendering requested storefront text and a brown giraffe, and replacing distorted painterly textures with coherent scenes).
• Any pipeline that must ship a deterministic, fixed-step sampler for reproducibility or latency reasons, since RPO does not require changing the inference path.

Industry relevance. The efficiency numbers and the "remove the masks at inference" property are directly relevant to teams that have already deployed a distilled or few-step generator and cannot afford to swap its sampler during fine-tuning.

Future Directions

• Extending beyond attention-based backbones. The action is a mask over key-value blocks, so it requires a backbone that computes attention; the authors state they have not generalized routing to architectures with no attention mass to propose routes.

• Finer routing proposals and likelihoods. Each route logit averages attention mass over a whole block, across all routed layers and heads, hiding token-level and layer-specific structure; a finer proposal and a likelihood that also credits the route choice itself could give a more faithful policy.

• Co-designing architecture with the RL method. The authors position RPO as adapting an existing generator rather than shaping it, and name architecture-method co-design as a direction for both higher efficiency and larger RL gains.

• Strengthening the empirical evidence. The image learning-speed comparison uses one training seed, the four Wan2.2 overall VBench scores span only 0.81 points, and the compute comparisons are not compute-matched — all of which the paper itself flags as limits on how strongly the ranking should be read.

Target Audience

Researchers and engineers working on reward alignment and post-training of diffusion and flow-matching generative models, particularly those interested in reinforcement learning for image and video generation. It is also relevant to practitioners who need to fine-tune a deployed generator without altering its deterministic sampler, and to readers interested in structured/sparse attention as a control interface rather than only an efficiency trick. Readers without background in score-function policy gradients and transformer attention will find the method sections demanding.

Authors’ 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.