UniGRPO: Unified Policy Optimization for Reasoning-Driven Visual Generation

ArXiv: 2603.23500

🎯 Pitch

A shared reward signal over thinking-and-painting collapses when reward hacking breaks the link, unless you ban guidance during RL and regularize on velocity fields. UniGRPO couples text reasoning and flow-matching image generation under one GRPO loop, hitting 0.90 on GenEval—but only if the thought actually shapes the image.

---

1. Executive Summary

This paper proposes UniGRPO, a unified reinforcement learning framework that jointly optimizes text and image generation policies for reasoning-driven visual generation by formulating the entire "Prompt → Thinking → Image" sequence as a single Markov Decision Process. Validated on the Bagel multimodal model using the GenEval and internal TA benchmarks, UniGRPO integrates standard GRPO for discrete text reasoning with FlowGRPO for continuous flow-matching image synthesis—coupling these two optimization processes under sparse terminal rewards derived from a VLM-based reward model. To ensure scalability to future multi-round and multi-condition generation, the framework introduces two critical modifications: eliminating classifier-free guidance during training to enforce linear unbranched rollouts, and replacing the standard latent KL penalty with an MSE penalty directly on velocity fields for more robust reward hacking mitigation. UniGRPO achieves state-of-the-art performance, reaching 0.8381 on TA and 0.90 on GenEval—substantially outperforming both single-modal baselines (FlowGRPO at 0.8208 on TA) and hybrid approaches—while establishing that joint optimization of reasoning and visual synthesis improves image generation quality only when the unified reward signal effectively couples the thought process to the final visual output.

2. Context and Motivation

The Core Problem: We Don't Know How to Jointly Optimize Multimodal Generation Policies

The fundamental problem this paper addresses is deceptively simple: how do you apply reinforcement learning to a single model that generates both text and images in an interleaved fashion? The community is rapidly converging on unified multimodal architectures — particularly the combination of autoregressive (AR) transformers for text and flow-matching models for visual synthesis — as exemplified by systems like Bagel, Show-o, Transfusion, and Mogao. These models can, in principle, produce sequences like "reason about the prompt → generate an image → reflect on the output → edit the image" in a single coherent process. But there is currently no established framework for post-training these models with RL to align their interleaved generation behavior with human preferences or task objectives.

This gap is significant for several reasons the paper highlights implicitly throughout Sections 1–3:

• Reasoning as a test-time compute lever for visual generation. Recent work (e.g., Interleaving Reasoning for Better Text-to-Image Generation) has shown that allowing models to "think" before generating images — expanding user prompts through chain-of-thought reasoning, then conditioning image synthesis on that reasoning — can substantially improve output quality. But such reasoning chains are not a free lunch: they need to be aligned so that the thinking actually helps the image generation, rather than producing verbose but irrelevant text that the visual pipeline ignores. RL provides a natural mechanism for this alignment, but no prior work had applied it jointly to both the reasoning and synthesis components within a single model.

• The multi-turn, multi-condition scaling bottleneck. The paper is explicitly forward-looking in its motivation. The authors are not just solving today's single-round "Prompt → Think → Image" problem; they are laying groundwork for complex multi-round interleaved generation (interactive editing, visual storytelling, multi-turn dialogue with images). In these scenarios, the model must maintain context consistency across alternating text and image generation phases, and standard inference techniques like Classifier-Free Guidance (CFG) — which requires multiple forward passes per denoising step and branches the computation graph with multiple conditioning contexts — become computationally prohibitive and architecturally messy. A scalable RL framework for interleaved generation must therefore address these architectural constraints from the start, not as an afterthought.

• The post-training pipeline gap. The paper positions itself at a specific inflection point: the community has built capable pretrained interleaved models (Bagel, Mogao, etc.), but these models produce suboptimal outputs out of the box — oversaturated colors, synthetic artifacts, reasoning traces disconnected from visual intent. Supervised fine-tuning (SFT) can partially address these issues (as shown in Table 1: SFT improves Bagel from 0.6810 to 0.7486 on TA), but SFT alone hits a ceiling because it trains on static demonstration data rather than optimizing for the actual reward signal that captures human preferences. RL-based post-training is the natural next step, but no unified recipe existed for doing this on models that generate both discrete text tokens and continuous image latents.

Where Existing Approaches Fall Short

The paper identifies specific limitations across three categories of prior work, which together motivate the need for UniGRPO:

Single-modal RL methods treat text and images in isolation. Standard GRPO (Shao et al., 2024) has been enormously successful for post-training LLMs on reasoning tasks (DeepSeek-R1), but it operates exclusively on discrete token sequences. Similarly, recent work has adapted GRPO to flow-matching models: FlowGRPO (Liu et al., 2025a) and DanceGRPO (Xue et al., 2025a) reformulate the deterministic flow ODE into a stochastic SDE to enable policy gradient exploration, and apply group-relative advantages to optimize image generation. But these methods are designed for single-modality optimization — they assume either a pure text generation problem or a pure image generation problem, not an interleaved sequence where the quality of generated text directly conditions the subsequent image synthesis. The paper shows empirically that optimizing only the text component (TextGRPO) or only the image component (FlowGRPO) leaves performance on the table compared to joint optimization (UniGRPO achieves 0.8381 on TA vs. 0.8208 for FlowGRPO with thinking).

Two-stage or separate-model approaches introduce fragility and misalignment. Several concurrent and prior works attempt to combine language and visual RL by training separate models in sequence or via alternating updates. For example:

• ReFL + TextGRPO (Table 1): A two-stage paradigm where the image generation component is first optimized using Reward Feedback Learning (ReFL, which backpropagates reward gradients through a randomly-picked late denoising timestep), and then the text generation module is subsequently fine-tuned with TextGRPO. This sequential approach achieves only 0.7987 on TA — worse than FlowGRPO alone with thinking (0.8208) and substantially below UniGRPO's 0.8381. The sequential optimization creates a mismatch: the image policy is optimized assuming the original (unaligned) reasoning distribution, then the reasoning policy is changed, making the image policy's optimization stale.

• DualGRPO (Kou et al., 2026): This concurrent work jointly optimizes a separate LLM model and diffusion backbone via a tree-structured rollout. The paper explicitly notes this design "is incompatible with true interleaved multimodal generation" because the language and visual components are architecturally decoupled — they cannot share representations or jointly condition on each other's outputs in a tight loop.

• SepGRPO (Jiao et al., 2025): Also built on Bagel, but proposes alternating RL between the MLLM and DiT modules — training them separately rather than jointly optimizing end-to-end. The paper positions UniGRPO in contrast: "our method is built on a single unified model, trained with general-purpose prompts at 1024 resolution, with a scalable algorithm design built upon an improved FlowGRPO."

These separate-module approaches suffer from a fundamental limitation: the language model's reasoning is optimized without direct feedback on how it affects image generation, and the image model is optimized without the ability to shape the language model's reasoning to better serve visual synthesis. The result is suboptimal coordination between the two modalities.

Gradient-based RL methods are incompatible with interleaved generation or impose restrictive requirements. ReFL (Xu et al., 2024) directly backpropagates reward gradients through a single late denoising timestep, but this approach has two critical limitations for interleaved generation: (1) it requires a differentiable reward function, which excludes the use of black-box verifiers or VLM-based judges that are often more flexible and powerful; (2) it only optimizes the image generation component, treating the reasoning text as fixed input rather than something to be jointly optimized. The paper notes that while ReFL is "restricted to such differentiable objectives, UniGRPO is compatible with a broader range of verifier-based rewards." This flexibility is crucial for scaling to complex tasks where reward signals may come from multiple, potentially non-differentiable sources.

FPO-based approaches suffer from training instability in the unified setting. FPO (Flow Matching Policy Gradients, McAllister et al., 2025) serves as an alternative to FlowGRPO's SDE-based exploration. Instead of injecting noise during the forward denoising process, FPO uses the forward (noising) process to obtain $x_t$ and treats the evidence lower bound (ELBO) of the denoising process as a surrogate for $\log p_\theta(x_0 | c)$ to compute importance sampling weights. The paper tests a unified variant called UniFPO (where the text component uses TextGRPO and the image component uses FPO), and reports that it failed to converge (Table 1, marked with "—" indicating training collapse). This is a striking finding: the specific choice of RL algorithm for the continuous visual component is not a minor detail — it determines whether joint optimization is even possible. The paper attributes this to "the stability advantages of our GRPO-based formulation," suggesting that FlowGRPO's SDE exploration + RatioNorm clipping provides more stable gradient signals when coupled with text GRPO in a unified loop.

How This Paper Positions Itself

The paper frames its contribution not as proposing a fundamentally new RL algorithm, but rather as providing the first working recipe for unified RL post-training of interleaved text-image models, with design choices explicitly motivated by scalability to future multi-round scenarios. This positioning has several key aspects:

A "minimalist methodology to avoid over-design." The paper repeatedly emphasizes this phrase (Section 1, Section 4). The approach is deliberately simple: take the standard GRPO recipe that works for LLMs, take the FlowGRPO recipe that works for flow-matching models, and combine them by formulating the entire interleaved generation as a single MDP with a unified reward signal. The $\lambda$ weighting factor between text and image objectives is set to 1 — no careful tuning, no complex scheduling. This minimalism is a strength: it demonstrates that the conceptual framework (unified MDP + group-relative advantages) is what matters, not elaborate hyperparameter engineering.

Validating on the "fundamental unit" before scaling. The paper explicitly declines to tackle full multi-round interleaved generation immediately, instead focusing on a single round of reasoning-driven image generation (Prompt → Thinking → Image). The justification (Section 1) is that this setting "already encompasses both text and image generation, covering the essential components of interleaved generation" and "serves as a meaningful and principled testbed" in the absence of open-source base models natively capable of full interleaved generation. This is a pragmatic choice that allows rigorous ablation and comparison while the architectural foundations are being validated.

Forward-looking design for multi-turn and multi-condition scenarios. Despite validating on single-round generation, the two key modifications to FlowGRPO — CFG elimination and velocity MSE regularization — are explicitly motivated by future scaling requirements rather than immediate performance gains:

• CFG elimination is described as "critical" not because it improves single-round results (the ablation in Figure 4 shows comparable or slightly better performance with CFG during training), but because "this computational burden scales with the number of conditions; for multi-condition generation such as image editing, CFG demands at least three evaluations per step" and "this complexity compounds in multi-round interleaved generation, where the system must continuously manage and branch multiple conditional contexts across alternating text and image phases." The paper is making a deliberate tradeoff: accept slightly harder optimization during training (compensating for CFG removal by internalizing alignment into the policy weights through RL) in exchange for a dramatically simpler, more scalable architecture for future use cases.

• Velocity MSE regularization is motivated by the observation that the standard latent KL divergence in flow-matching RL applies an uneven penalty across timesteps (weighted by $1/\sigma_{t_k}^2$, which becomes excessively small at high-noise timesteps), creating "temporal vulnerabilities that the RL optimizer can easily exploit." The velocity MSE penalty provides uniform regularization across all noise levels, which is more important in longer-horizon multi-turn scenarios where the model must remain stable over many alternating text-image generation steps.

This forward-looking motivation distinguishes the paper from concurrent work that optimizes for immediate benchmark performance without considering architectural scalability.

A bridge between two RL traditions that have evolved independently. The paper explicitly notes (Section 2) that "multimodal understanding and image generation have long evolved independently" — RL for LLMs developed GRPO, PPO, and related methods for discrete token sequences; RL for diffusion/flow models developed DPO variants, ReFL, and most recently FlowGRPO for continuous latent trajectories. UniGRPO bridges these traditions by showing that they can be combined in a single MDP formulation, with the same group-relative advantages driving both policy updates. This unification is conceptually important even beyond the immediate empirical results because it provides a template for how future interleaved systems should be post-trained.

3. Technical Approach

3.1 Reader Orientation

What the system is: UniGRPO is a joint reinforcement learning training procedure that takes a single pretrained multimodal model capable of generating both text reasoning and images, and fine-tunes it so that the generated reasoning actively helps produce better images, rather than being decorative or disconnected. What problem it solves and the shape of the solution: The core challenge is that standard RL methods for LLMs (discrete token optimization) and for flow-matching image models (continuous latent optimization) operate on fundamentally different action spaces, yet in an interleaved generation model, the quality of the reasoning text directly determines the conditioning signal for image synthesis — so optimizing them separately, sequentially, or with incompatible algorithms creates misalignment and instability. UniGRPO solves this by formulating the entire "Prompt → Thinking → Image" sequence as a single Markov Decision Process where both the discrete text policy and the continuous image policy receive gradient updates from the same group-relative advantage signal, computed from a sparse terminal reward on the final generated image. The framework is deliberately minimalist — it takes established recipes (standard GRPO for text, FlowGRPO for images), makes two forward-looking modifications to FlowGRPO for multi-round scalability (CFG removal and velocity MSE regularization), and shows that simply coupling them under a unified advantage signal is sufficient to achieve substantial gains over any single-modality or sequential optimization approach.

3.2 Big-Picture Architecture (Diagram in Words)

The UniGRPO training system has five major components, connected in a single unified loop:

1. Unified Multimodal Model (Bagel, after SFT): A single transformer-based architecture that handles both text token autoregression and flow-matching-based image denoising. This is the policy $\pi_\theta$ being optimized — it takes a user prompt, generates reasoning tokens, then conditions on those tokens to denoise an image latent.

2. Group Sampling Module: For a given prompt, the system samples $G$ complete multimodal trajectories (reasoning chain + generated image) from the current policy. A trajectory consists of: user prompt → autoregressively sampled reasoning tokens → SDE/ODE hybrid denoising trajectory → final decoded image.

3. Reward Model (VLM-based): A differentiable reward model (fine-tuned InternVL, from RewardDance) that takes the original prompt and the final generated image as input, and outputs a scalar score measuring text-image alignment and visual quality. This is the terminal reward $R(x_0, c)$ assigned at the end of each trajectory; all intermediate text tokens and denoising steps receive zero reward.

