UniGRPO: Unified RL for Multimodal Models

• UniGRPO is a unified reinforcement learning algorithm that integrates policy gradients for both textual reasoning and image synthesis under a single training loop.
• It employs group-relative rewards and variance reduction techniques to ensure robust credit assignment even with sparse or terminal reward signals.
• Empirical results demonstrate superior performance with a text alignment score of 0.838 and a GenEval score of 0.90 compared to traditional RL fine-tuning methods.

UniGRPO (Unified Group-Relative Policy Optimization) is a reinforcement learning (RL) algorithm that unifies policy-gradient post-training across both reasoning (textual and multimodal) and generative (text-to-image) tasks within multimodal diffusion architectures. The method extends Group-Relative Policy Optimization (GRPO) to operate natively over discrete diffusion trajectories, allowing a single policy and training loop to optimize for diverse task objectives using reward functions differentiated only by end-goal metrics. UniGRPO enables joint and efficient post-training of unified diffusion models for interleaved generation, making it central to advanced frameworks such as MMaDA and reasoning-driven visual generation models (Yang et al., 21 May 2025, Liu et al., 24 Mar 2026).

1. Motivation and Conceptual Foundations

UniGRPO addresses the longstanding challenge of unified post-training for foundation models that span both autoregressive and diffusion-based generation modalities. Traditional RLHF and PPO/GRPO methods are typically tailored to autoregressive LLMs, while ad hoc approaches to diffusion often suffer from inefficiencies or lack seamless alignment to diverse end tasks. UniGRPO adapts policy-gradient RL to the unique structure of high-dimensional diffusion and flow-based models, providing a scalable mechanism for end-to-end learning over reasoning (text, multimodal) and synthesis (images) branches. This is achieved by:

• Treating diffusion models as stochastic policies over masked-token or flow trajectories.
• Defining task-specific rewards that are directly incorporated into the same clipped-surrogate plus divergence-penalty loss.
• Using group-relative rewards for variance reduction and robust credit assignment, even under sparse or terminal-only reward conditions (Yang et al., 21 May 2025, Liu et al., 24 Mar 2026).

2. Formalization as Unified RL Objective

The UniGRPO objective generalizes PPO/GRPO to unified multimodal diffusion settings. For a prompt $q$, output $o$, random mask ratio $p_i$, and batch of $G$ rollouts, the optimization objective is:

$$\mathcal{J}_\text{UniGRPO}(\theta) = \mathbb{E}_{q,\{o_i\},\{p_i\}} \left[ \frac{1}{G}\sum_{i=1}^G \frac{1}{|o_i|}\sum_{t=1}^{|o_i|} \min \left( r'_{i,t}(\theta) \hat{A}_{i,t}, \mathrm{clip}(r'_{i,t}(\theta),1-\varepsilon,1+\varepsilon)\hat{A}_{i,t} \right) - \beta D_\mathrm{KL}(\pi'_\theta \| \pi'_\mathrm{ref}) \right]$$

where $r'_{i,t}(\theta)$ is the per-token (or per-timestep) importance ratio, $\hat{A}_{i,t}$ the group-relative advantage, $\varepsilon$ the clipping factor, and $\beta$ the KL penalty towards a reference policy. Task alignment is exclusively determined by the reward $R^\mathrm{Uni}(o)$ for each output modality. Full derivation and gradient estimator details are provided in (Yang et al., 21 May 2025), with an analogous formulation for continuous flow SDEs in (Liu et al., 24 Mar 2026).

3. Algorithmic Implementation

3.1 Core Steps

The following summarizes the typical UniGRPO implementation flow (see pseudocode in (Yang et al., 21 May 2025, Liu et al., 24 Mar 2026)):

1. Initialize model weights from a supervised or SFT checkpoint producing coherent chain-of-thought (CoT) traces.
2. For each batch, sample prompts and generate $o$0 completions/trajectories per prompt.
3. Compute task-specific rewards and normalize advantages across the group for each prompt.
4. Sample random mask ratios or SDE windows, covering different stages of the denoising process.
5. Perform local updates to minimize the surrogate RL loss with KL (or velocity-MSE in flow settings) regularization, using a single policy across reasoning and generative modalities.
6. Iterate to convergence, typically validated using domain metrics (correctness, preference scores, CLIP alignment).

