Source-linked AI summary
Step Back to Move Forward: Reflection-Aware Preference Optimization for Visual Generation
Junlong Wu, Jiuzhou Lin, Jia Sun, Boheng Zhang, Huaiqing Wang, Dewen Fan, Houde Liu, Qianqian Gan, Fan Yang, Tingting Gao
TL;DR
Policy-gradient RL for diffusion models can explore inefficiently and become vulnerable to poor local optima. RA-GRPO reflects on sampled trajectories and distills the resulting paths into training; experiments report improvements over RL baselines across image and video generation.
Problem
Policy-gradient RL can be restricted to the current generative manifold, and GRPO’s cold-start groups may offer only low-quality relative positives, limiting effective exploration.
Source: §1 · pp. 1–2Method
RA-GRPO uses Diffusion Reflection to refine sampled trajectories and Counterfactual Path Synthesis to convert them into policy-training supervision.
Source: §1 · p. 2Results
Takeaways & Limitations
The authors present reflection-driven exploration as a scalable paradigm for alignment and identify diverse weak-to-strong estimator pairs as future work.
Source: §5 · p. 8The reflection direction assumes stronger text guidance produces samples better aligned with the prompt and closer to the target conditional data distribution.
Source: §3.2 · p. 4
Abstract
from arXiv · showhide
Diffusion models have become the mainstream paradigm for modern visual generation and have substantially advanced multimedia content synthesis, especially in text-to-image and text-to-video tasks. To further align such generative models with human preferences, reinforcement learning (RL) has recently shown strong potential as a post-training strategy. Nevertheless, existing policy gradient-based methods often explore inefficiently, making them vulnerable to local optima that may degrade semantic faithfulness and visual realism. To address these challenges, we present Reflection-Aware GRPO (RA-GRPO), a new RL-based preference alignment framework for diffusion generative models. The core idea is to improve "forward" generation by incorporating "backward" reflection during optimization. We first introduce Diffusion Reflection, which rectifies intermediate sampling trajectories by inverting the diffusion process with a weak estimator, guiding latent states toward higher-probability regions of the true data manifold. Furthermore, we introduce Counterfactual Path Synthesis to implicitly distill these rectified trajectories into the policy, enabling the model to internalize the benefits of search-based exploration without incurring inference-time overhead. Extensive experiments on T2I and T2V models demonstrate that RA-GRPO significantly outperforms existing methods, particularly in mitigating reward hacking and improving generalization. The method remains architecture-agnostic and integrates seamlessly with standard pipelines, suggesting a promising direction for stable preference alignment.
1 Introduction
The introduction identifies inefficient exploration and weak cold-start signals as obstacles to preference alignment with diffusion models. RA-GRPO combines trajectory reflection with counterfactual supervision, and is reported to outperform strong baselines in alignment efficiency and generation quality.
1 Introduction
- Existing policy-gradient RL is largely limited to the model’s current generative manifold, while GRPO’s cold-start groups may contain only low-quality samples, making relative positives ineffective and gradients noisy.These conditions can leave exploration inefficient and optimization prone to poor local optima.Source: §1 · pp. 1–2
- Diffusion Reflection alternates denoising and inversion to redirect sampled latent trajectories toward regions better aligned with the data distribution and associated with higher rewards.The introduction says this improves sample quality, alleviates current-policy exploration bias, and expands available trajectories.Source: §1 · p. 2
- Counterfactual Path Synthesis turns reflection-refined trajectories into coherent training paths associated with the original generation context.This implicit distillation lets the policy learn reflection’s benefits without additional inference cost.Source: §1 · p. 2
- RA-GRPO consistently outperforms strong baselines in alignment efficiency and generation quality.The introduction reports extensive experiments validating the method.Source: §1 · p. 2
2 Related Work
Related work covers RL methods for aligning generative models and reflective methods for improving generation. The paper distinguishes its approach by aiming to make reflection a persistent training-time signal.
2.1 RL-based Alignment for Generative Models
- GRPO and other RL alignment methods optimize generative models through reward-driven denoising, while structural variants target sampling efficiency and temporal credit assignment.Prior work also includes offline preference optimization adapted from DPO.Source: §2.1 · p. 2
2.2 Self-Reflection in Generative Models
- Existing visual-generation reflection methods refine outputs at inference time or use feedback for fine-tuning, but typically do not turn reflection computation into a persistent training signal.Source: §2.2 · p. 2
3 Method
RA-GRPO combines diffusion-based reflection with policy optimization: it corrects selected sampling trajectories, distills those corrections into the policy, and uses local updates intended to mitigate reward hacking.
3 Method
- The method section introduces RA-GRPO as a framework that combines trajectory refinement, counterfactual supervision, and an analysis of reward-hacking mitigation.Its stated components are Diffusion Reflection, Counterfactual Path Synthesis, and discussion of how the framework may reduce reward hacking.Source: §3 · pp. 2–3
3.1 Preliminaries: RL for Diffusion Models
- The preliminaries use stochastic sampling to obtain diverse rollouts for group-relative advantages, since deterministic Flow Matching sampling offers limited exploration for GRPO.The stochastic differential equation uses a predefined noise schedule and a correction term to preserve the deterministic flow's marginal distributions.Source: §3.1 · p. 3
3.2 Active Exploration via Diffusion Reflection
- Randomized Single-Step Reflection applies a strong-forward, weak-inverse correction at one sampled timestep to selected trajectories, extending exploration beyond the current policy support.Only a fraction of each group is reflected, limiting the added computational overhead; the weak-to-strong correction assumes stronger guidance better aligns with the prompt.Source: §3.2 · p. 4
3.3 Implicit Distillation via Counterfactual Path Synthesis
- Counterfactual Path Synthesis stitches a reflection-refined trajectory into a counterfactual training path and supervises the policy at the reflection timestep.The advantage-weighted target is intended to help the policy reproduce high-quality trajectories at inference without explicit reflection.Source: §3.3 · pp. 4–5
3.4 Why Reflection Can Mitigate Reward Hacking
- RA-GRPO combines standard GRPO with a reflection loss that locally corrects trajectories, a design the paper links to stronger robustness and cross-reward generalization.The paper frames the reflection correction as a way to improve reward alignment while retaining locally consistent denoising updates.Source: §3.4 · p. 5, §3.3 · p. 5
4 Experiment
RA-GRPO is evaluated against diffusion and reinforcement-learning baselines on image and video generation, with results covering reward alignment, generation quality, ablations, human preference, and reflection settings.
4.1 Experiments Setup
- The image experiments use FLUX.1-Dev on HPS v2.1 prompts, comparing RA-GRPO with the base model, DanceGRPO, and MixGRPO across five alignment and quality metrics.The metrics are HPSv2.1, PickScore, CLIPScore, ImageReward, and HPSv3.Source: §4.1 · p. 5
4.2 Implementation Details
- Reflection uses guidance levels 3.5 and 1.0 for the strong and weak estimators, respectively, at a randomized timestep; RA-GRPO trains for 300 steps.Other methods train for 5% more steps; training uses AdamW at 1 × 10−5 learning rate and 1 × 10−4 weight decay on 8 NVIDIA H800 GPUs.Source: §4.2 · p. 5
4.3 Main Results
- RA-GRPO ranks first across evaluated metrics under multi-objective optimization and shows the strongest overall performance under both HPSv2.1 and PickScore optimization.Table 1 compares reward objectives using multiple automatic metrics; the reported pattern suggests RA-GRPO's gains generalize better across evaluators than DanceGRPO's target-reward gains.Source: §4.3 · p. 5, Table 1 · p. 6
- For atypical compositions and spatial prompts, RA-GRPO yields more coherent structures, accurate symbol arrangements, and stronger semantic alignment than baselines.Figure 2 compares Flux.1-dev, DanceGRPO, MixGRPO, and RA-GRPO across prompts A–D.Source: §4.3 · pp. 5–6, Figure 2 · p. 6
4.4 Ablation Study
- Removing Reflection brings performance near baseline, while removing Counterfactual Synthesis slightly raises the target reward but degrades several generalization metrics.The full model performs best across all metrics; the ablation attributes the reflection effect to weak-to-strong guidance and notes a mismatch when directly optimizing reflected states.Source: §4.4 · p. 7, Table 2 · p. 7
4.5 Human Preference Evaluation
- 73.12% text-alignment, 69.35% image-quality, and 67.91% aesthetic-preference win rates favor Ours over DanceGRPO in 360 randomized pairwise judgments.Qualitatively, DanceGRPO shows abnormal lighting and highlights under HPSv2.1, and dark tones with blurry contours under PickScore; Ours is described as more natural and coherent.Source: §4.5 · p. 7, Figure 3 · p. 8
4.6 Optimal Configurations
- A reflection ratio of r = 0.5 gives the best overall performance, while reflecting every trajectory slightly degrades quality.The paper attributes the trade-off to partial reflection retaining unmodified trajectories for diverse, stable groupwise comparisons.Source: §4.6 · p. 7, Table 3 · p. 7
4.7 Extension to T2V Model
- On Wan2.1, RA-GRPO outperforms the vanilla model and DanceGRPO in both video-level and first-frame image-level evaluation protocols.The video-level setting reports gains across visual quality, motion quality, and text alignment; the image-level setting evaluates first-frame generation and reports gains across five image metrics.Source: §4.7 · pp. 7–8, Table 4 · p. 7
5 Conclusion
The conclusion presents RA-GRPO as reflection-driven exploration for diffusion-model reinforcement learning, with reported gains across image and video generation and future work on broader estimator pairs.
5 Conclusion
- RA-GRPO uses weak-to-strong estimator discrepancy to guide exploration and synthesizes counterfactual paths for training, rectifying sampling trajectories.The conclusion says this approach navigates the optimization landscape more efficiently and addresses the cold-start dilemma by discovering high-probability data modes early in training.Source: §5 · p. 8
- Experiments across image and video generation show consistent improvements in visual quality, semantic alignment, and temporal coherence over existing reinforcement-learning baselines.The conclusion also reports that RA-GRPO mitigates reward hacking.Source: §5 · p. 8
- The authors identify testing diverse weak-to-strong estimator pairs as future work to broaden the framework's applicability.Source: §5 · p. 8