Routing Policy Optimization for Online Diffusion Reinforcement Learning

Routing Policy Optimization for Online Diffusion Reinforcement Learning. By Qiuhong Shen and ten others, from the National University of Singapore and ByteDance. Published on the fifteenth of October. No code repository is linked to it.

Reinforcement learning can align visual generators with reward signals that go beyond their pretraining objectives. Existing score-function methods, however, typically introduce the needed stochasticity by perturbing the reverse denoising trajectory, and that leads to costly optimization built on rollouts. The exploration that reinforcement learning needs is bought by changing how the generator samples, and every rollout pays for that change. The authors ask a different question. Can stochastic exploration be introduced inside the generator instead, while its original sampler is left intact? If it can, the generator explores during training without its sampling procedure being disturbed.

Routing Policy Optimization, or RPO, answers by turning attention routing into a stochastic exploration action. At a few selected denoising levels, each query block keeps its own block and then makes a choice. It either samples further key blocks, in proportion to the generator's native attention mass, or it keeps dense attention. The randomness lives in which blocks attend to which, rather than in a perturbed denoising trajectory. Sampling in proportion to native attention mass means the routes follow where the generator already attends, and the option of staying dense is always part of the action.

A sampled route perturbs the velocity of that transition. A Gaussian likelihood centred on the dense velocity scores the perturbation, under a score-function objective. RPO adds a controlled two-rollout estimator. The two rollouts share the same prompt, the same initial latent, the same routed denoising levels, and the same dense anchor layers, and they differ only in their sampled attention routes. The difference between their rewards then provides the advantage signal for the routed transitions. Because everything else is held equal, that difference can be traced to the routes alone.

One detail makes the training carry over to deployment. Because the likelihood differentiates the dense velocity itself, the learned updates remain effective after routing is removed. And routing is removed: at inference, RPO restores dense attention and retains the original deterministic sampler, so the generator samples just as it did before training began.

Experiments cover both text-to-image and text-to-video generation. In both, RPO improves generation quality, and it does so with nearly four times higher training efficiency than Flow-GRPO-style methods. Visual results are on the project page.

The contrast with earlier methods comes down to where exploration happens. Perturbing the denoising trajectory alters the sampler and makes every rollout expensive. Sampling attention routes leaves the sampler alone, isolates the effect of the routing choice through a matched pair of rollouts, and lets the routes disappear at inference while the learning they drove remains in the dense model. Exploration becomes a property of training only.