PaperlayerView exact revision on arXiv

Source-linked AI summary

UniGRPO: Unified Policy Optimization for Reasoning-Driven Visual Generation

Jie Liu, Zilyu Ye, Linxiao Yuan, Shenhan Zhu, Yu Gao, Jie Wu, Kunchang Li, Xionghui Wang, Xiaonan Nie, Weilin Huang, Wanli Ouyang

arXiv:2603.23500v1cs.CV

TL;DR

Interleaved multimodal generation needs a unified RL framework that jointly trains reasoning and image-generation policies. UniGRPO models a reasoning-driven image-generation round as one MDP, combining GRPO and FlowGRPO with CFG-free rollouts and velocity-based regularization. Experiments report improved image quality and establish a scalable baseline, while the current single-round, sparse-reward setting leaves multi-round scaling and reasoning credit assignment for future work.

  • Problem

    Interleaved generation requires a unified RL framework that jointly optimizes text and image generation policies.

  • Method

    UniGRPO models Prompt → Thinking → Image as one MDP, combining GRPO for reasoning, FlowGRPO for synthesis, CFG-free rollouts, and velocity-field MSE regularization.

  • Results

    UniGRPO improves image-generation quality through reasoning and provides a robust, scalable baseline for future post-training of fully interleaved models.

  • Takeaways & Limitations

    The recipe offers a unified training approach for reasoning-driven image generation and is intended to extend toward fully interleaved models.

  • Takeaways & Limitations

    The framework is currently validated on single-round generation with sparse terminal rewards, which can produce inefficient credit assignment for intermediate reasoning.

Abstract

from arXiv · show

Unified models capable of interleaved generation have emerged as a promising paradigm, with the community increasingly converging on autoregressive modeling for text and flow matching for image generation. To advance this direction, we propose a unified reinforcement learning framework tailored for interleaved generation. We validate our approach on its fundamental unit: a single round of reasoning-driven image generation, where the model first expands the user prompt through reasoning, followed by image synthesis. Formulating this multimodal generation process as a Markov Decision Process with sparse terminal rewards, we introduce UniGRPO to jointly optimize text and image generation policies using GRPO. Adopting a minimalist methodology to avoid over-design, we leverage established training recipes for both modalities by seamlessly integrating standard GRPO for reasoning and FlowGRPO for visual synthesis. To ensure scalability to multi-round interleaved generation, we introduce two critical modifications to the original FlowGRPO: (1) eliminating classifier-free guidance to maintain linear, unbranched rollouts, which is essential for scaling to complex scenarios involving multi-turn interactions and multi-condition generation (e.g., editing); and (2) replacing the standard latent KL penalty with an MSE penalty directly on the velocity fields, providing a more robust and direct regularization signal to mitigate reward hacking effectively. Our experiments demonstrate that this unified training recipe significantly enhances image generation quality through reasoning, providing a robust and scalable baseline for the future post-training of fully interleaved models.

1 Introduction

UniGRPO is a unified RL framework that models reasoning-driven image generation as one MDP and jointly optimizes language reasoning with visual synthesis. It introduces FlowGRPO adaptations aimed at scalable, robust training and reports improved image-generation quality.

  • Framework: The framework validates unified RL on a single round of reasoning-driven image generation, covering both text and image generation.This setting is presented as the fundamental unit of interleaved generation.
  • Framework: UniGRPO formulates the Prompt → Thinking → Image sequence as a single MDP for joint text-and-image policy optimization.It combines standard GRPO for reasoning with FlowGRPO for visual synthesis under sparse terminal rewards.
  • Scalable Flow Matching RL Adaptations: Removing classifier-free guidance keeps rollouts linear and unbranched, supporting future multi-turn and multi-condition generation.The design targets scenarios such as complex image editing.
  • Scalable Flow Matching RL Adaptations: Replacing the latent KL penalty with velocity-field MSE provides a more direct regularization signal for mitigating reward hacking.This modification is part of the FlowGRPO adaptations introduced for scalable training.
  • Framework: UniGRPO jointly optimizes autoregressive reasoning and flow-matching image policies to improve image generation quality.The paper positions this as a minimalist unified training recipe for reasoning-driven image generation.
  • Results: The unified training recipe effectively optimizes the model under sparse terminal rewards and establishes a baseline for future fully interleaved-model post-training.The reported contribution is a robust and scalable training baseline.