3.2 Diffusion and Flow Settings

• Text/multimodal: Discrete denoising steps, masking tokens at random positions/rates.
• Image/flow: Flow-matching SDEs, with RL updates in a sub-window, and stochastic velocity perturbations for exploration.
• Rollout Efficiency: Forward passes are performed per mask/SDE ratio, with no need to enumerate combinatorial Monte Carlo samples, contributing to stability and scalability (Yang et al., 21 May 2025).

4. Key Design Adaptations and Scalability Features

UniGRPO incorporates several crucial adaptations for efficiency and scalability in unified generation settings:

Feature	Description	Source
Group-relative reward	Reduces RL variance; enables robust updates from sparse signals	(Yang et al., 21 May 2025)
Single policy/state	Unified $o$1 applies to both AR (text) and diffusion (image)	(Yang et al., 21 May 2025, Liu et al., 24 Mar 2026)
No classifier-free guidance (CFG) in training	Eliminates quadratic rollout complexity for multi-round generation	(Liu et al., 24 Mar 2026)
Velocity-MSE (vs. latent-KL)	Uniform MSE regularization on velocity fields, preventing reward hacking	(Liu et al., 24 Mar 2026)

Classical classifier-free guidance (CFG) is omitted during RL updates to maintain linear, unbranched trajectory graphs. The stepwise KL penalty typically used in diffusion RL is replaced by a velocity-MSE loss, which provides uniform regularization across diffusion noise schedules and is empirically shown to mitigate reward hacking and failure modes associated with latent-KL weighting (Liu et al., 24 Mar 2026). This suggests UniGRPO is well-suited to multi-round, multi-condition generation regimes.

5. Diversified Reward Modeling and Objectives

UniGRPO enables seamless integration of diversified reward signals, unified within a single policy-gradient backpropagation. The algorithm supports:

• Textual reasoning: Numeric correctness (+2.0), format adherence (+0.5).
• Multimodal reasoning: Combination of correctness, format, and CLIP-based alignment.
• Text-to-image generation: CLIP image-text alignment, human preference scores via ImageReward or VLM-based reward models (Yang et al., 21 May 2025, Liu et al., 24 Mar 2026).

All reward types are plugged into the same policy gradient loss, reinforcing alignment across all modalities and enabling generalization beyond modality-specific paradigms.

6. Experimental Results and Benchmarks

UniGRPO achieves empirically strong results on both reasoning and generative tasks. For example, in the MMaDA framework (Yang et al., 21 May 2025), UniGRPO post-training lifts performance across text, multimodal, and image benchmarks in the following configuration: 50K RL steps, 64xA100-80GB GPUs, batch size 1280, AdamW optimizer (initial lr $o$2), and diffusion horizon $o$3. In reasoning-driven visual generation (Liu et al., 24 Mar 2026), UniGRPO outperforms previous RL fine-tuning (ReFL), flow-only or text-only GRPO, and other unified FPO variants on text alignment and generalization metrics:

Method	Thinking?	Text Alignment (TA)	GenEval
Bagel (vanilla)	✗	0.681	0.78
Bagel + Thinking	✓	0.713	0.79
SFT (no RL)	✓	0.777	0.82
ReFL	✗	0.779	0.85
FlowGRPO	✗	0.811	0.88
TextGRPO	✓	0.808	0.88
ReFL + TextGRPO	✓	0.799	0.87
UniGRPO	✓	0.838	0.90

Ablations confirm the negligible effect of omitting CFG in training and the superiority of velocity-MSE regularization over latent-KL or absence of regularization (which leads to reward hacking or collapse) (Liu et al., 24 Mar 2026).

7. Significance and Impact

UniGRPO establishes a general, practical baseline for unified RL post-training of multimodal diffusion models, enabling joint optimization of LLM-style reasoning and diffusion-based generation under a single policy. Its adaptation of group-relative advantage normalization, task-sensitive reward modeling, and efficient rollout structure make it suitable for future development of interleaved, multi-round, and multi-condition generation architectures. Open-source implementations as part of the MMaDA and Bagel-family frameworks provide reproducibility and a platform for continued research (Yang et al., 21 May 2025, Liu et al., 24 Mar 2026).