Online reinforcement learning (RL) fine-tunes diffusion models with black-box rewards, but reverse-process methods such as Flow-GRPO are limited in efficiency by poor credit assignment, while forward-process methods such as DiffusionNFT are efficient yet unstable without carefully tuned exponential-moving-average (EMA) off-policy training. We show that NFT-style objectives contain a hidden forward--reverse consistency force that joint positive--negative training cancels only partially, and that EMA stabilizes training through restoration rather than off-policy rollouts. We introduce DiffusionART (Diffusion Reinforcement with Aggregation-First Reward Tilting), which aggregates each prompt's rewards into a shared reward-improving direction for reverse-process states, removing the consistency force by construction, and adds a trust-region-masked second update for drift-adaptive restoration without EMA tuning. It requires no differentiable reward model, remains effective with fewer samples per prompt, and supports pairwise reward models (DiffusionART-Pairwise). Across image and video generation, DiffusionART trains efficiently and stably without per-task tuning. In particular, on HunyuanVideo-13B with equally weighted visual and motion quality rewards, DiffusionART reaches the late-stage mean reward levels of DiffusionNFT and Flow-GRPO-Fast with 12.02 and 11.72 times fewer GPU hours, respectively, while achieving 28.8% and 53.0% higher mean training rewards.
Figures & tables
Figure 1: Efficient learning and improved generations. Left: Mean of raw VQ and MQ (equal weights). Right: Selected same-prompt images and video frames before and after DiffusionART.
Figure 2: (a) Ablation results for positive and negative objectives. (b) The positive and negative objectives partially cancel the consistency force.
Figure 3: (a) On-policy DiffusionNFT without EMA ( β=1 ): increasing number of micro-batches per rollout from 1 to 2 results in 15.2× slower training. The dashed curve is the one-update run’s gradient norm, spiking as the reward collapses. (b) EMA-policy rollouts without the old-policy anchor still collapse. (c) At the second update, averaged over rollouts 0–17, restoration has a 9.4× smaller output-force RMS but a 4.9× larger parameter-gradient norm than guidance. (d) The second MSE update decomposes into guidance and restoration.
Figure 4: Aggregation-first guidance. Left: NFT’s positive and negative branches introduce opposing consistency forces that may not fully cancel (Section 3.1 ). DiffusionART avoids these forces by applying G(c) directly at the reverse velocity. Right: We aggregate each prompt’s endpoints and rewards into normalized guidance G(c) (Equation 5 ) and apply it to states visited along this rollout.
Figure 5: Training rewards: (a–c) single- and (d) multi-reward SD3.5-M.
Figure 6: (a) FLUX.1-dev on GenEval. (b) Guidance and restoration ablation on OCR.
Figure 7: Efficient reward improvement across video generation. (a–c) T2V: DiffusionART improves VQ and MQ fastest while maintaining TA comparable to the strongest baseline. (d–e) I2V: DiffusionART rapidly improves both rewards, whereas the baselines improve slowly or plateau.
Figure 9: Training reward curves for four jointly optimized objectives on SD3.5-M over 500 steps. PickScore and CLIPScore are shown as unscaled model logits.
Figure 10: T2V with matched VQ:MQ =1:1 weights. The left pair shows VQ and the right pair MQ, each versus training steps and GPU hours.
Figure 11: DiffusionART versus ReFL with matched VQ:MQ =1:1 weights. The left pair shows VQ and the right pair MQ, each versus training steps and GPU hours.
Figure 13: Trust-region diagnostics on HunyuanVideo-13B. (a) Parameter gradient norm before clipping. (b) Scheduler-reference-weighted velocity drift energy, averaged over training states before each update.
Figure 14: Guidance assignment on GenEval.
Figure 15: Text-to-video generation before and after DiffusionART. Each pair uses the same prompt, seed, and sampling settings. Four chronologically ordered frames are shown at identical frame positions within each pair: pretrained outputs on the left and DiffusionART outputs on the right.
Figure 16: Text-to-image generation with multi-reward optimization. Selected SD3.5-M outputs before and after DiffusionART, jointly optimized with GenEval, CLIPScore, PickScore, and HPSv2. Each pair shows the same prompt, with pretrained outputs on the left and DiffusionART outputs on the right.
Figure 17: Additional text-to-image comparisons. The top two rows use multi-reward optimization; the bottom two use PickScore. Each pair shows pretrained and DiffusionART outputs under the same prompt.
Diffusion and flow models achieve State-Of-The-Art (SOTA) generative performance, yet many practically important behaviors such as fine-grained prompt fidelity, compositional correctness, and text rendering are weakly specified by score or flow matching pretraining objectives. Reinforcement Learning (RL) fine-tuning with external, black-box rewards is a natural remedy, but diffusion RL is often brittle. Trajectory-based methods incur high memory cost and high-variance gradient estimates; forward-process approaches converge faster but can suffer from distribution drift, and hence reward hacking. In this work, we present Centered Reward Distillation (CRD), a diffusion RL framework derived from KL-regularized reward maximization built on forward-process-based fine-tuning. The key insight is that the intractable normalizing constant cancels under within-prompt centering, yielding a well-posed reward-matching objective. To enable reliable text-to-image fine-tuning, we introduce techniques that explicitly control distribution drift: (i) decoupling the sampler from the moving reference to prevent ratio-signal collapse, (ii) KL anchoring to a CFG-guided pretrained model to control long-run drift and align with the inference-time semantics of the pre-trained model, and (iii) reward-adaptive KL strength to accelerate early learning under large KL regularization while reducing late-stage exploitation of reward-model loopholes. Experiments on text-to-image post-training with GenEval and OCR rewards show that CRD achieves competitive SOTA reward optimization results with fast convergence and reduced reward hacking, as validated on unseen preference metrics.
Reinforcement learning fine-tuning has become the dominant approach for aligning diffusion models with human preferences. However, assessing images is intrinsically a multi-dimensional task, and multiple evaluation criteria need to be optimized simultaneously. Existing practice deal with multiple rewards by training one specialist model per reward, optimizing a weighted-sum reward R(x)=∑kwkRk(x), or sequentially fine-tuning with a hand-crafted stage schedule. These approaches either fail to produce a unified model that can be jointly trained on all rewards or necessitates heavy manually tuned sequential training. We find that the failure stems from using a naive weighted-sum reward aggregation. This approach suffers from a sample-level mismatch because most rollouts are specialist samples, highly informative for certain reward dimensions but irrelevant for others; consequently, weighted summation dilutes their supervision. To address this issue, we propose MARBLE (Multi-Aspect Reward BaLancE), a gradient-space optimization framework that maintains independent advantage estimators for each reward, computes per-reward policy gradients, and harmonizes them into a single update direction without manually-tuned reward weighting, by solving a Quadratic Programming problem. We further propose an amortized formulation that exploits the affine structure of the loss used in DiffusionNFT, to reduce the per-step cost from K+1 backward passes to near single-reward baseline cost, together with EMA smoothing on the balancing coefficients to stabilize updates against transient single-batch fluctuations. On SD3.5 Medium with five rewards, MARBLE improves all five reward dimensions simultaneously, turns the worst-aligned reward's gradient cosine from negative under weighted summation in 80% of mini-batches to consistently positive, and runs at 0.97X the training speed of baseline training.
Canyu Zhao, Hao Chen, Yunze Tong +3
Zhejiang University · Work done during an internship at HiThink. · HiThink +1
We consider time discretization for score-based diffusion models to generate samples from a learned reverse-time dynamic on a finite grid. Uniform and hand-crafted grids can be suboptimal given a budget on the number of time steps. We introduce Adaptive Reparameterized Time (ART), which controls the clock speed of a reparameterized time variable to redistribute computation along the sampling trajectory while preserving the terminal time, with the objective of minimizing the aggregate Euler discretization error. We derive a randomized companion ART-RL that recasts ART as a continuous-time reinforcement learning problem with Gaussian policies, and prove a two-directional bridge between the two: the deterministic ART optimum lifts to an optimal Gaussian policy, and conversely any optimal Gaussian policy must recover the ART control through its mean. This bridge turns continuous-time actor--critic learning into a principled, rather than heuristic, route to the deterministic timestep optimum. Within the official EDM pipeline, ART-RL improves FID on CIFAR--10 across a wide range of budgets; after one-time offline training, the distilled deterministic schedule transfers without retraining to AFHQv2, FFHQ, and ImageNet at no extra inference cost.
Yilie Huang, Wenpin Tang, Xunyu Zhou
Department of Industrial Engineering and Operations Research, Columbia University, New York, NY 10027, USA. · Department of Industrial Engineering and Operations Research & Data Science Institute, Columbia University, New York, NY 10027, USA.