4. Advantage Computation Module: Takes the $G$ terminal rewards for a group, standardizes them into relative advantages $\hat{A}_i = (R_i - \text{mean})/\text{std}$, and supplies these advantages as the learning signal for both text and image policy updates.

5. Unified Policy Update: Backpropagates through both the text generation trajectory (using standard GRPO with importance sampling and clipping) and the image denoising trajectory (using FlowGRPO with SDE exploration, RatioNorm clipping, and velocity MSE regularization), with both objectives weighted equally ($\lambda = 1$).

Information flow: A batch of prompts enters → for each prompt, the model generates $G$ text reasoning chains autoregressively → each reasoning chain conditions the model's flow-matching head to generate one image via hybrid SDE/ODE denoising → the reward model scores each of the $G$ images → group-relative advantages are computed → the text policy is updated to increase the probability of reasoning tokens that led to high-reward images → the image policy is updated to increase the probability of denoising steps that led to high-reward images (while staying close to the reference model via velocity MSE) → repeat.

3.3 Roadmap for the Deep Dive

Understanding UniGRPO requires mastering several interlocking mechanisms, each building on the previous. Here is the order of explanation and why:

• First, the unified MDP formulation (Section 3.4.1): This is the conceptual foundation — how interleaved generation is formalized as a sequential decision process with distinct text and image phases, shared state, and sparse terminal rewards. Without this abstraction, the joint optimization objective makes no sense.

• Second, text GRPO (Section 3.4.2): We review the standard GRPO objective for discrete tokens, since UniGRPO uses it without modification for the reasoning phase. Understanding the importance ratio, clipping, and group-relative advantage is essential before seeing how they extend to images.

• Third, FlowGRPO for images (Section 3.4.3): This is the most technically involved component. We cover the SDE formulation that enables exploration in deterministic flow models, the RatioNorm mechanism that prevents reward hacking in the clipping regime, and the hybrid SDE-ODE sampling strategy. These are all inherited from prior work but are necessary context for understanding UniGRPO's modifications.

• Fourth, the two key UniGRPO modifications (Section 3.4.4): CFG elimination and velocity MSE regularization. These are UniGRPO's novel contributions to the FlowGRPO recipe, and understanding why they matter requires understanding what they replace and what problems the replaced components cause in multi-round settings.

• Fifth, putting it together — the unified objective and training loop (Section 3.4.5): How the text and image objectives are combined with equal weighting, how advantages are shared, and what the full training iteration looks like in practice.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a training methodology paper whose core idea is that interleaved text-image generation must be post-trained with RL that jointly optimizes both modalities under a single unified advantage signal, using a deliberately minimalist combination of established per-modality recipes with two forward-looking modifications to the image recipe.

---

3.4.1 The Unified MDP: Formalizing Interleaved Generation as Sequential Decision-Making

The paper's foundational contribution is the formulation of interleaved generation as a single Markov Decision Process (MDP) defined by the tuple $(\mathcal{S}, \mathcal{A}, P, R)$, where each step $k$ corresponds to either a single text token prediction or a single image denoising step. This formulation is essential because it provides a common mathematical language for the discrete and continuous components, enabling them to share a single advantage signal and a single optimization loop.

State Space $\mathcal{S}$. The state evolves through two distinct phases, with different representations for text generation versus image synthesis:

• Text phase state: $s_k^{\text{txt}} = (c, y_{<k})$, where $c$ is the input user prompt and $y_{<k}$ is the sequence of all previously generated reasoning tokens. This is the standard autoregressive state: the model conditions on the prompt and its own partial output to predict the next token.

• Image phase state: $s_k^{\text{img}} = (c, y, x_{t_k}, t_k)$, where $c$ is the original prompt, $y$ is the completed reasoning trace (all text tokens have been generated), $x_{t_k}$ is the current noisy image latent at flow time $t_k$, and $t_k$ represents the current position along the denoising trajectory. The critical design choice here is that the reasoning text $y$ is "frozen" once the image phase begins — it becomes part of the conditioning context for all subsequent denoising steps, but is not further modified during image generation.

The transition between phases is implicit: once the text policy emits a special end-of-thinking token, the state representation switches from $s_k^{\text{txt}}$ to $s_k^{\text{img}}$, with the completed reasoning trace $y$ packaged into the image phase state.

Action Space $\mathcal{A}$. The action depends on which phase the MDP is in:

• Text phase action: $a_k^{\text{txt}} \in \mathcal{V}$ is a single discrete token drawn from the vocabulary $\mathcal{V}$. This is sampled from the model's autoregressive next-token distribution $\pi_\theta(a_k^{\text{txt}} \mid s_k^{\text{txt}})$.

• Image phase action: $a_k^{\text{img}} = x_{t_k - \Delta t} \in \mathbb{R}^d$ is the denoised latent at the next flow time step, a continuous vector in the latent space of the VAE used by the flow-matching model. This is sampled from the model's flow-matching denoising distribution $\pi_\theta(a_k^{\text{img}} \mid s_k^{\text{img}})$.

The paper is explicit that actions in both phases are stochastic during RL training — text tokens are sampled with non-zero temperature, and image latents are sampled via SDE (rather than the deterministic ODE used at inference). This stochasticity is essential for exploration: without it, the policy gradient would have no variance to learn from.

Transition $P$. Both phases have deterministic transitions given the action, which simplifies the MDP structure:

• Text transition: Appending $a_k^{\text{txt}}$ to the token sequence: $y_{<k+1} = (y_{<k}, a_k^{\text{txt}})$. The new state is $s_{k+1}^{\text{txt}} = (c, y_{<k+1})$.

• Image transition: Advancing the latent from $x_{t_k}$ to $x_{t_k - \Delta t} = a_k^{\text{img}}$. The new state is $s_{k+1}^{\text{img}} = (c, y, x_{t_k - \Delta t}, t_k - \Delta t)$.

The deterministic transition means that all stochasticity comes from the policy itself, not from environment dynamics — the "environment" for this MDP is effectively just the mechanism that accumulates tokens and advances flow time.

Reward $R$. The reward structure is sparse and terminal: $R(x_0, c)$ is assigned only after the image latent has been fully denoised to the clean latent $x_0$, which is then decoded to a pixel image and evaluated by the reward model. All intermediate text tokens and all intermediate denoising steps receive zero reward. This sparsity is a deliberate choice — it forces the model to learn long-range credit assignment across both the text and image phases, and it means the advantage signal $\hat{A}_i$ for a trajectory must propagate backward through the entire text chain and the entire denoising trajectory to assign credit to individual tokens and denoising steps.

Why sparse terminal rewards work here. The paper relies on two mechanisms to make sparse rewards effective despite the long credit assignment horizon: (1) group-relative advantages normalize the reward signal within each batch, making the relative quality of different reasoning chains and denoising trajectories immediately comparable; (2) the policy gradient formulation automatically handles temporal credit assignment — the importance ratio $r_{i,k}$ at each step $k$ multiplies the same trajectory-level advantage $\hat{A}_i$, so steps that contributed to high-reward trajectories get reinforced proportionally to how much more likely they became under the current policy relative to the old policy.

Why a unified MDP rather than two separate MDPs. The paper could have formulated text generation and image generation as two separate sequential problems with their own rewards, but this would break the critical coupling: the reward signal for image quality must inform the text policy about which reasoning patterns lead to better images. In a separate-MDP formulation, the text policy would need its own reward — but what reward could you assign to intermediate reasoning tokens without evaluating the final image? The unified MDP with sparse terminal reward elegantly solves this: the reasoning tokens are optimized solely through their downstream effect on image quality, as captured by the same terminal reward that evaluates the image. This is the mathematical expression of the paper's core insight: reasoning should be optimized for its utility to visual synthesis, not for its standalone quality.

---

3.4.2 Text GRPO: The Discrete Token Optimization Component

For the autoregressive text component of the unified model, UniGRPO adopts the standard Group Relative Policy Optimization (GRPO) formulation from Shao et al. (2024) without modification. Understanding this component is necessary before seeing how it couples with the image objective.

Group sampling. Given a prompt $c$, the current policy $\pi_\theta$ generates a group of $G$ reasoning chains $\{y_i\}_{i=1}^G$, where each chain $y_i = (y_{i,1}, y_{i,2}, \ldots, y_{i,|y_i|})$ is a sequence of tokens. The paper uses a group size of $G = 24$ (Table 3).

Group-relative advantage. The advantage for the $i$-th trajectory is computed purely from within-group statistics:

$\hat{A}_i = \frac{R_i - \text{mean}(\{R_j\}_{j=1}^G)}{\text{std}(\{R_j\}_{j=1}^G)}$

where $R_i$ is the terminal reward (image quality score) for the $i$-th trajectory, $\text{mean}(\{R_j\})$ is the average reward across the group of $G$ trajectories, and $\text{std}(\{R_j\})$ is the standard deviation of rewards across the group.

What it computes: this normalizes the raw rewards into standardized scores where positive values indicate above-average performance within the group, negative values indicate below-average performance, and the magnitude represents how many standard deviations away from the group mean a trajectory falls. This eliminates the need for a learned value function (as in PPO) and makes the optimization invariant to reward scale shifts.

Why this form: group-relative advantages provide an automatic curriculum — even if the absolute quality of all trajectories improves over training (shifting the mean reward upward), the within-group ranking still provides a meaningful learning signal because some trajectories will be better than others. This is also computationally efficient because it requires no separate value network and no reward normalization across batches. The division by standard deviation automatically adapts the advantage scale to the current reward variance, preventing gradients from becoming too large when reward variance is high or too small when variance shrinks.

GRPO objective for text. The text component's loss function is:

$\mathcal{J}_{\text{Text}}(\theta) = \frac{1}{G} \sum_{i=1}^{G} \frac{1}{|y_i|} \sum_{k=1}^{|y_i|} \left( \min \left( r_{i,k} \hat{A}_i, \; \text{clip}(r_{i,k}, 1-\epsilon, 1+\epsilon) \hat{A}_i \right) - \beta_{\text{txt}} D_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}}) \right)$

where:

• $|y_i|$ is the length in tokens of the $i$-th reasoning chain (the averaging normalizes for variable-length sequences).
• $r_{i,k} = \frac{\pi_\theta(y_{i,k} \mid y_{i,<k})}{\pi_{\theta_{\text{old}}}(y_{i,k} \mid y_{i,<k})}$ is the per-token importance ratio — the probability of token $y_{i,k}$ under the current policy divided by its probability under the policy that generated the trajectory (the "old" policy, frozen at the start of the batch).
• $\epsilon$ is the clipping parameter (set to 0.2 in standard GRPO, though the paper does not specify the exact value for UniGRPO).
• $\beta_{\text{txt}}$ is the KL penalty coefficient for the text component (set to 0 in Table 3 — meaning no KL penalty is applied to the text policy).
• $D_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}})$ is the KL divergence between the current and reference text policies.

What it computes: for each token in each reasoning chain, the objective multiplies the token's importance ratio by the trajectory's group-relative advantage. The min-and-clip operation is PPO's standard conservative policy iteration: if the current policy makes a token more likely than the old policy ($r_{i,k} > 1$) and the trajectory had a positive advantage, the update is capped at $1 + \epsilon$ to prevent the policy from moving too far in one update; if the trajectory had a negative advantage and the current policy makes the token less likely ($r_{i,k} < 1$), the update is capped at $1 - \epsilon$ symmetrically. The result is summed over tokens and averaged over trajectories in the group.

Why this form: the per-token importance ratio allows fine-grained credit assignment — tokens within a high-reward trajectory are reinforced proportionally to how "surprising" they were (how much the policy had to change to make them likely). A token that was already highly probable under the old policy gets a ratio near 1 and contributes little to the gradient; a token that the policy learned to make much more probable gets a larger update. The KL penalty (even when $\beta_{\text{txt}} = 0$, as in this paper) is conceptually replaced by the clipping mechanism, which acts as a trust-region constraint per token, preventing any single token's probability from changing by more than a factor of $1 \pm \epsilon$.

The zero KL penalty for text. Table 3 shows $\beta_{\text{txt}} = 0$ for the text component. This is notable because standard GRPO often uses a small positive KL coefficient to prevent the policy from diverging too far from the reference. The paper does not explicitly justify this choice, but it is consistent with the "minimalist methodology" philosophy — the clipping alone appears sufficient for stability in this setting, possibly because the reasoning chains are relatively short and the policy starts from a well-behaved SFT checkpoint.

---

3.4.3 FlowGRPO: The Continuous Image Denoising Optimization Component

For the flow-matching image component, UniGRPO adopts FlowGRPO (Liu et al., 2025a) with its Fast variant and RatioNorm mechanism. This component is substantially more complex than text GRPO because flow-matching models are deterministic at their core — they define a probability path via an ODE — and must be stochastified for RL exploration.

Why flow matching needs to be stochastified. Flow matching models (Lipman et al., 2022; Liu et al., 2022) learn a velocity field $v_\theta(x_t, t)$ that transports samples from a simple prior distribution (Gaussian noise at $t=1$) to the data distribution (clean latents at $t=0$) via the ODE:

$\frac{dx_t}{dt} = v_\theta(x_t, t)$

This ODE is deterministic: given the same initial noise $x_1$ and the same conditioning, the model always produces the same output. For RL, this is problematic because there is no exploration — the policy gradient would have zero variance, and the model could never discover better outputs through trial and error. FlowGRPO solves this by reformulating the generation process as a Stochastic Differential Equation (SDE) that adds controlled noise during denoising, creating a distribution over trajectories that the policy gradient can optimize.

SDE sampling formulation. The denoising step from time $t_k$ to $t_{k-1} = t_k - \Delta t$ is formulated as a discrete Euler-Maruyama step of an SDE:

$\Delta x_{t_k} = \left[ v_\theta(x_{t_k}, t_k) + \frac{\sigma_{t_k}^2}{2 t_k} \bigl( x_{t_k} + (1 - t_k) v_\theta(x_{t_k}, t_k) \bigr) \right] \Delta t + \sigma_{t_k} \sqrt{\Delta t} \; \epsilon$

