OP-GRPO: Efficient Off-Policy GRPO for Flow-Matching Models

1. Problem Statement

Flow-matching models have become a dominant paradigm for generative AI (e.g., Stable Diffusion 3.5, Wan), offering superior sampling efficiency and visual fidelity compared to traditional diffusion models. However, aligning these models with human preferences using Group Relative Policy Optimization (GRPO) faces significant efficiency bottlenecks:

• Low Sample Efficiency: Standard GRPO is strictly on-policy. It discards all generated trajectories after a single training iteration, even if they are high-quality. This leads to massive computational waste, as trajectory sampling constitutes a major portion of the cost.
• Reward Degeneration: In difficult tasks, models often fail to produce successful trajectories, causing rewards to collapse to near-zero. This results in vanishing gradients and stalled learning.
• Challenges of Off-Policy Learning: Simply reusing old trajectories (off-policy) in flow-matching models introduces two critical issues:
  1. Distributional Shift: Reusing data from an older policy ($\pi_{off}$) instead of the current policy ($\pi_{old}$) distorts the importance sampling ratios, breaking the clipping mechanism that ensures training stability in GRPO.
  2. Numerical Instability: As the denoising process approaches the final steps (low noise), the transition distribution becomes nearly deterministic. This causes importance weights to become ill-conditioned (extremely large or small), leading to unstable updates.

2. Methodology: OP-GRPO

The authors propose OP-GRPO, the first off-policy GRPO framework tailored specifically for flow-matching models. The framework consists of three core components:

A. High-Quality Replay Buffer

Instead of discarding trajectories, OP-GRPO maintains a replay buffer ($\mathcal{B}_{off}$) to store and reuse high-quality samples.

• Selection Strategy: During rollout, the system selects the highest-reward trajectory from a group to retain in the buffer.
• Uniqueness & Decay: To prevent the buffer from collapsing into repetitive samples, it enforces a uniqueness constraint (one trajectory per conditioning prompt) and applies stepwise reward decay. Older samples become easier to replace, ensuring the buffer remains close to the current policy distribution.
• Hybrid Rollout: In each training batch, prompts are sampled from the dataset, but a portion of the generated trajectories are replaced by samples from the replay buffer (e.g., $G-1$ new samples + 1 buffer sample).

B. Sequence-Level Importance Sampling Correction

Directly applying standard importance sampling to off-policy data in GRPO fails because it cancels out the reference policy ($\pi_{old}$) in the ratio, causing the clipping mechanism to measure the wrong discrepancy (between current policy and old off-policy data).

• The Solution: OP-GRPO retains the standard per-step importance ratio between the current policy ($\pi_\theta$) and the reference policy ($\pi_{old}$) to preserve the clipping logic.
• Correction Term: To compensate for the distribution shift introduced by the off-policy data, a sequence-level correction term is applied as a multiplicative factor to the objective function:
  $$\frac{P_{\pi_{old}}(\tau)}{P_{\pi_{off}}(\tau)} = \prod_{t=1}^{T} \frac{p_{\theta_{old}}(z_{t-1} | z_t, c)}{p_{\theta_{off}}(z_{t-1} | z_t, c)}$$
• Result: This approach reduces the clipping rate of off-policy samples from >40% (in naive approaches) to 11.8%, significantly improving sample utilization while maintaining stability.

C. Trajectory Truncation for Numerical Stability

The authors identify that the ill-conditioning of importance weights occurs primarily in the late denoising steps (low noise regime) where the transition distribution becomes deterministic.

• Truncation Strategy: The off-policy trajectory is only used up to a specific truncation step ($t_{off}$). For the remaining steps ($t_{off} \to 0$), the trajectory is re-generated using the current policy.
• Benefit: This eliminates the numerically unstable log-probability terms in the low-noise regime without significantly impacting perceptual quality, as late steps contribute minimally to visual fidelity.

3. Key Contributions

1. First Off-Policy GRPO for Flow-Matching: Introduces a framework that breaks the on-policy constraint of GRPO for flow-matching models, enabling trajectory reuse.
2. Stable Distribution Correction: Proposes a sequence-level importance sampling correction that decouples the distribution shift compensation from the GRPO clipping mechanism, ensuring stable policy updates.
3. Theoretical & Empirical Truncation: Demonstrates theoretically and empirically that late denoising steps cause numerical instability in off-policy settings and proposes a truncation strategy to mitigate this.
4. Efficiency Gains: Achieves comparable or superior performance to on-policy baselines using significantly fewer training steps.

4. Experimental Results

The method was evaluated on Stable Diffusion 3.5-Medium (SD3.5-M) for image generation and Wan2.1-1.4B for video generation across three tasks: Compositional Image Generation, Visual Text Rendering, and Human Preference Alignment.

• Training Efficiency: OP-GRPO requires only 34.2% of the training steps needed by Flow-GRPO (on-policy) to reach the same final performance. In video generation, this efficiency gap widens to 30.1% of the training time.
• Performance:
  • OP-GRPO consistently outperforms Flow-GRPO at intermediate training stages.
  • It achieves comparable or superior final performance (e.g., higher PickScore, better OCR accuracy, and improved aesthetic scores) compared to the on-policy baseline.
  • Generalization: Models trained with OP-GRPO show better generalization on unseen metrics (e.g., T2I-CompBench++) and do not overfit to specific task rewards.
• Ablation Studies:
  • Removing the sequence correction leads to training divergence.
  • Removing trajectory truncation causes instability due to ill-conditioned weights.
  • An off-policy sample proportion of ~25% was found to offer the best balance between convergence speed and stability.

5. Significance

This work addresses a critical bottleneck in the post-training of large-scale generative models. By enabling off-policy learning in flow-matching models, OP-GRPO drastically reduces the computational cost (GPU hours) required for alignment. It makes the training of high-fidelity image and video models more scalable and accessible. Furthermore, the proposed techniques for handling distribution shift and numerical instability in continuous generative processes provide a robust theoretical foundation for future reinforcement learning applications in diffusion and flow-based models.