Co-GRPO: Co-Optimized Group Relative Policy Optimization for Masked Diffusion Model

TL;DR

This work tackles the mismatch between single-step BERT-style training and trajectory-driven inference in Masked Diffusion Models (MDMs) by formulating generation as a unified Markov Decision Process and applying trajectory-level Group Relative Policy Optimization to jointly optimize both the denoising model and the inference schedule. The proposed Co-GRPO framework introduces a trainable scheduling policy that co-adapts with the model, employing an alternating optimization strategy to stabilize training and avoid backpropagation through multi-step generation. Empirical results on ImageReward, HPSv2, GenEval, and DPG-Bench show substantial gains over Naive GRPO and baselines, including zero-shot improvements on unseen benchmarks, demonstrating strong generalization. The approach offers a practical path to align training and inference in MDMs, enabling higher-quality, more faithful text-to-image generation with learnable schedules and improved efficiency at inference time.

Abstract

Recently, Masked Diffusion Models (MDMs) have shown promising potential across vision, language, and cross-modal generation. However, a notable discrepancy exists between their training and inference procedures. In particular, MDM inference is a multi-step, iterative process governed not only by the model itself but also by various schedules that dictate the token-decoding trajectory (e.g., how many tokens to decode at each step). In contrast, MDMs are typically trained using a simplified, single-step BERT-style objective that masks a subset of tokens and predicts all of them simultaneously. This step-level simplification fundamentally disconnects the training paradigm from the trajectory-level nature of inference, leaving the inference schedules never optimized during training. In this paper, we introduce Co-GRPO, which reformulates MDM generation as a unified Markov Decision Process (MDP) that jointly incorporates both the model and the inference schedule. By applying Group Relative Policy Optimization at the trajectory level, Co-GRPO cooperatively optimizes model parameters and schedule parameters under a shared reward, without requiring costly backpropagation through the multi-step generation process. This holistic optimization aligns training with inference more thoroughly and substantially improves generation quality. Empirical results across four benchmarks-ImageReward, HPS, GenEval, and DPG-Bench-demonstrate the effectiveness of our approach. For more details, please refer to our project page: https://co-grpo.github.io/ .

Figures (6)

• Figure 1: Qualitative and quantitative comparison of Co-GRPO against baseline approaches. Through cooperative optimization of the MDM model and inference schedule, Co-GRPO produces images with markedly superior quality compared to baseline. Detailed prompts are provided in \ref{['supptab:teaser-prompt']}.
• Figure 2: Comparison between the conventional MDM post-training framework and our Co-GRPO. Naive GRPO collects trajectories using a trainable MDM model under a fixed, predefined inference schedule. Our proposed Co-GRPO challenges this convention by cooperatively optimizing both the MDM model and the inference schedule based on the reward feedback.
• Figure 3: Overview of our proposed Co-GRPO. During the trajectory collection phase, both the sampled visual tokens ($\mathbf{V}$) and their associated inference schedule ($\mathcal{A}$) are collected at each step. These trajectories are evaluated by the reward model, and the resulting scores are aggregated and normalized at the group level to compute individual advantages. In the subsequent policy optimization phase, the joint policy is explicitly factorized into a model policy$\pi_{\theta}$ and a scheduling policy$\pi_{\phi}$. By estimating their respective likelihoods and applying an alternating optimization strategy, our approach enables the cooperative refinement of both policies toward improved generation quality.
• Figure 4: Qualitative comparisons of the base model and models optimized with GRPO and Co-GRPO. Co-GRPO generates images with superior aesthetics while better preserving fine-grained visual details compared to both the base model and GRPO.
• Figure 5: Comparison between Meissonic base model and our Co-GRPO trained model across different steps. Our Co-GRPO improves the performance of the model over every steps, with a 4.2 times faster comparied to the base model on the same score.