where $\Delta x_{t_k} = x_{t_k - \Delta t} - x_{t_k}$ is the change in the latent over the step, $v_\theta(x_{t_k}, t_k)$ is the model's predicted velocity at the current state and time, $\sigma_{t_k}$ is the noise level at time $t_k$ (controlling the amount of stochastic exploration), and $\epsilon \sim \mathcal{N}(0, I)$ is standard Gaussian noise.

What it computes: this equation produces the next latent $x_{t_k - \Delta t}$ by combining three terms. The first term $v_\theta \Delta t$ is the deterministic flow matching drift — this is what the ODE would do, moving the latent toward the data distribution along the learned probability path. The second term $\frac{\sigma_{t_k}^2}{2 t_k} (x_{t_k} + (1-t_k) v_\theta) \Delta t$ is a correction term that ensures the SDE maintains the same marginal distribution as the ODE at each time $t$ (it is derived from the Fokker-Planck equation to keep the process consistent with the flow matching probability path). The third term $\sigma_{t_k} \sqrt{\Delta t} \epsilon$ is the injected Gaussian noise that creates stochasticity — this is what enables exploration. The output is a noisy next latent that is close to (but not exactly) what the deterministic ODE would produce.

Why this form: the specific form of the correction term is chosen so that the SDE preserves the marginal distributions of the flow matching process — meaning that if the model's velocity field is perfect, the SDE still generates samples from the correct data distribution despite the added noise. This is crucial because it means the SDE doesn't introduce bias; it only adds variance around the deterministic path. The noise term $\sigma_{t_k} \sqrt{\Delta t}$ scales with $\sqrt{\Delta t}$ (consistent with Brownian motion) so that as $\Delta t \to 0$, the process converges to a well-defined continuous-time SDE.

Noise level and SDE window. The paper specifies $\sigma_{t_k} = 0.8$ (Table 3) and applies the SDE only within a "window" of $[0, 5]$ — meaning only the first 3 denoising steps out of 25 total training timesteps use SDE sampling (SDE Window Size = 3, Table 3). The remaining steps use standard deterministic ODE sampling. This is the "FlowGRPO-Fast" variant: by confining stochastic exploration to early timesteps (where the latent is most noisy and small perturbations have the largest downstream effect on the final image) and using deterministic ODE for later timesteps, the method reduces computational overhead while preserving optimization effectiveness. Early timesteps determine the high-level structure of the image; later timesteps refine details. Perturbing early steps explores different image compositions and layouts; perturbing later steps would only add fine-grained noise with limited impact on the reward.

Importance ratio for flow steps. The analog of the text GRPO importance ratio for a flow-matching denoising step is:

$r_{t_k}(\theta) = \frac{p_\theta(x_{t_k - \Delta t} \mid c, y, x_{t_k}, t_k)}{p_{\theta_{\text{old}}}(x_{t_k - \Delta t} \mid c, y, x_{t_k}, t_k)}$

where $p_\theta(x_{t_k - \Delta t} \mid c, y, x_{t_k}, t_k)$ is the probability density of generating the next latent $x_{t_k - \Delta t}$ under the current policy, and $p_{\theta_{\text{old}}}$ is the probability under the old (frozen) policy. Under the SDE formulation, this probability is a Gaussian distribution with mean $\mu_\theta(x_{t_k}, t_k)$ (the drift term from the SDE equation above) and variance $\sigma_{t_k}^2 \Delta t$ (from the injected noise).

The RatioNorm problem and solution. A critical complication arises because the distribution of importance ratios $r_{t_k}(\theta)$ in flow-matching RL is systematically left-shifted (mean $< 1$) and exhibits inconsistent variance across timesteps. The paper cites GRPO-Guard (Wang et al., 2025) for this observation. This left-shift means that most steps have ratios below 1, which — under standard PPO clipping — would cause the conservative clipping to only constrain negative updates (when $\hat{A}_i < 0$ and $r_{t_k} < 1-\epsilon$) while letting positive updates pass through unclipped (since even "overconfident" steps still have $r_{t_k} < 1+\epsilon$), thereby failing to prevent reward hacking.

RatioNorm (Ratio Normalization). To fix this, the paper adopts the RatioNorm mechanism from GRPO-Guard, which standardizes the log-importance ratio to center its distribution around zero:

$\log \tilde{r}_{t_k}(\theta) = \sigma_{t_k} \sqrt{\Delta t} \left( \log r_{t_k}(\theta) + \frac{\|\Delta \mu_\theta(x_{t_k}, t_k)\|^2}{2 \sigma_{t_k}^2 \Delta t} \right)$

where $\tilde{r}_{t_k}(\theta)$ is the normalized (and exponentiated) importance ratio used in the clipping objective, $\log r_{t_k}(\theta)$ is the raw log-ratio of Gaussian densities, $\Delta \mu_\theta(x_{t_k}, t_k) \triangleq \mu_{\theta_{\text{old}}}(x_{t_k}, t_k) - \mu_\theta(x_{t_k}, t_k)$ is the mean drift between the current and reference policies (the difference in their predicted denoising directions), and $\sigma_{t_k} \sqrt{\Delta t}$ is the standard deviation of the SDE noise.

What it computes: the normalization term $\frac{\|\Delta \mu_\theta\|^2}{2 \sigma_{t_k}^2 \Delta t}$ inside the parentheses is exactly the KL divergence between the current and old Gaussian policies at this step. Adding this to the log-ratio then multiplying by the standard deviation $\sigma_{t_k} \sqrt{\Delta t}$ produces a standardized log-ratio whose distribution is approximately centered at zero. After exponentiating to get $\tilde{r}_{t_k}$, the normalized ratio has mean near 1, restoring the symmetry of the clipping mechanism — both positive and negative updates are now constrained by the $[1-\epsilon, 1+\epsilon]$ bounds.

Why this form: the specific correction term is derived from the observation that the log-importance ratio and the KL divergence have opposite signs on average — when the policy moves away from the reference, the KL increases (positive contribution) while the log-ratio of a trajectory generated under $\theta_{\text{old}}$ becomes negative (since the trajectory is less likely under $\theta$). Adding them cancels the systematic bias, centering the normalized distribution. The multiplication by $\sigma_{t_k} \sqrt{\Delta t}$ ensures the normalization is consistent across timesteps with different noise levels.

FlowGRPO objective. The final FlowGRPO loss is computed only over the SDE timestep subset $\mathcal{T}_{\text{SDE}}$:

$\mathcal{J}_{\text{Flow}}(\theta) = \frac{1}{G} \sum_{i=1}^{G} \frac{1}{|\mathcal{T}_{\text{SDE}}|} \sum_{t_k \in \mathcal{T}_{\text{SDE}}} \left( \min \left( \tilde{r}_{i,t_k} \hat{A}_i, \; \text{clip}(\tilde{r}_{i,t_k}, 1-\epsilon, 1+\epsilon) \hat{A}_i \right) - \beta_{\text{img}} D_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}}) \right)$

where $|\mathcal{T}_{\text{SDE}}|$ is the number of denoising steps within the continuous SDE window (3 steps in this paper), $\tilde{r}_{i,t_k}$ is the RatioNorm-corrected importance ratio for the $i$-th trajectory at time $t_k$, and the rest of the notation mirrors the text GRPO objective. The key difference from text GRPO: the importance ratio is at the level of denoising steps (not tokens), the RatioNorm correction is applied to fix the distributional bias, and the objective is only computed over stochastic steps.

Why only SDE steps get gradients: the ODE steps are deterministic — given $x_{t_k}$, the next latent $x_{t_k - \Delta t}$ is uniquely determined, so there is no distribution to differentiate through. The policy's influence on ODE steps is captured implicitly through the SDE steps that preceded them, since the initial SDE states determine which deterministic trajectory the ODE follows. Computing gradients only on SDE steps significantly reduces memory and computation while still capturing the policy's effect on the final image.

---

3.4.4 UniGRPO's Two Modifications: CFG Elimination and Velocity MSE Regularization

These are the paper's novel contributions to the FlowGRPO training recipe, both motivated by the need for scalability to multi-round and multi-condition interleaved generation scenarios.

Eliminating Classifier-Free Guidance (CFG)

What CFG does and why it's standard at inference. Classifier-Free Guidance (Ho and Salimans, 2022) is the dominant technique for improving prompt adherence in diffusion and flow-matching models. At each denoising step, instead of using the model's conditional velocity prediction $v_\theta(x_t, t, c)$ directly, CFG computes a guided velocity as:

$v^{\text{CFG}} = v_\theta(x_t, t, \varnothing) + w \cdot (v_\theta(x_t, t, c) - v_\theta(x_t, t, \varnothing))$

where $v_\theta(x_t, t, \varnothing)$ is the unconditional prediction (with null conditioning), $v_\theta(x_t, t, c)$ is the conditional prediction, and $w$ is the guidance scale (typically $w > 1$, often 7.5 for diffusion models). This linearly extrapolates away from the unconditional prediction toward the conditional prediction, sharpening the conditioning signal and improving prompt adherence at the cost of reduced diversity. Crucially, each CFG step requires two forward passes through the model: one conditional, one unconditional.

The scaling problem in multi-condition and multi-round settings. The paper identifies that CFG's computational cost scales multiplicatively with the number of conditioning signals. For standard text-to-image generation, CFG requires 2 evaluations per step. For image editing (conditioning on both text and a reference image), CFG would require at least 3 evaluations per step (unconditional, text-only, text+image). For multi-round interleaved generation — where the model alternates between generating text and refining images, each step conditioning on the accumulated dialogue history, previous images, and the current intent — the number of required conditioning contexts grows combinatorially. Moreover, in an RL setting, CFG creates a branched computation graph: the same model must be called with multiple different conditioning inputs, each producing different velocity predictions that are then combined linearly. This branching makes gradient estimation more complex and memory-intensive, since each branch must be tracked through the computation graph.

UniGRPO's solution: train without CFG, internalize alignment through RL. The paper trains UniGRPO entirely without CFG, setting the CFG scale to 1 (Table 3: "CFG Scale = 1"). This enforces a linear, unbranched rollout: at each denoising step, the model is called exactly once with the actual conditioning, and the resulting velocity is used directly. This is computationally simpler, memory-efficient, and produces a clean gradient path through the model — exactly what is needed for stable multi-round RL.

How alignment is recovered without CFG. The paper's key insight is that RL training itself can substitute for CFG's alignment function. CFG works by sharpening the conditioning signal at inference time, making the model attend more strongly to the prompt. RL training works by updating the model's weights to maximize a reward that explicitly measures prompt adherence and visual quality — effectively baking the alignment directly into the policy parameters rather than achieving it through a post-hoc guidance trick. The reward model (InternVL-based, fine-tuned on user preference data) provides the optimization pressure toward better prompt adherence; the policy update internalizes this pressure into the velocity predictions themselves, so that even without CFG, the model's conditional predictions are sharp enough to produce high-quality outputs.

Empirical evidence. The ablation in Figure 4 shows that training without CFG yields comparable or slightly better final performance when evaluated with CFG (which is the standard inference-time practice). The paper states: "although CFG during training yields images with higher rewards, removing CFG results in comparable or better final performance when evaluated with CFG." This confirms that CFG is useful during training only for exploration purposes (it helps generate higher-quality samples that set a higher reward baseline), but the policy itself learns to achieve equivalent quality without it.

Velocity-Based Regularization via MSE Penalty

The problem with standard latent KL divergence. In the SDE formulation, the step-wise transition probabilities are Gaussian with known mean and variance. This allows the exact local KL divergence at each step to be computed analytically:

$D_{\text{KL}}(p_\theta \| p_{\text{ref}}) = \frac{\| \mu_\theta(x_{t_k}, t_k) - \mu_{\text{ref}}(x_{t_k}, t_k) \|^2}{2 \sigma_{t_k}^2 \Delta t}$

where $\mu_\theta$ and $\mu_{\text{ref}}$ are the mean drifts (predicted next states) of the current and reference policies respectively, and $\sigma_{t_k}^2 \Delta t$ is the variance of the Gaussian transition. This is the standard KL penalty used in the FlowGRPO objective (Equation 6 in the paper, where it appears as $D_{\text{KL}}(\pi_\theta \| \pi_{\text{ref}})$).

Why this form causes temporal vulnerabilities. The denominator $2 \sigma_{t_k}^2 \Delta t$ means the KL penalty is weighted by the inverse noise variance. At timesteps with high noise variance (early in denoising, when $t_k$ is large), $\sigma_{t_k}$ is large, making $1/(\sigma_{t_k}^2 \Delta t)$ very small — the KL penalty becomes negligible. At timesteps with low noise variance (late in denoising), the penalty becomes very large. This uneven weighting creates temporal vulnerabilities: the RL optimizer can make large changes to the velocity field at high-noise timesteps (where the KL penalty is weak) that cascade through the deterministic ODE phase to produce dramatically different final images, potentially exploiting reward model blind spots without being penalized. The optimizer effectively learns which timesteps have weak regularization and concentrates its "reward hacking" effort there.

UniGRPO's solution: direct MSE on velocity fields. The paper replaces the latent KL divergence with an unweighted mean squared error penalty applied directly to the model's velocity predictions:

$\mathcal{L}_{\text{MSE}}(\theta) = \| v_\theta(x_{t_k}, t_k, y) - v_{\text{ref}}(x_{t_k}, t_k, y) \|^2$

where $v_\theta(x_{t_k}, t_k, y)$ is the current policy's predicted velocity at state $x_{t_k}$ and time $t_k$ (conditioned on the reasoning trace $y$), and $v_{\text{ref}}(x_{t_k}, t_k, y)$ is the reference (pre-RL) model's predicted velocity at the same state. The MSE penalty weight is set to $1.5 \times 10^{-5}$ (Table 3: "MSE Loss Weight").