2 Related Work

Related work spans RL for language reasoning, reward-based optimization of text-to-image models, unified multimodal architectures, and concurrent RL approaches for joint generation.

  • RL for LLMs: GRPO removes the value model through group-relative baselines and is used to efficiently optimize reasoning-intensive language models.UniGRPO adapts GRPO to intermediate thinking tokens before visual synthesis.
  • RL for Text-to-Image Models: Text-to-image alignment has used reward optimization, RWR, DPO, PPO-style policy gradients, and training-free guidance methods.Flow-matching architectures require specific stochastic formulations when adapting established RL paradigms.
  • Unified Multimodal Models: Unified multimodal research includes vector-quantized models that place image and text tokens in a shared autoregressive training space.The cited examples include Chameleon, Emu3, and VILA-U.
  • Concurrent Multimodal RL: Concurrent RL methods differ in their rollout or model designs, with some using benchmark-specific prompts or tree-structured separate backbones.The passage contrasts these designs with true interleaved multimodal generation.

3 Preliminary

The preliminary framework combines standard GRPO for autoregressive text with FlowGRPO for flow-matching visual generation, using stochastic exploration and ratio normalization to stabilize diffusion-policy updates.

  • Text GRPO: Standard GRPO maximizes expected reward while constraining autoregressive policy updates with importance-sampling clipping.
  • Text GRPO: Group-relative advantages compare each generated text sample against the rewards of its group.
  • Flow GRPO: FlowGRPO adapts reinforcement learning to flow-matching models by converting deterministic ODE sampling into stochastic SDE sampling for exploration.
  • Flow GRPO: FlowGRPO-Fast restricts gradient-tracked optimization to an SDE time window while using ODE sampling elsewhere to reduce computational overhead.
  • Flow GRPO: RatioNorm standardizes left-shifted, variable-variance log-importance ratios so clipping can constrain overconfident positive updates.

4 Method

UniGRPO formulates reasoning-driven image generation as one sequential MDP and jointly optimizes text and image policies under sparse terminal rewards. It combines group-relative policy optimization with CFG removal and velocity-field MSE regularization to support scalable rollouts and mitigate reward hacking.

  • MDP Formulation: UniGRPO treats each text-token prediction and image-denoising step as one action in a sequential MDP.Text states track the prompt and reasoning prefix; image states additionally track the completed reasoning, noisy latent, and flow time.
  • MDP Formulation: Text and image transitions are deterministic given their actions, respectively appending a token or advancing the latent to the next flow step.
  • MDP Formulation: A sparse terminal reward is assigned only after the image latent is fully denoised, while intermediate steps receive zero reward.
  • UniGRPO Optimization: For each prompt, UniGRPO samples reasoning chains, generates corresponding image trajectories, computes group-relative advantages from terminal rewards, and updates one unified policy.
  • Scalable Training Adaptations: UniGRPO sets the image-objective weight λ to 1 across experiments to balance reasoning and synthesis, while removing CFG to preserve linear, unbranched rollouts.CFG removal is intended to support multi-round and multi-condition generation.
  • Scalable Training Adaptations: Replacing timestep-weighted latent KL with unweighted velocity-field MSE keeps the RL-tuned vector field close to the reference model across noise levels.The paper reports that this uniform regularization leaves fewer exploitation loopholes and more effectively mitigates reward hacking.

5 Experiments