What it computes: for each denoising step in the SDE window, this penalty computes the squared Euclidean distance between the current and reference models' velocity vectors. The velocity $v_\theta$ is the fundamental output of the flow-matching model — it specifies the direction and magnitude of denoising at each point in the latent space. By penalizing deviations from the reference velocity uniformly across all timesteps (no $\sigma_{t_k}$ weighting), the regularization is temporally homogeneous — every step gets the same pressure to stay close to the reference model.

Why this form is more robust. Three reasons:

1. Uniform penalty across noise levels: Unlike the latent KL, which is inversely weighted by noise variance, the velocity MSE applies the same penalty magnitude at all timesteps. This eliminates the temporal vulnerabilities where high-noise steps could be exploited. Empirically, Figure 5 shows that latent KL produces "grid-like artifacts as early as step 250" despite using a sufficiently large KL coefficient, while velocity MSE maintains stable training and high-fidelity images.

2. Direct constraint on the model's core function: The velocity field is what the flow-matching model actually learns and outputs. Regularizing it directly means the RL update cannot change the model's fundamental behavior in ways that produce plausible-looking latents (passing the KL check in latent space) but represent degenerate solutions in velocity space. For example, a model could learn to produce latents that have the right statistical properties (satisfying the latent KL) but are generated through entangled, non-smooth velocity fields that fail to generalize.

3. Preservation of generative priors: The reference model's velocity field encodes the pretrained understanding of how to denoise images — it knows that certain latent-space directions correspond to realistic image features while others correspond to artifacts. The MSE penalty explicitly forces the RL-tuned model to stay close to these learned directions, rather than discovering reward-maximizing but unrealistic paths through the latent space. The paper states this "leaves fewer loopholes for policy exploitation."

The MSE weight magnitude. The penalty weight of $1.5 \times 10^{-5}$ is very small, suggesting that only a gentle nudge toward the reference is needed — the clipping mechanism (with a loss clip range of $1 \times 10^{-6}$, Table 3) provides the primary regularization against large policy updates, while the velocity MSE provides an additional safety net against more subtle forms of reward hacking that might slip through the clipping bounds.

---

3.4.5 Putting It Together: The Unified Objective and Training Loop

The combined objective. UniGRPO's final training objective is a simple weighted sum:

$\mathcal{J} = \mathcal{J}_{\text{Text}} + \lambda \mathcal{J}_{\text{Flow}}$

where $\lambda = 1$ (Table 3), meaning the text and image objectives are equally weighted. The paper explicitly states this is a deliberately simple choice to "avoid over-design" — no careful tuning, no dynamic scheduling, no per-module adaptive weighting. The equal weighting works because both objectives are already normalized: the text objective averages over tokens and trajectories, and the Flow objective averages over SDE steps and trajectories.

What it computes: the text objective $\mathcal{J}_{\text{Text}}$ encourages the language model component to produce reasoning tokens that lead to high-reward images (via the shared advantage $\hat{A}_i$), while keeping token-level probability changes within the clipping bounds. The Flow objective $\mathcal{J}_{\text{Flow}}$ encourages the flow-matching component to produce denoising trajectories that lead to high-reward images, while keeping denoising-step-level probability changes within the RatioNorm-corrected clipping bounds and staying close to the reference velocity field via the MSE penalty. Both objectives use the identical group-relative advantages $\hat{A}_i$ computed from the same terminal rewards — this is the critical coupling mechanism.

Training loop mechanics (one iteration):

1. Sample a batch of prompts (batch size 32, Table 3). For each prompt:

   • Generate reasoning chains: The current policy $\pi_\theta$ autoregressively generates $G = 24$ reasoning chains (one per prompt? or 24 per prompt? The paper says "for a given prompt $c$, we first sample $G$ reasoning chains" — suggesting 24 chains per prompt, but with batch size 32, this would mean $32 \times 24 = 768$ trajectories per iteration. The exact batching is not fully specified.)
   • Generate images: For each reasoning chain, the model's flow-matching component denoises an image latent from noise to clean data, using hybrid SDE-ODE sampling: the first 3 steps (SDE Window Size = 3, within the window $[0, 5]$) use the SDE formulation with noise level $\sigma = 0.8$, and the remaining 22 steps (25 total training timesteps minus 3 SDE steps) use deterministic ODE. Total training timesteps = 25; at evaluation, 50 timesteps are used (Table 3).
   • Compute rewards: The terminal reward model (InternVL-based, from RewardDance) scores each generated image against its prompt, producing $G$ scalar rewards.

2. Compute group-relative advantages: For each group of $G$ trajectories, standardize the rewards into advantages $\hat{A}_i$ using mean and standard deviation.

3. Compute text GRPO loss: For each reasoning chain, compute per-token importance ratios between current and old policies, apply clipping, multiply by the trajectory's advantage, sum over tokens, average over the group. No KL penalty ($\beta_{\text{txt}} = 0$).

4. Compute FlowGRPO loss: For each denoising trajectory, compute per-step (SDE steps only) importance ratios using the RatioNorm correction, apply clipping, multiply by the trajectory's advantage, sum over SDE steps, average over the group. Add velocity MSE penalty weighted by $1.5 \times 10^{-5}$.

5. Backpropagate the combined loss $\mathcal{J} = \mathcal{J}_{\text{Text}} + \mathcal{J}_{\text{Flow}}$ through the entire unified model, updating both the text and image parameters jointly.

6. PPO epochs: The paper specifies 2 PPO epochs (Table 3), meaning the same batch of trajectories is used for two gradient update steps before resampling. This is the standard PPO practice of making multiple updates per batch of collected data while the importance ratios remain valid.

Key hyperparameters from Table 3:

Design choice: different learning rates for text and image components. The reasoning (text) component uses a learning rate of $1 \times 10^{-6}$, while the denoising (image) component uses $3 \times 10^{-5}$ — a 30× difference. The paper does not explicitly justify this, but it likely reflects the different sensitivity of the two components: the text policy operates on discrete tokens with a well-behaved softmax output, so small changes in logits produce meaningful changes in token probabilities; the flow-matching policy operates on high-dimensional continuous velocity vectors where gradients can be much smaller per dimension, necessitating a higher learning rate to achieve comparable policy updates.

Design choice: why the loss clip range is extremely small for Flow ($1 \times 10^{-6}$). This is the clipping parameter $\epsilon$ in the FlowGRPO objective. A clip range of $10^{-6}$ means the policy is constrained to make extremely conservative updates — the effective importance ratio is clamped to $[1 - 10^{-6}, 1 + 10^{-6}]$, which is nearly 1.0. This extreme conservatism reflects the fragility of flow-matching RL: the continuous, high-dimensional action space makes it easy for the policy to drift into regions of velocity space that produce degenerate images, and even small per-step changes can compound over the denoising trajectory to produce large effects on the final image. The tiny clip range, combined with the velocity MSE penalty, creates a very tight trust region around the reference model.

Design choice: 1024 resolution training. The paper trains at 1024 × 1024 pixel resolution (Table 3), which is notably high for RL-based image generation training. Most prior RL for image generation work trains at 512 or lower due to computational constraints. Training at 1024 means the model learns to optimize for high-resolution quality directly, rather than optimizing at low resolution and hoping the improvements transfer when evaluated at high resolution. This choice is made possible by the efficiency of the FlowGRPO-Fast variant (only 3 SDE steps with gradients) and the CFG-free training (halving the forward passes per step), which together keep the computational cost manageable.

What happens at inference time. At inference (evaluation), the model:

• Generates reasoning tokens autoregressively.
• Denoises images using 50 deterministic ODE steps (no SDE noise, since exploration is only needed for training).
• Optionally applies CFG (since the evaluation in Figure 4 and Table 1 uses CFG at inference, even though the model was trained without it). The paper shows that CFG still provides a benefit at inference even for CFG-free-trained models, likely because the reward-maximizing policy has learned to produce sharper conditional predictions, and CFG further amplifies this sharpness.

This training-vs-inference asymmetry — SDE for training, ODE for inference; no CFG for training, CFG for inference — is a deliberate design pattern: training uses mechanisms that enable exploration and scalability, while inference uses mechanisms that maximize output quality. The RL training ensures the model's weights are well-aligned even without inference-time tricks, and the tricks provide an additional quality boost on top of the aligned base.

4. Key Insights and Innovations

Innovation 1: The Unified MDP as a Conceptually Necessary Abstraction for Interleaved Generation

Before UniGRPO, the standard approach to applying RL to models that generate both text and images was to treat the two modalities as separate optimization problems — either by training disjoint models (a separate LLM for reasoning, a separate diffusion model for images, as in DualGRPO and PromptRL), by alternating training between modules (SepGRPO), or by optimizing them sequentially in a two-stage pipeline (ReFL + TextGRPO). The field's de facto assumption was that text and image generation are sufficiently different — discrete tokens versus continuous latents, autoregressive versus flow-matching objectives — that they require separate RL formulations, separate reward signals, or separate training phases.

UniGRPO challenges this assumption at the conceptual level by demonstrating that the only thing needed to unify them is the MDP abstraction itself. By formalizing the entire "Prompt → Thinking → Image" sequence as a single MDP — where the action space alternates between discrete token selection and continuous latent denoising, where the state transitions between text and image phases carry the reasoning trace forward as conditioning context, and where a single sparse terminal reward on the final image propagates backward through both phases — the paper shows that text and image policies can be optimized jointly under the identical group-relative advantage signal. The mechanism (combining standard GRPO with FlowGRPO) is deliberately minimalist; what's novel is the recognition that no additional machinery is needed beyond the MDP framing.

This is a conceptual advance, not an algorithmic one. The paper does not propose a new policy gradient estimator, a new credit assignment scheme, or a new architecture. It simply observes that if interleaved generation is viewed as a sequential decision process where the "action type" changes mid-trajectory, then existing per-modality RL recipes slot naturally into the same optimization loop without modification. The significance lies in what this framing rules out as unnecessary: no separate reward model for reasoning quality, no value function for intermediate states, no alternating training phases, no tree-structured rollouts across separate models. The negative result on UniFPO (training collapse in Table 1) reinforces the point — this unification works specifically because FlowGRPO's SDE-based exploration and RatioNorm clipping provide stable gradients that couple cleanly with text GRPO, while FPO's alternative importance sampling approach does not. The MDP framing is what makes this comparison legible: it reveals that the choice of image RL algorithm is not an isolated design decision but a compatibility constraint on the entire unified system.

The intellectual move here is analogous to what MDP formulations did for robotics two decades ago: by abstracting diverse sensorimotor modalities (joint torques, visual observations, tactile feedback) into a common state-action-reward language, they enabled a single reinforcement learning loop to optimize policies spanning fundamentally different action spaces. UniGRPO does the same for multimodal generation, and the paper's forward-looking emphasis on multi-round scaling (editing, dialogue, storytelling) suggests this abstraction will become increasingly necessary as interleaved generation horizons grow.

Evidence: The joint optimization results in Table 1, where UniGRPO (0.8381 TA, 0.90 GenEval) substantially outperforms both single-modality optimization (FlowGRPO with thinking at 0.8208 TA, TextGRPO at 0.8078 TA) and sequential two-stage approaches (ReFL + TextGRPO at 0.7987 TA), confirming that the unified MDP coupling is the active ingredient, not just the sum of independent optimizations.

---

Innovation 2: CFG Elimination as a Strategic Bet on RL-Hardenable Alignment

Classifier-free guidance is arguably the most universally adopted technique in diffusion and flow-matching inference — it is the standard way to make generated images adhere to text prompts, and nearly every deployed text-to-image model uses it. Removing CFG during training might therefore seem like accepting a handicap, and indeed the paper's ablation (Figure 4) confirms that training with CFG produces higher in-training rewards. The conventional wisdom would be: if CFG helps, keep it.

UniGRPO's counterintuitive move is to deliberately remove CFG during training, not because it improves single-round results (the ablation shows it is roughly performance-neutral), but because CFG is architecturally incompatible with the computational and memory scaling properties demanded by multi-round interleaved generation. This is a strategic bet: accept a harder optimization problem now (training without the inference-time alignment crutch) in exchange for a linear, unbranched computation graph that scales gracefully to editing, dialogue, and multi-turn interaction. The paper frames this explicitly as a forward-looking design choice: "this computational burden scales with the number of conditions; for multi-condition generation such as image editing, CFG demands at least three evaluations per step" and "this complexity compounds in multi-round interleaved generation."

What makes this intellectually distinctive is the reconceptualization of CFG from an inference-time necessity to a training-time convenience. The paper's implicit argument is that CFG's alignment function — sharpening the conditioning signal by extrapolating away from the unconditional distribution — is something the model can learn to internalize during RL training. The reward model provides explicit optimization pressure toward prompt adherence and image quality; the policy update bakes this pressure into the model weights. By the end of training, the model's conditional velocity predictions are sufficiently sharp that CFG at inference provides only a marginal additional boost (and the ablation shows this empirically). This reframes alignment as a weight-level property that RL can optimize, rather than an inference-time trick that must be applied during generation.

The significance extends beyond the immediate empirical result. If this finding generalizes — if RL-trained flow-matching models can reliably achieve strong prompt adherence without CFG — it would remove one of the primary computational bottlenecks for deploying these models in interactive settings. A CFG-free model at inference requires half the forward passes (or one-third for multi-condition editing), directly translating to 2× or 3× throughput improvements at scale. The paper doesn't fully explore this inference-time efficiency gain (all evaluations still use CFG), but the training-time evidence that it's possible is a provocative finding for the deployment-focused community.

Evidence: Figure 4 shows that training without CFG achieves comparable or better final validation reward compared to training with CFG, even though in-training rewards are higher with CFG. The paper notes this means CFG is "unnecessary for RL-based alignment."

---

Innovation 3: Diagnosing and Solving the Temporal Vulnerability in Latent-Space KL Regularization

The standard regularization approach in flow-matching RL — penalizing the KL divergence between the current and reference policy in latent space — has an internal structure that the field had not previously problematized. Because the latent KL at each denoising step is inversely weighted by the noise variance ($1/\sigma_{t_k}^2$), it applies dramatically different regularization pressure at different timesteps: early high-noise steps receive negligible penalty, while late low-noise steps receive massive penalty. The paper identifies this as a "temporal vulnerability" — the RL optimizer can concentrate its reward-hacking effort at timesteps where the KL is weakest, making large changes to the velocity field early in denoising that cascade through the deterministic ODE phase to produce reward-maximizing but degenerate images.

This is an instance of a broader phenomenon in RL for generative models: regularization that is formally correct (the exact KL divergence between Gaussian transition distributions) can be practically inadequate because it is non-uniform across the trajectory, creating exploit-prone regions that the optimizer will find. The paper's diagnosis goes beyond saying "reward hacking happens" to pinpointing why the standard remedy fails and where the vulnerability is located — the temporal dimension of the denoising trajectory.

The solution — replacing latent KL with unweighted MSE on the velocity fields — is conceptually clean but its motivation is what's novel. Rather than trying to design a better weighting scheme for the KL (which would preserve the probabilistic interpretation), the paper abandons the probabilistic KL entirely in favor of a direct behavioral constraint: keep the model's velocity predictions close to the reference model's, uniformly across all timesteps. This shift from distributional closeness (KL between transition probabilities) to functional closeness (MSE between velocity outputs) reflects a pragmatic judgment: what matters for preventing reward hacking is not that the new policy's latent transition distribution remains close to the reference in an information-theoretic sense, but that the model's fundamental behavior — the direction and magnitude of denoising — does not drift into degenerate regions.

This innovation is significant because it identifies a failure mode in the standard mathematical machinery of diffusion/flow RL that is not obvious from the equations alone. The latent KL is derived correctly from the SDE formulation; the problem only manifests empirically, in the form of grid-like artifacts that "emerge as early as step 250" (Figure 5) despite using a sufficiently large KL coefficient. The paper's contribution is connecting this empirical degradation to the specific mathematical structure of the penalty — the inverse noise variance weighting — and providing a fix that is both simpler to implement and more robust in practice.

Evidence: Figure 5 provides the direct comparison. Training with no KL leads to reward hacking (validation reward drops, images become oversaturated and unnatural). Training with latent KL shows a "significant drop in training reward" (indicating the KL penalty is constraining the policy) yet "grid-like artifacts still emerge as early as step 250." Training with velocity MSE maintains stable dynamics and produces high-fidelity images. The MSE weight of $1.5 \times 10^{-5}$ is notably small, suggesting the uniform regularization is effective even at very low strength because it closes the temporal loopholes that the latent KL left open.

---

Innovation 4: The "Fundamental Unit" Validation Strategy as a Principled Research Methodology

This innovation is methodological rather than algorithmic: the paper's decision to validate its unified RL framework on single-round reasoning-driven image generation ("Prompt → Thinking → Image") rather than immediately tackling full multi-round interleaved generation. The paper explicitly frames this as validating on "the fundamental unit" of interleaved generation, arguing that this setting "already encompasses both text and image generation, covering the essential components" and "serves as a meaningful and principled testbed" in the absence of open-source models natively capable of longer-horizon generation.

This might seem like an obvious practical constraint — they didn't have a multi-round base model, so they couldn't test multi-round RL. But the paper elevates this constraint into a research design principle with two subtle but important implications:

First, it enables rigorous ablation of the joint optimization claim. By keeping the generation horizon fixed (single round), the paper can cleanly compare joint optimization (UniGRPO) against optimizing only text (TextGRPO), only images (FlowGRPO), or optimizing them sequentially (ReFL + TextGRPO). In a multi-round setting, the comparison space explodes combinatorially — different rounds might need different optimization strategies, credit assignment across rounds becomes entangled with within-round credit assignment, and the baselines become harder to define. The single-round testbed isolates the core question: does jointly optimizing reasoning and synthesis under a unified advantage signal outperform any scheme that optimizes them separately? The answer (yes, by substantial margins in Table 1) is cleanly attributable to the unified MDP framework, not to any confounding multi-round dynamics.

Second, it separates architectural scalability from algorithmic validation. The paper makes two design choices (CFG elimination and velocity MSE) that are explicitly motivated by multi-round scalability, but tests them in a single-round setting where those benefits are latent. This is an unusual research strategy: introduce complexity-motivated simplifications, validate that they don't hurt performance on the simpler task, and position the simpler system as "scalable by construction" to the more complex task. The risk is that the scalability benefits might not materialize (the multi-round setting might introduce new failure modes that the single-round validation didn't surface), but the paper accepts this risk in exchange for being able to validate the core RL framework thoroughly on a tractable problem.

This methodology contrasts sharply with concurrent work: DualGRPO tests on a tree-structured rollout that is "incompatible with true interleaved multimodal generation" (the paper's critique), PromptRL validates on "limited training datasets," and SepGRPO alternates training between modules rather than joint optimization. These works attempt to solve the harder multi-component problem directly, but at the cost of making it difficult to isolate why their approaches work (or don't). UniGRPO's staged approach — validate the unified RL principle on the simplest problem that contains both modalities, then scale — provides a cleaner scientific narrative, even if it leaves the multi-round results as future work.

Evidence: The experimental design in Section 5, where all baselines operate on the same single-round task (Prompt → Reasoning → Image) using the same Bagel backbone after SFT, and the primary comparisons in Table 1 isolate the effect of joint vs. separate vs. sequential optimization. The paper's explicit statement in the conclusion — "future work will apply UniGRPO to complex scenarios such as interactive image editing, visual storytelling, and multi-turn dialogue" — acknowledges that the scaling claim remains unvalidated, but the framework is designed with that trajectory in mind.

5. Experimental Analysis

Evaluation Methodology

• Dataset. The paper uses two evaluation benchmarks. The first is an internal Text Alignment (TA) benchmark consisting of 150 diverse prompts; for each prompt, 4 images are generated and a VLM evaluates the outputs against multiple specific exam points per prompt, with each exam point receiving a binary score (1 for correct, 0 for incorrect). The second is GenEval (Ghosh et al., 2023), a standard public benchmark assessing Text-to-Image models on complex compositional capabilities including object counting, spatial relations, and attribute binding across six sub-categories.

• Base model. All experiments use Bagel (Deng et al., 2025), a unified multimodal model with inherent interleaved generation potential, combining autoregressive text modeling with flow-matching image synthesis. The paper notes that "vanilla Bagel exhibits limited instruction-following capabilities and suboptimal image generation quality," so the authors first performed Supervised Fine-Tuning (SFT) on Bagel using a curated internal dataset before any RL experiments. Unless stated otherwise, all baselines and experiments start from this SFT checkpoint as the reference model. The paper uses this model family specifically because it is "representative of the architectural direction the community is converging on" — AR text + flow-matching images — and because it natively supports the "Thinking" mode (chain-of-thought reasoning before image generation) that is central to the framework's validation.

• Metrics. The primary metrics are TA Score (average binary score across all exam points and evaluated images on the internal 150-prompt benchmark) and GenEval overall score (macro-average across the six compositional sub-categories: Single Object, Two Objects, Counting, Colors, Position, Attribute Binding). For GenEval, detailed per-category breakdowns are provided in Table 2. The reward model used during training is a differentiable InternVL-based model fine-tuned on collected user preference data (from RewardDance, Wu et al., 2025), explicitly designed to measure consistency between generated images and user prompts. The paper notes this specific reward model choice is to "ensure a fair comparison with gradient-based baselines like ReFL (which necessitates differentiable rewards)," while emphasizing that UniGRPO itself is "compatible with a broader range of verifier-based rewards."

• Baselines. The paper compares against a substantial set of methods, all starting from the same SFT Bagel checkpoint:

  • ReFL (Xu et al., 2024): Directly fine-tunes diffusion models by backpropagating reward model gradients through a randomly-picked late denoising timestep. Tested with (ReFL ✓) and without (ReFL ×) explicit thinking tokens.
  • ReFL (w/ Thinking): Generates thinking prompts during training but optimizes only the image generation component using the ReFL objective. This isolates the effect of reasoning-driven conditioning without joint optimization.
  • ReFL (w/ Thinking) + TextGRPO: A two-stage pipeline: first train with ReFL (w/ Thinking) on images, then optimize the text module using TextGRPO. Tests the sequential optimization hypothesis.
  • FPO (Flow Matching Policy Gradients, McAllister et al., 2025; also related to AWR, Peng et al., 2019): An alternative to FlowGRPO that uses the forward noising process to obtain $x_t$ and treats the ELBO of denoising as a surrogate for $\log p_\theta(x_0|c)$ to compute importance sampling weights. Tested with and without thinking tokens.
  • FlowGRPO (Liu et al., 2025a): The standard single-modal FlowGRPO recipe, applied only to the image component. Tested with and without thinking tokens.
  • TextGRPO: Standard GRPO applied only to the text reasoning component. Only tested with thinking tokens (since without thinking there are no text tokens to optimize).
  • UniFPO: A unified framework analogous to UniGRPO, but using FPO for the image component instead of FlowGRPO. This tests whether the specific image RL algorithm matters for joint optimization stability.
  • Bagel and SFT baselines: The original Bagel and the SFT-finetuned version, both tested with and without thinking tokens, serving as non-RL lower bounds.

• Generation budget / compute accounting. The paper does not report FLOPs or wall-clock time. Instead, compute is implicitly accounted for through the training configuration (Table 3), which specifies group size ($G = 24$), batch size (32), number of training timesteps (25), SDE window size (3), and PPO epochs (2). Comparisons between methods are made at the same training stage (all methods start from the same SFT checkpoint and are trained under comparable conditions). The paper notes that CFG elimination during training halves the forward passes per denoising step compared to CFG-based methods, but this efficiency gain is not quantified in the main results. Inference-time evaluation uses 50 timesteps (Table 3) and applies CFG for all methods during final evaluation, regardless of whether CFG was used during training.

• Cross-validation / statistical protocol. The paper does not report explicit cross-validation or statistical significance testing on the evaluation benchmarks. The TA benchmark (150 prompts with 4 images each = 600 evaluated images) and GenEval (standard test split) are fixed evaluation sets. The training and validation reward curves (Figure 3) show training dynamics across gradient update steps, providing evidence of convergence behavior. The ablation on CFG (Figure 4) is noted as not directly comparable to the main training curves because it uses the original Bagel (not the SFT checkpoint) at resolution 512 rather than 1024.

Main Quantitative Results

Headline Benchmark Performance (Table 1)

The central quantitative result is that UniGRPO achieves state-of-the-art performance on both evaluation benchmarks, reaching 0.8381 on TA and 0.90 overall on GenEval, substantially outperforming all baselines that optimize text and image components separately or sequentially.

The performance hierarchy across methods reveals a clear pattern:

• SFT provides the foundational boost. Bagel (no thinking, no RL): 0.6810 TA, 0.78 GenEval. After SFT with thinking tokens (SFT ✓): 0.7769 TA, 0.82 GenEval — a gain of approximately +0.0959 on TA and +0.04 on GenEval. This establishes SFT as a necessary but insufficient step; the paper notes SFT "compromises image sharpness, resulting in noticeable blurriness upon close inspection."

• Single-modal RL provides incremental gains over SFT. The best single-modality RL method with thinking is FlowGRPO ✓ at 0.8208 TA and 0.86 GenEval — an improvement over SFT ✓ of approximately +0.0439 on TA and +0.04 on GenEval. TextGRPO alone achieves 0.8078 TA and 0.88 GenEval, suggesting that optimizing reasoning text primarily helps compositional understanding (GenEval) while optimizing image generation primarily helps visual quality (TA).

• Joint optimization (UniGRPO) provides a synergistic leap. UniGRPO achieves 0.8381 TA and 0.90 GenEval. Relative to the best single-modality baseline (FlowGRPO ✓ at 0.8208 TA), this is a gain of +0.0173 on TA and +0.04 on GenEval. The GenEval gain is particularly notable: UniGRPO's 0.90 represents a substantial improvement over the next-best GenEval score (0.88 for both FlowGRPO × and TextGRPO ✓), suggesting that joint optimization of reasoning and synthesis is especially beneficial for compositional tasks that require precise attribute binding, spatial reasoning, and object relationships — exactly the capabilities GenEval measures.

• Sequential optimization underperforms joint optimization. The two-stage pipeline ReFL (w/ Thinking) + TextGRPO achieves only 0.7987 TA and 0.87 GenEval — worse than FlowGRPO ✓ alone (0.8208 TA) and substantially below UniGRPO (0.8381 TA). This is strong evidence that the coupling matters: optimizing the image policy first, then optimizing the text policy on top of the (now-stale) image optimization, creates a mismatch that joint optimization avoids.

• Algorithm choice determines whether joint optimization is possible at all. UniFPO (the FPO-based analog of UniGRPO) is marked with "—" in Table 1, indicating training collapse. The paper states "UniFPO failed to converge, underscoring the stability advantages of our GRPO-based formulation." This is a critical finding: the specific RL algorithm for the continuous visual component is not an interchangeable detail — FlowGRPO's SDE-based exploration and RatioNorm clipping provide stable gradients that couple cleanly with text GRPO, while FPO's alternative approach does not.

Detailed GenEval Sub-Category Analysis (Table 2)

Table 2 provides the fine-grained breakdown across GenEval's six sub-categories, revealing where UniGRPO's gains concentrate:

• Single Object and Two Objects: All methods perform near-ceiling (0.98–1.00 for Single Object, 0.92–0.99 for Two Objects). UniGRPO achieves 0.99 on both, matching the best baselines. These categories do not discriminate between methods.

• Counting: This category shows the largest variance. SFT ✓ exhibits a dramatic drop to 0.63 (from SFT × at 0.83), and several RL methods with thinking also degrade on counting compared to their non-thinking counterparts (ReFL ✓ drops to 0.63 vs. ReFL × 0.86; FPO ✓ drops to 0.69 vs. FPO × 0.90). UniGRPO achieves 0.91 — matching FPO × but substantially above FlowGRPO ✓ (0.76). The paper does not analyze this counting degradation for thinking-enabled methods; it may reflect that Bagel's reasoning module was "primarily trained for knowledge-based reasoning and may not be ideally suited for short prompt rewriting tasks found in benchmarks like GenEval" (the paper cites a Bagel GitHub issue comment for this observation). UniGRPO appears to overcome this limitation through joint RL optimization, recovering strong counting performance.