Experiments evaluate UniGRPO against multiple baselines on text alignment and compositional generation, then analyze qualitative outputs and ablate CFG and regularization. UniGRPO achieves the strongest reported benchmark results while velocity-field MSE supports stable, high-fidelity training.

  • Experimental Settings: The evaluation uses a 150-prompt Text Alignment benchmark with four images per prompt and GenEval for compositional capabilities.TA scores average binary exam-point assessments across all evaluated images; GenEval covers counting, spatial relations, and attribute binding.
  • Benchmark Performance: UniGRPO achieves 0.8381 on TA and 0.90 on GenEval, reaching state-of-the-art performance among the evaluated RL methods.UniFPO failed to converge, while joint optimization outperformed optimizing reasoning or synthesis alone.
  • Qualitative Analysis: UniGRPO produces photorealistic, finely detailed images with stronger aesthetic quality and text-image alignment than the original Bagel and SFT models.The original model shows oversaturation and artifacts, while SFT reduces artifacts but introduces blurriness.
  • Qualitative Analysis: UniGRPO reasoning traces become purposeful and task-oriented, tightly coupling intermediate thoughts with subsequent image synthesis.The joint RL stage addresses the gap between extensive textual reasoning and fine-grained visual detail translation observed after SFT.
  • Impact of CFG-Free Training: Removing CFG during training yields comparable or better final performance when CFG is applied at evaluation, making branched training rollouts unnecessary.The ablation uses the original Bagel model at 512 resolution, so its results are not directly comparable with Figure 3’s 1024-resolution curves.
  • Regularization Strategies: Velocity MSE provides the best regularization balance, whereas no KL causes reward hacking and latent KL still produces grid-like artifacts.The latent-KL run was terminated early after artifacts emerged around step 250 despite a sufficiently large KL coefficient.

6 Conclusion and Future Work

UniGRPO unifies reinforcement learning for interleaved text-and-image generation and improves image quality through chain-of-thought reasoning. The paper establishes a scalable baseline while identifying multi-round generation and denser process rewards as future directions.

  • Conclusion: UniGRPO formulates multimodal generation as an MDP integrating autoregressive reasoning with flow-based visual synthesis in one optimization loop.The framework eliminates CFG for linear rollouts and uses velocity-based regularization to mitigate reward hacking.
  • Conclusion: UniGRPO enhances image generation quality through chain-of-thought reasoning and provides a scalable baseline for future post-training of fully interleaved models.Figure 6 depicts task-oriented reasoning guiding faithful, photorealistic synthesis.
  • Scaling to Multi-Round Interleaved Generation: The current validation is limited to single-round Prompt → Reasoning → Image generation, with future work targeting interactive editing, visual storytelling, and multi-turn dialogue.These settings require maintaining context consistency across longer horizons.
  • Multimodal Process Reward Modeling: UniGRPO currently uses sparse terminal rewards from the final image, which can assign credit inefficiently when reasoning is flawed but the image is high quality.Multimodal process reward models are proposed to provide dense feedback on intermediate reasoning and improve sample efficiency and interpretability.

A.1 Detailed GenEval Results

This section presents GenEval results with fine-grained metrics across six categories and identifies how evaluated RL methods are represented in the comparison.

  • GenEval is broken down into Single Object, Two Objects, Counting, Colors, Position, and Attribute Binding categories.
  • Table 2 reports quantitative results for all evaluated models and baselines, with RL methods initialized from the Bagel checkpoint after SFT.
  • The table marks whether each method explicitly generates intermediate reasoning tokens and uses × to denote training collapse.

A.2 Qualitative Results Across Training Stages

The qualitative results compare reasoning traces and generated images across training stages. They indicate that RL training improves the coherence of reasoning and its translation into visual details after SFT remains suboptimal.

  • The qualitative comparison covers reasoning traces and generated images across different training stages.
  • After SFT, the model produces extensive reasoning, but translating textual thoughts into fine-grained visual details remains suboptimal.
  • The subsequent RL-based training stage effectively addresses this gap, yielding more coherent reasoning traces.

B Implementation Details and Hyperparameters

This section documents UniGRPO’s joint text-and-image training configuration and provides qualitative examples of prompt expansion, reasoning, and generated-image inputs.

  • Table 3 lists joint-training hyperparameters for the TextGRPO and FlowGRPO experts and the joint objective weight λ.
  • The examples pair prompts with intermediate thinking texts and resulting images as part of the implementation materials.
  • The paired prompt descriptions include six cups arranged in two rows of three and a boy holding wood while a girl draws on it.
  • Additional examples describe a three-platform climbing frame occupied by a cat, dog, and mouse, plus a tiger chasing two sheep on grassland.
  • Figures 7 and 8 provide baseline Bagel and post-SFT reasoning traces with generated images, while Figure 8 notes suboptimal visual-detail translation after SFT.
  • Reasoning examples expand prompts by specifying scene structure, viewpoint, style, composition, and visual details.
Loading 2603.23500v1…