• Colors: Relatively narrow spread (0.84–0.94). UniGRPO achieves 0.91, slightly below the best baseline (ReFL ✓ at 0.94) but well above FlowGRPO ✓ (0.90).

• Position: This category shows substantial differentiation. TextGRPO ✓ achieves 0.76 — the highest among all methods, including UniGRPO (0.73). This suggests that position-understanding is primarily a function of the reasoning text quality, and optimizing the text policy alone is sufficient for this sub-task. FlowGRPO ✓ achieves 0.71, and SFT ✓ achieves 0.68. The fact that UniGRPO (0.73) does not match TextGRPO ✓ (0.76) on position is notable: joint optimization may create a slight trade-off where position reasoning is slightly muted in favor of other compositional aspects that contribute more to the overall reward.

• Attribute Binding: Another category where UniGRPO excels. UniGRPO achieves 0.86, matching FlowGRPO × (0.86) and FPO × (0.86). FlowGRPO ✓ achieves 0.84, TextGRPO ✓ achieves 0.84. The gain from 0.84 to 0.86 is modest but consistent across the top methods.

Training Dynamics (Figure 3)

The training and validation reward curves in Figure 3 show UniGRPO's learning dynamics on the SFT Bagel base model at 1024 resolution. The x-axis represents gradient update steps (not epochs or iterations), and both curves show a characteristic RL training pattern:

• Training reward increases monotonically and smoothly, showing that the policy is successfully learning to generate images that score higher under the reward model. There is no evidence of the sudden drops or oscillations that would indicate training instability.

• Validation reward also increases but appears to plateau or show slower growth in later steps compared to the training curve, suggesting that some degree of overfitting to the reward model's preferences may be occurring. The paper does not provide explicit early-stopping criteria based on validation performance, but the final evaluation uses the trained model at its terminal state (as reflected in Table 1).

• The curves do not exhibit the "reward hacking" signature (validation reward declining while training reward continues to rise) that the regularization ablation (Figure 5) demonstrates for the no-KL and latent-KL conditions. This confirms that the velocity MSE regularization effectively prevents the catastrophic degradation seen in the ablation.

Qualitative Analysis (Figures 2, 6, 7, 8)

The paper provides several qualitative comparisons that complement the quantitative results:

• Image quality progression (Figure 2): The original Bagel generates images with "oversaturated colors and noticeable synthetic artifacts." SFT "helps mitigate these synthetic artifacts" but "compromises image sharpness, resulting in noticeable blurriness upon close inspection." UniGRPO "significantly enhances both aesthetic quality and text-image alignment, yielding photorealistic, finely detailed images that faithfully reflect complex user prompts."

• Reasoning trace evolution (Figures 6, 7, 8): The paper presents side-by-side comparisons of reasoning traces across training stages. The original Bagel and SFT models "generate detailed reasoning texts" that "can sometimes lose focus or become disconnected from the core visual generation task." In contrast, "UniGRPO optimizes the reasoning phase to be highly purposeful and task-oriented," producing "reasoning traces that tightly couple with and effectively guide the subsequent image synthesis." The paper's Figure 6 caption states: "Our joint RL optimization produces task-oriented reasoning that guides the synthesis policy toward faithful, photorealistic image generation."

This qualitative evidence supports the paper's claim that joint optimization does not merely improve images or reasoning independently, but rather aligns the reasoning process to serve visual synthesis — the thinking becomes more instrumental, more focused on visual details, and more directly helpful for generating the image.

Ablation Studies and Robustness Checks

CFG-free training (Figure 4): The paper compares training UniGRPO with and without CFG, evaluating both conditions with CFG applied at inference. The key finding is that removing CFG during training yields comparable or better final performance when evaluated with CFG, even though "CFG during training yields images with higher [in-training] rewards." This demonstrates that CFG is a training-time convenience (it helps generate higher-quality exploration samples) but not a necessary component for the policy to learn alignment — the reward signal and policy update internalize the alignment directly into the model weights. This ablation uses the original Bagel (not the SFT checkpoint) at resolution 512, so the absolute reward values are not directly comparable to the main results in Figure 3.

Regularization strategies (Figure 5): Three conditions are compared: (1) no regularization, (2) standard latent KL divergence, and (3) the proposed velocity MSE penalty. The results demonstrate a clear gradient of effectiveness:

• No KL: The validation reward drops after an initial increase, and generated images show "unnatural, oversaturated textures" — the classic signature of reward hacking. The policy learns to exploit blind spots in the reward model, producing images that score highly but are visually degraded.
• Latent KL: The "significant drop in training reward indicates that a sufficiently large KL coefficient has been used" — the penalty is constraining the policy. However, "grid-like artifacts still emerge as early as step 250, prompting us to terminate this run early." This is the crucial negative result: even with a KL coefficient large enough to substantially reduce the training reward (and thus constrain policy updates), the latent KL fails to prevent a specific form of visual degradation. The paper attributes this to the temporal non-uniformity of the latent KL penalty — the inverse noise variance weighting $(1/\sigma_{t_k}^2)$ makes regularization weakest at high-noise early timesteps, creating "temporal vulnerabilities that the RL optimizer can easily exploit."
• Velocity MSE (proposed): Training dynamics are stable, and generated images maintain "high-fidelity images with realistic textures." The uniform penalty across all timesteps closes the loopholes that the latent KL left open, despite using a very small penalty weight ($1.5 \times 10^{-5}$).

UniFPO training collapse (Table 1): The failure of UniFPO (marked as "—" for all metrics) serves as an implicit ablation demonstrating that the choice of image RL algorithm is critical for joint optimization stability. FPO's approach — using the forward process to compute ELBO-based importance sampling weights — does not couple stably with TextGRPO in the unified MDP framework. The paper attributes this to "the stability advantages of our GRPO-based formulation," specifically the combination of SDE exploration (injecting controlled stochasticity) and RatioNorm clipping (centering the importance ratio distribution).

Thinking vs. non-thinking comparisons (Table 1, Table 2): While not presented as a formal ablation, the comparison between methods with (✓) and without (×) thinking tokens reveals an important pattern: enabling thinking does not consistently improve GenEval scores for non-RL or single-modal RL methods. As the paper notes, "enabling the explicit reasoning chain on the Bagel model does not consistently improve GenEval scores" — the Bagel authors' reasoning module "is primarily trained for knowledge-based reasoning and may not be ideally suited for short prompt rewriting tasks found in benchmarks like GenEval." However, UniGRPO (which uses thinking) achieves the best GenEval score (0.90), suggesting that the RL optimization repurposes the reasoning chain from general knowledge reasoning to task-specific visual planning, overcoming the initial misalignment.

Resolution scaling: The main experiments use 1024 × 1024 resolution training (Table 3), while the CFG ablation uses 512 × 512. The paper does not provide a direct resolution ablation, but the fact that UniGRPO trains and achieves SOTA at 1024 resolution — notably higher than most prior RL-for-image-generation work — is itself an implicit demonstration that the framework scales to high-resolution training without instability.

Critical Assessment

The paper makes three central claims that can be evaluated against the experimental evidence:

Claim: Joint optimization of reasoning and visual synthesis under a unified MDP improves generation quality beyond single-modality optimization. This claim is strongly supported by Table 1, which shows UniGRPO (0.8381 TA, 0.90 GenEval) outperforming both the best text-only optimization (TextGRPO at 0.8078 TA, 0.88 GenEval) and the best image-only optimization (FlowGRPO with thinking at 0.8208 TA, 0.86 GenEval). The sequential two-stage approach (ReFL + TextGRPO at 0.7987 TA) performs worse than image-only optimization, demonstrating that the coupling must be simultaneous, not sequential. The qualitative evidence (Figures 2, 6, 7, 8) further supports that joint optimization produces reasoning that is more task-oriented and images that are more faithful.

However, there is an important nuance: the paper demonstrates joint optimization on a single model (Bagel after SFT) on a single task type (text-to-image generation from prompts). The claim that joint optimization is superior has been demonstrated only in this specific setting. Whether the same gains would hold on a different unified architecture (Show-o, Transfusion, Mogao), on different generation tasks (editing, multi-round dialogue), or with different reward models remains untested. The paper acknowledges this limitation by framing the work as validation on "the fundamental unit" and explicitly naming multi-round scaling as future work.

Claim: CFG elimination during training enables linear unbranched rollouts essential for scaling to multi-condition generation, without sacrificing final quality. The ablation in Figure 4 partially supports this claim: training without CFG achieves comparable or better performance when evaluated with CFG. However, the evidence has important limitations:

1. The ablation uses a different base model and resolution (original Bagel at 512, not SFT Bagel at 1024), making it an imperfect proxy for the main experiments.

2. The scalability benefit is asserted, not demonstrated. The paper argues that CFG elimination is "essential for scaling to complex scenarios involving multi-turn interactions and multi-condition generation," but no multi-round or multi-condition experiment is conducted. The claim that CFG removal "naturally scales" to these scenarios is a design argument, not an empirical finding. Until UniGRPO is tested on image editing (3+ conditions) or multi-round generation (alternating text-image phases), the scalability claim remains a well-motivated but unvalidated hypothesis.

3. All evaluation still uses CFG at inference. The paper's claim is that CFG is unnecessary for training, not for inference. This means the practical deployment benefit (reduced inference cost) is not realized in the current evaluation setup. A stronger demonstration would show that CFG-free inference after CFG-free RL training achieves comparable quality — which would directly translate to 2× inference speedup. The paper does not report this comparison.

Claim: Velocity MSE regularization provides more robust reward hacking mitigation than latent KL divergence. The ablation in Figure 5 provides compelling evidence for this claim, demonstrating a clear gradient of failure modes: no regularization → reward hacking (oversaturated textures); latent KL → grid-like artifacts despite constraining the policy; velocity MSE → stable training and high-fidelity images. The specific diagnosis — that latent KL's inverse noise variance weighting creates temporal vulnerabilities — is supported by the empirical observation that artifacts emerge even when the KL coefficient is large enough to significantly reduce training reward.

However, the paper does not explore alternative fixes within the KL framework. For example, a timestep-dependent KL coefficient (increasing the penalty weight at high-noise timesteps to compensate for the $1/\sigma_{t_k}^2$ weighting) might address the temporal vulnerability without abandoning the KL entirely. The velocity MSE is presented as a superior alternative, but the comparison is limited to three conditions (none, latent KL, velocity MSE) without exploring intermediate designs that might combine the probabilistic interpretation of KL with the temporal uniformity of MSE. This leaves open the question of whether the velocity MSE's effectiveness comes from the uniform weighting specifically, or from regularizing the velocity field rather than the latent transition distribution.

Genuine weaknesses and gaps in the experimental design:

• Single model family, single architecture. All experiments use Bagel, a specific unified architecture combining AR text with flow-matching images. The paper positions this as representative of the architectural direction the community is converging on, but does not demonstrate transfer to other unified architectures (e.g., Show-o's combined AR-diffusion training, Transfusion's mixed next-token and diffusion losses). The findings about CFG elimination, velocity MSE, and joint optimization stability may be specific to Bagel's particular architectural choices and training dynamics.

• Small and partially internal evaluation set. The TA benchmark is an internal 150-prompt set, and while the paper references RewardDance for details on the scoring mechanism, the prompts themselves and their distribution (easy vs. hard, compositional complexity, domain coverage) are not characterized. GenEval is a standard benchmark but evaluates specific compositional capabilities; it does not measure aesthetic quality, photorealism, or diversity — all of which are claimed improvements in the qualitative analysis. A more comprehensive evaluation would include benchmarks like DrawBench, T2I-CompBench, or human preference studies.

• The reward model is fixed and shared across all methods. The paper uses the same differentiable InternVL-based reward model from RewardDance for all methods, specifically to enable fair comparison with gradient-based baselines like ReFL. However, this means all methods are optimizing for the same (potentially imperfect) reward signal. The paper's claim that UniGRPO is "compatible with a broader range of verifier-based rewards" is stated but never tested — no experiment uses a non-differentiable reward model, a VLM judge with different architecture, or a multi-aspect reward model. The robustness of the unified framework to different reward formulations remains unvalidated.

• No human evaluation. Despite claiming improvements in "photorealism" and "aesthetic quality" (Section 5.2, Figure 2), no human preference study or user study is reported. The reward model scores and GenEval metrics serve as proxies for human judgment, but the paper's qualitative claims about blurriness, artifacts, and photorealism would be strengthened by direct human evaluation, especially given the known phenomenon of reward model over-optimization in RL for generative models.

• Training cost and compute scaling are not reported. The paper does not report GPU hours, FLOPs, or wall-clock time for training UniGRPO or any baseline. The CFG-free training is claimed to reduce computational cost (halving forward passes per denoising step), but this efficiency gain is never quantified. For practitioners deciding whether to adopt this framework, the training cost relative to alternatives is a critical piece of missing information.

• The thinking module's initial misalignment is handled by RL, but not analyzed. The paper notes that Bagel's reasoning module "may not be ideally suited for short prompt rewriting tasks found in benchmarks like GenEval" and that GenEval scores sometimes degrade when thinking is enabled (e.g., SFT ✓ counting drops to 0.63 from SFT × 0.83). UniGRPO recovers this performance (0.91 counting), but the paper does not analyze how the RL optimization resolves this misalignment — does it change what the reasoning module attends to? Does it shorten the reasoning traces? Does it learn to produce different types of reasoning for different task categories? Understanding this recovery mechanism would strengthen the claim that joint RL optimization is the right tool for aligning reasoning to visual synthesis.

Missing experiments that would strengthen the paper:

• Multi-round or multi-condition experiments. The paper's design choices (CFG elimination, velocity MSE) are motivated by multi-round scalability, but no multi-round experiment is conducted. Even a simple two-round setting (Prompt → Think → Image → Reflect → Edit Image) would provide preliminary evidence that the framework scales as claimed.

• CFG-free inference evaluation. Directly comparing inference quality with and without CFG after CFG-free RL training would quantify the practical deployment benefit and test whether the model truly internalizes alignment into its weights.

• Ablation on the SDE window size and noise level. The paper sets SDE Window Size = 3 and noise level = 0.8 based on prior work (FlowGRPO-Fast), but these choices might interact with the unified optimization dynamics. A sweep over window sizes (1, 3, 5, 10) would reveal whether the exploration-vs-stability tradeoff changes when text and image policies are jointly optimized.

• Ablation on the text-image objective weight λ. The paper sets λ = 1 "to avoid over-design," but the different learning rates for text (1e-6) and image (3e-5) components suggest asymmetric sensitivity to gradient updates. A sweep over λ values (0.1, 0.5, 1, 2, 10) would reveal whether the equal weighting is truly optimal or merely a convenient default.

• Comparison with preference-based methods (DPO variants). The baseline set includes ReFL, FPO, and FlowGRPO, but does not include DPO-based alignment methods (e.g., Diffusion-DPO, SPO) that are widely used for diffusion model alignment. A DPO baseline would contextualize the RL-based approach within the broader alignment landscape.

• Diversity metrics. The paper focuses exclusively on quality and prompt-alignment metrics. RL fine-tuning is known to potentially reduce output diversity (mode collapse). Reporting diversity metrics (e.g., LPIPS diversity, CLIP embedding variance across samples for the same prompt) would address the concern that improved TA and GenEval scores come at the cost of generating more homogeneous outputs.

In summary, the experimental evidence strongly supports the paper's core finding — that joint optimization under a unified MDP with compatible per-modality RL algorithms outperforms separate or sequential optimization — within the tested scope (single-round text-to-image on Bagel). The evidence for the scalability-motivated design choices (CFG elimination, velocity MSE) is solid for the single-round setting but the scalability claims themselves remain unvalidated. The velocity MSE regularization ablation is the most rigorous and compelling part of the experimental section, clearly demonstrating a failure mode in the standard approach and providing a well-motivated fix. The paper's candid acknowledgment of limitations (single benchmark, single model family, unexplored multi-round scaling) is commendable and appropriately scopes the contributions.

6. Limitations and Trade-offs

The Difficulty Estimation Overhead Is Not Accounted For

The assumption or constraint. The compute-optimal scaling framework requires estimating each prompt's difficulty before allocating the test-time compute budget. The paper's method for doing this — generating 2048 samples per question and computing pass@1 (oracle) or averaging PRM final-answer scores (predicted) — is extraordinarily expensive. The paper acknowledges this openly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The consequence. The paper's headline efficiency gains (e.g., "more than 4× better efficiency over best-of-N") are computed after difficulty is already known, without amortizing the cost of estimating it. Generating 2048 samples per question is itself comparable to or greater than the largest test-time compute budgets studied in the paper (256–512 generations). In a realistic deployment, the total cost would be difficulty_estimation + strategy_execution, and the former could dominate the latter, potentially eroding or eliminating the reported efficiency advantages. A system that spends 2048 generations estimating difficulty and then saves 48 generations by choosing the right strategy is worse than one that simply runs best-of-N on all problems. This is the single most practically consequential limitation, because it means the $4\times$ figure should be understood as an upper bound on achievable efficiency given perfect (and free) difficulty information, not as a realized deployment gain.

What evidence exists in the paper. The paper provides no experiment that accounts for difficulty estimation cost in the total budget. The compute-optimal scaling curves in Figures 4 and 8 treat difficulty bins as given. The predicted difficulty bins are generated using the same expensive 2048-sample procedure, simply replacing ground-truth correctness with PRM final-answer scores — the sampling cost is identical. There is no ablation testing performance with fewer difficulty-estimation samples (e.g., 16, 64, 256 rather than 2048), and no experiment where the difficulty estimation budget is subtracted from the per-question compute allowance.

Mitigation status. The paper explicitly flags this as a limitation and suggests future work on "training models to directly predict difficulty of a question" (Section 8). However, no such model is developed or evaluated. The paper also frames difficulty estimation as an exploration-exploitation tradeoff (Section 3.2) — spending compute on assessing difficulty vs. solving the problem — but does not study this tradeoff empirically. Until a cheap difficulty estimator is demonstrated, the compute-optimal framework remains a conceptual contribution rather than a deployable system.

---

The $14\times$ Larger Model Baseline Is Not Compute-Optimally Trained

The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters while holding training data fixed, following the LLaMA paradigm rather than Chinchilla-optimal training. The paper states:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the larger model uses only greedy decoding — no majority voting, best-of-N, or any form of test-time compute augmentation.

The consequence. The reported advantages of test-time compute over the larger model (e.g., "+27.8% on easy questions at $R \ll 1$" in Figure 1) are measured against a baseline that may be systematically weaker than what compute-optimal pretraining would produce. A Chinchilla-optimal model trained with $14\times$ more total FLOPs — scaling both parameters and data equally — would likely achieve higher accuracy than a parameter-only-scaled model, potentially shrinking or reversing the reported advantages. Furthermore, the larger model receives no test-time compute budget of its own, making the comparison asymmetric: the smaller model gets sophisticated adaptive allocation while the larger model gets greedy decoding. A fairer comparison would give the larger model some test-time compute (e.g., best-of-8, majority voting, or its own compute-optimal strategy at proportional budget), which would raise the bar and potentially narrow the advantage margins.

What evidence exists in the paper. The FLOPs-matched results appear in Figure 9 and the bar charts in Figure 1, broken out by difficulty and by $R = D_{\text{inference}} / D_{\text{pretrain}}$ regime. The paper reports that test-time compute with the smaller model outperforms the $14\times$ larger model primarily on easy-to-medium questions at low-to-moderate $R$ values. On the hardest questions (bin 5) and at high $R$, pretraining is clearly preferable — even against this potentially suboptimal baseline. The paper does not provide results for a Chinchilla-optimal pretraining baseline or for a larger model augmented with even modest test-time compute.

Mitigation status. The paper is transparent about this limitation, explicitly acknowledging the departure from compute-optimal pretraining and leaving the fairer comparison to future work. However, the asymmetry in the comparison (small model gets RL, large model gets greedy) is never discussed as a weakness. The claim that "test-time compute can substitute for pretraining compute" should be understood as conditional on the specific pretraining scaling paradigm used (parameter-only scaling with greedy decoding), not as a universal statement about the training-inference tradeoff.

---

Results Are Restricted to a Single Benchmark and Single Model Family

The assumption or constraint. All experiments use the MATH benchmark (500 test questions) and the PaLM 2-S* model family. The paper states it believes this model is "representative of the capabilities of many contemporary LLMs" (Section 4), but provides no replication on other benchmarks, model families, or task domains.

The consequence. Several aspects of the paper's findings may not generalize:

• The PRM's behavior and over-optimization characteristics depend on PaLM 2-S*'s output distribution. A model with different calibration, error patterns, or knowledge boundaries might exhibit different difficulty-dependent scaling curves, potentially shifting the optimal strategy per bin or changing the difficulty thresholds at which beam search becomes preferable to best-of-N.

• The MATH benchmark consists exclusively of competition-level math problems requiring multi-step symbolic reasoning. The core finding — that beam search hurts on easy problems but helps on medium ones, and that sequential revisions dominate on easy problems — might not transfer to other reasoning domains (code generation, logical deduction, scientific QA) where the structure of the solution space differs, or to tasks requiring factual recall rather than inference.

• The difficulty bin definitions are model-relative (computed from PaLM 2-S*'s pass@1 on the test set). A model with different overall MATH accuracy (higher or lower) would produce different bin boundaries, and the strategy rankings per bin might shift. Without replication on other models, the paper's bin-to-strategy mappings cannot be assumed to hold.

• The 500-question test set, when split into five difficulty quintiles and further split by two-fold cross-validation, means the compute-optimal policy is selected based on approximately 50 questions per fold per bin. This is a small sample, and the selected strategies may not be robust — a different random split could produce different policy recommendations. The paper does not report confidence intervals or variance estimates for the compute-optimal scaling curves.

What evidence exists in the paper. The paper's experimental scope is characterized in Section 4 (dataset and model) and all quantitative results (Figures 3, 4, 6, 7, 8, 9; Tables 1, 2) are on MATH with PaLM 2-S*. The paper provides no out-of-domain evaluation, no cross-model evaluation, and no sensitivity analysis for the test-set split. The cross-validation protocol (Section 3.2) is described as two-fold within each bin but no error bars are shown on any result.

Mitigation status. The paper does not address this limitation explicitly — it does not claim generalizability to other domains or models, but also does not caution against over-interpreting the MATH-specific findings. The paper frames its contributions as systematic analysis and framework development, not as a universal recipe, but a reader unfamiliar with the domain-specificity caveats might reasonably assume the difficulty-dependent patterns generalize. Future work on code generation, scientific reasoning, or multimodal tasks would be needed to establish broader applicability.

---

Hard Problems Remain Fundamentally Unsolved by Test-Time Compute

The assumption or constraint. The paper's entire framework operates under the implicit assumption that the base model can produce correct solutions at some non-trivial rate. When the base model's pass@1 is near zero — the hardest problems (bin 5) — test-time compute provides essentially no benefit regardless of strategy.

The consequence. The compute-optimal framework offers no path forward for problems that genuinely exceed the base model's capability range. Across all methods — search, revisions, and their compute-optimal combinations — performance on the hardest problems (difficulty bin 5) hovers near 1–3% regardless of budget allocation. In Figure 3 (right), bin 5 accuracy is essentially flat across all search methods and all budgets (4, 16, 64, 256 generations). In Figure 7 (right), bin 5 shows approximately 2–3% accuracy irrespective of the sequential-to-parallel ratio at 128 generations. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line remains essentially flat near 0–5% even as test-time compute is scaled up, while the $14\times$ larger model (even with greedy decoding) performs substantially better.

This is a fundamental capability boundary: test-time compute can amplify existing capability — finding correct solutions that exist in the proposal distribution and refining near-correct answers — but it cannot create capability that was not acquired during pretraining. For problems where the model's initial pass@1 is genuinely zero (no correct solutions in 2048 samples), no amount of search or revision will produce correct answers because there are no correct seeds to find or refine.

What evidence exists in the paper. The difficulty-bin analysis across all major results (Figures 3 right, 7 right, 9) consistently shows bin 5 as an outlier where no method makes meaningful progress. The paper's takeaway box in Section 7 explicitly states: "on the hardest problems, pretraining compute is almost always more effective." The near-zero absolute accuracy on bin 5 also means that the relative percentages reported in the FLOPs-matched comparison (e.g., "-52.9% relative disadvantage" for test-time compute vs. pretraining on hard problems at $R \gg 1$) are computed from very small absolute numbers and should be interpreted carefully.

Mitigation status. The paper is candid about this limitation, explicitly stating that "test-time compute amplifies existing capability but does not create it from nothing" and that "some capabilities can only be acquired through pretraining." However, this boundary condition has practical implications that the paper does not fully explore: if the deployment distribution includes a non-trivial fraction of problems in the bin 5 regime, the compute-optimal policy cannot help, and the only viable approach is deploying a larger pretrained model or routing those problems to human experts. The paper also does not investigate whether there is a detectable signal that a problem is in bin 5 (i.e., truly beyond capability) before spending the full inference budget — such early-detection capability would be valuable for routing decisions.

---

Sequential Revision Strategies Create Latency Bottlenecks Not Accounted For

The assumption or constraint. The paper measures compute in "generations" (number of complete solutions sampled), which is a reasonable proxy for total FLOPs but ignores wall-clock time. Sequential revision strategies — where each revision conditions on the previous output — are inherently serial and cannot be parallelized across independent hardware, regardless of available compute.

The consequence. The compute-optimal policy's recommendations for easy and medium problems favor sequential-heavy allocations. For example, Section 6 reports that easy problems perform best with purely sequential revisions, and medium problems with a balanced sequential-to-parallel ratio. A deployment using 64 sequential revisions requires $64\times$ the wall-clock time of a parallel strategy generating 64 independent samples simultaneously with sufficient hardware. This means the reported $4\times$ efficiency gains in terms of total computation may correspond to substantially worse latency than the parallel baselines being compared against, making the compute-optimal strategy impractical for latency-sensitive applications (interactive assistants, real-time systems, user-facing APIs).

This is a fundamental tradeoff the paper does not discuss: the optimal allocation for throughput (total FLOPs per correct answer) may be very different from the optimal allocation for latency (wall-clock time per query). Sequential revisions trade latency for sample efficiency — each revision can build on previous attempts, potentially requiring fewer total samples, but the serial dependency means longer response times. The paper measures and optimizes only the sample efficiency side of this tradeoff.

What evidence exists in the paper. The paper's compute accounting is based entirely on generation counts (Section 3.1 defines the budget $N$ in generations; Section 5.3 sweeps budgets as powers of 2; Section 6 compares sequential, parallel, and hybrid allocations at equal total generation counts). The paper never reports wall-clock time, throughput, or latency for any configuration. The compute-optimal policies in Figures 4 and 8 select strategies based solely on accuracy at a given generation budget, with no latency constraint. The FLOPs-matched comparison in Section 7 similarly uses total FLOPs as the only resource metric.

Mitigation status. The paper does not address this limitation. The term "latency" does not appear in the paper; "wall-clock" or "response time" are never mentioned. The paper frames the contributions around computational efficiency, not latency optimization, and the omission is understandable for a methods analysis paper — but practitioners deploying these techniques in production would need to evaluate the latency-accuracy tradeoff separately. A latency-constrained version of the compute-optimal allocation (e.g., maximize accuracy subject to a maximum number of sequential steps) is suggested by the framework but never explored.

---

Revisions and Search Are Studied Independently, Not Combined

The assumption or constraint. The paper studies two complementary mechanisms for test-time compute scaling — PRM-guided search (Section 5) and iterative revisions (Section 6) — but never combines them into a single system. The paper acknowledges this gap:

"we did not experiment with PRM tree-search techniques in combination with revisions" (Section 8)

The consequence. The paper's two main scaling axes have complementary, difficulty-dependent strengths. Revisions (proposal distribution modification) are most effective on easy problems where the model's initial output is roughly correct and needs targeted refinement — essentially a local search in answer space. PRM-guided beam search (verifier optimization) is most effective on medium-hard problems where the model needs to explore qualitatively different solution strategies — a global search. The paper demonstrates that each mechanism individually, when allocated optimally per difficulty, achieves $4\times$ efficiency gains over best-of-N (Figures 4 and 8). But the two mechanisms are never combined: there is no experiment using the revision model as the proposal distribution within beam search (where each search step conditions on previous rejected branches), no experiment using the PRM to guide which revision directions to pursue, and no experiment that jointly selects the search-revision allocation per difficulty bin.

This means the current results represent a lower bound on what a fully integrated system could achieve. The optimal strategy for a given difficulty bin might involve both revision depth and search breadth — for example, using a small number of parallel revision chains (exploiting the revision model's refined proposal distribution) with beam search applied to each chain (exploiting the PRM's ability to select promising intermediate steps). The paper's compute-optimal policies select between search and revisions, but never combine them within a single strategy.

What evidence exists in the paper. Section 5 presents results for search methods (best-of-N, beam search, lookahead search) applied to the base LLM's outputs. Section 6 presents results for revision methods (sequential, parallel, hybrid) applied to the revision model's outputs. Neither section references the other's methods. The compute-optimal policies in Figures 4 and 8 are derived from separate experiments — the search policy selects among search algorithms, and the revision policy selects among sequential-to-parallel ratios, with no cross-contamination. Table 1 in the FLOPs-matched comparison shows improvements for both search and revisions separately, but never a combined estimate.

Mitigation status. The paper explicitly acknowledges this gap in Section 8 ("Conclusion and Future Work") and identifies it as a direction for future work. However, the paper's central claim — that difficulty-conditioned allocation yields $4\times$ efficiency gains — is demonstrated only for each axis independently. A reader might reasonably assume that the gains would be additive or multiplicative if combined, but the paper provides no evidence for this. It is also possible that combining search and revisions would create new interaction effects (e.g., the revision model's outputs might be distributed differently, changing the PRM's over-optimization threshold) that could reduce or alter the gains. The separate analysis is methodologically clean but leaves the practical question of how to build a fully unified system unresolved.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper makes a methodological intervention rather than an algorithmic breakthrough. It does not propose a new policy gradient estimator, a new reward model, or a new architecture for multimodal models. Instead, it demonstrates that the standard RL recipes that already work for text-only and image-only generation can be combined into a single unified loop with no additional machinery beyond an MDP formulation — and that doing so produces substantially better results than any scheme that optimizes the two modalities separately or sequentially. This is a reframing of the multimodal post-training problem: from "how do we design a new RL algorithm for interleaved generation?" to "how do we couple existing per-modality algorithms under a shared signal so they reinforce each other?"

The paper's most important conceptual contribution is the demonstration that the unified MDP abstraction is sufficient and that algorithm compatibility is the bottleneck, not reward design or credit assignment. The negative result on UniFPO (training collapse in Table 1) is arguably as important as the positive UniGRPO result. It shows that the specific choice of image RL algorithm determines whether joint text-image optimization is even possible: FlowGRPO's SDE-based exploration and RatioNorm clipping couple stably with text GRPO, while FPO's ELBO-based importance sampling does not. This narrows the search space for future multimodal RL methods: researchers now know that the image component's exploration mechanism and importance ratio normalization are not interchangeable details but compatibility requirements that must be verified against the text policy's optimization dynamics.

The paper also reframes CFG as a training-time convenience rather than a necessary alignment mechanism, with implications extending beyond multimodal generation. The finding that a model trained without CFG can match or exceed the quality of a CFG-trained model when evaluated with CFG at inference (Figure 4) suggests that RL optimization can internalize alignment into model weights, reducing dependence on inference-time guidance tricks. If this transfers to standard (non-multimodal) text-to-image models, it would mean that a substantial fraction of the inference compute currently spent on CFG's dual forward passes could be eliminated by RL post-training — a 2× inference throughput improvement at scale. The paper does not fully validate this claim (all evaluation still uses CFG at inference), but the training-time evidence establishes it as a testable hypothesis that the field should pursue.

The paper also provides a new diagnostic for reward hacking in flow-matching RL: the temporal vulnerability in latent-space KL regularization. By identifying that the inverse noise variance weighting ($1/\sigma_{t_k}^2$) in the standard KL penalty creates exploit-prone regions at high-noise early timesteps, and by showing that an unweighted velocity MSE penalty closes these loopholes (Figure 5), the paper gives future flow-matching RL researchers both a problem to watch for and a simple fix. This diagnostic is not specific to multimodal settings — it applies to any flow-matching or diffusion model trained with policy gradients.

Finally, the paper's "fundamental unit" validation strategy — deliberately testing on single-round generation before scaling to multi-round — provides a methodological template for the field. By isolating the core question (does joint optimization outperform separate optimization?) from the confounds of multi-round credit assignment and context management, the paper produces clean, interpretable results that establish a baseline for future multi-round work. This is a sharp contrast with concurrent work (DualGRPO, SepGRPO) that jumps directly to more complex settings, making it harder to attribute performance to specific design choices.

Follow-Up Research This Work Enables

Multi-round interleaved generation on the same unified architecture. The most direct and urgent follow-up: apply UniGRPO to two-round or three-round interleaved generation (e.g., Prompt → Think → Image → Reflect → Edit Image → Final Output) on Bagel or a similar unified model. This would test the paper's central scalability claim — that CFG-free linear rollouts and velocity MSE regularization enable stable optimization over longer horizons. The key measurement would be whether the joint optimization advantage observed in single-round generation (UniGRPO outperforming TextGRPO and FlowGRPO separately, Table 1) compounds or diminishes over multiple rounds. A failure mode to watch for: as the horizon grows, errors in early-round reasoning may compound, and the sparse terminal reward at the final image may provide too weak a signal to assign credit across multiple text-image cycles. A strong positive result would be UniGRPO achieving comparable or larger relative gains over single-modality baselines in a two-round editing task compared to the single-round gains reported in Table 1. A negative result (joint optimization benefit shrinks with horizon) would indicate that intermediate rewards or process reward models (as the paper suggests in Section 6) are necessary for longer-horizon settings.

CFG-free inference evaluation after CFG-free RL training. The paper trains without CFG but evaluates with it. A critical follow-up is to close this loop: take the UniGRPO-trained model, evaluate it at inference time with CFG scale = 1 (no guidance), and compare against the same model evaluated with CFG scale > 1. This would directly measure how much of the alignment has been internalized into the weights versus how much still depends on the inference-time guidance trick. The specific experiment: evaluate UniGRPO on GenEval and the TA benchmark at CFG scales of 1, 2, 3, 4, 7.5, and compare against the baseline methods similarly swept. If the CFG-free-trained model achieves near-peak performance at CFG scale 1–2 (while CFG-trained baselines need scale 4–7.5 to reach their peaks), this demonstrates genuine internalization of alignment and implies a practical 2–3× inference throughput improvement. If the CFG-free-trained model still requires high CFG scales to match baselines, the internalization claim is weaker than the paper suggests.

Compatibility testing across unified architectures. The paper validates on Bagel only. A replication study on other unified multimodal architectures — specifically Show-o (which uses combined AR and discrete diffusion losses), Transfusion (mixed next-token prediction and continuous diffusion), and Mogao (another AR-text + flow-matching model) — would establish whether UniGRPO's stability and performance gains are architecture-specific or general. The specific comparison: replicate the Table 1 experiment (SFT → UniGRPO vs. FlowGRPO vs. TextGRPO vs. ReFL + TextGRPO) on each architecture, using the same reward model and roughly comparable model scales. If UniGRPO consistently outperforms baselines across architectures, the unified MDP framework is robust. If gains are specific to Bagel, the compatibility claims (about SDE exploration coupling stably with text GRPO) may depend on Bagel's particular text-image interface design, cross-attention mechanism, or pretraining data mixture.

Characterizing the text-image coupling mechanism in UniGRPO's gradients. The paper demonstrates that joint optimization works, but does not analyze how the text policy changes during training to better serve image generation. A mechanistic analysis would track: what properties of the reasoning traces change during UniGRPO training? Do traces become shorter or longer? Do they shift from general knowledge reasoning to visual detail specification? Can you intervene by replacing a trained reasoning trace with a random one and measure the image quality drop to quantify how much the image policy depends on the optimized reasoning? The specific experiment: for a held-out set of prompts, generate images using three conditions — (a) the UniGRPO-optimized reasoning trace for that prompt, (b) the SFT model's reasoning trace for the same prompt, (c) a randomly selected reasoning trace from a different prompt. Measure the reward drop from (a) to (b) (alignment gain from optimization) and from (b) to (c) (alignment gain from relevance). This would decompose the joint optimization benefit into "better reasoning content" versus "better conditioning of the image policy on whatever reasoning is provided," revealing whether UniGRPO primarily improves the text policy, the image policy, or their interaction.

Scaling the SDE window and noise level with horizon. The paper uses a fixed SDE window of 3 steps out of 25 training timesteps and a fixed noise level of 0.8, both inherited from FlowGRPO-Fast without modification. In a multi-round setting, the exploration requirements may change: longer horizons may need more exploration (larger SDE windows) to discover diverse multi-round strategies, or less exploration (smaller windows) to maintain stability across alternating text-image phases. A systematic sweep — SDE window sizes of 1, 3, 5, 10, all 25 steps with noise levels of 0.2, 0.5, 0.8, 1.2 — in both single-round and (if available) two-round settings would characterize how the exploration-stability tradeoff changes with horizon. The hypothesis to test: as the generation horizon grows, the optimal SDE window shrinks because accumulated stochasticity across multiple rounds creates excessive variance in the terminal reward signal, making velocity MSE regularization more important relative to exploration.

Multimodal process reward models (PRMs) for dense credit assignment. The paper explicitly names this as a future direction (Section 6), and UniGRPO's sparse terminal reward structure makes it a natural testbed. The specific research question: does providing intermediate rewards on the reasoning trace (e.g., a VLM judge evaluating whether the reasoning text correctly identifies the visual elements needed, before any image is generated) improve sample efficiency and/or final performance compared to terminal-only rewards? The experiment would train a separate VLM-based PRM to score reasoning traces, add its per-token or per-step rewards to the UniGRPO objective (breaking the sparsity of the terminal reward), and compare convergence speed and final GenEval/TA scores against the sparse-reward baseline. A positive result would demonstrate that PRMs are worth the additional training cost for multimodal RL. A negative result (dense rewards fail to improve or hurt performance) would suggest that the terminal reward already provides sufficient signal and that PRM-based dense rewards introduce harmful bias or reward hacking vulnerabilities.

Practical Applications and Downstream Use Cases

Post-training pipelines for unified multimodal models deployed in production. The most immediate application: any team that has pretrained or SFT-finetuned a unified text-image model (Bagel, Mogao, or a future open-source model) and wants to improve its generation quality and prompt alignment can apply UniGRPO as a drop-in post-training step. The paper's minimalist recipe — take the SFT checkpoint, run UniGRPO with equal text/image weighting, no CFG during training, velocity MSE regularization with weight 1.5e-5, group size 24 — provides a concrete configuration that has been demonstrated to work. The expected gain is on the order of +0.06 on TA score and +0.08 on GenEval over the SFT baseline (comparing SFT ✓ at 0.7769/0.82 to UniGRPO at 0.8381/0.90 in Table 1). The primary deployment consideration is that CFG-free training roughly halves the GPU memory and compute per denoising step during training (one forward pass instead of two), which partially offsets the cost of the RL loop's multiple samples per prompt.

High-resolution image generation with aligned reasoning. The paper demonstrates training at 1024 × 1024 resolution — higher than most prior RL-for-image-generation work. This means the framework is directly applicable to production-quality image generation systems, not just proof-of-concept low-resolution experiments. A team deploying a text-to-image API (where users provide prompts and expect high-quality, prompt-adherent images) could use UniGRPO to improve the model's ability to understand and execute complex compositional prompts. The GenEval gains are particularly relevant here: UniGRPO's 0.90 overall score, with strong performance on counting (0.91), position (0.73), and attribute binding (0.86), indicate improved handling of precisely the kind of detailed compositional instructions that users frequently provide and that current models often fail on.

Data generation for self-improvement loops in multimodal models. When generating training data for further SFT or RL iterations (a STaR-like or ReST-like self-improvement loop), the quality and diversity of generated reasoning-image pairs matter critically. UniGRPO provides a way to generate higher-quality reasoning traces that are more tightly coupled to visual synthesis — exactly the kind of data that would be valuable for training a model to produce better reasoning in future iterations. A self-improvement loop could alternate between: (1) using UniGRPO to generate high-quality reasoning-image pairs on a large prompt dataset, (2) filtering for the highest-reward pairs, (3) SFT-finetuning the base model on those pairs, and (4) repeating. The UniGRPO step provides the exploration and reward optimization that discovers better reasoning strategies, while the SFT step distills those strategies back into the base model weights efficiently.

Deployment of computationally efficient inference pipelines. If the CFG internalization hypothesis holds — i.e., if a UniGRPO-trained model can achieve competitive quality with CFG scale 1 or 2 at inference rather than the standard 4–7.5 — this translates to a direct 2–3× reduction in inference compute and latency per generated image. For a production image generation API serving millions of requests, this halves GPU costs or doubles throughput on the same hardware. The caveat is that the paper has not yet demonstrated CFG-free inference quality; this application depends on the follow-up experiment proposed above (CFG scale sweep at inference for UniGRPO-trained models). If validated, this would be the single highest-impact practical outcome of the paper.