Co-GRPO: Co-Optimized Group Relative Policy Optimization for Masked Diffusion Model
Co-GRPO optimizes both model and inference schedule, enhancing MDM generation quality.
Key Findings
Methodology
Co-GRPO redefines the generation process of masked diffusion models (MDMs) as a unified Markov Decision Process (MDP), applying Group Relative Policy Optimization at the trajectory level. This method optimizes model and schedule parameters cooperatively without costly backpropagation through the multi-step generation process, aligning training with inference more thoroughly.
Key Results
- On the ImageReward benchmark, Co-GRPO improved the score from 0.942 to 1.122, showing significant performance enhancement.
- On the HPSv2 benchmark, the score increased from 28.83 to 29.37, surpassing the Naive GRPO baseline.
- Co-GRPO demonstrated strong generalization capabilities in zero-shot scenarios on both GenEval and DPG-Bench benchmarks.
Significance
By optimizing the cooperation between model and inference schedule, Co-GRPO significantly enhances the generation quality of MDMs. This method addresses the fundamental mismatch between training and inference, offering substantial academic and industrial value, especially in fields requiring efficient high-quality image generation.
Technical Contribution
Co-GRPO provides a new theoretical framework that treats the model and inference schedule as cooperating policies, utilizing trajectory-level policy gradients. This approach not only improves generation quality but also opens new engineering possibilities.
Novelty
Co-GRPO is the first to treat the generation process of MDMs as a unified MDP, innovatively considering the inference schedule as a trainable action rather than a fixed hyperparameter, contrasting sharply with existing methods, especially in inference schedule optimization.
Limitations
- Co-GRPO may require longer training times in complex scenarios to achieve optimal performance.
- The optimization of the inference schedule may lead to overfitting in some cases.
Future Work
Future research could explore applying Co-GRPO to larger datasets and further optimizing the inference schedule to enhance model robustness.
AI Executive Summary
Masked Diffusion Models (MDMs) have shown great potential in vision, language, and cross-modal generation, but there is a significant discrepancy between their training and inference processes. MDM inference is a multi-step iterative process, while training typically uses a simplified single-step objective. This discrepancy leaves inference schedules unoptimized during training, affecting generation quality.
Co-GRPO addresses this issue by redefining the MDM generation process as a unified Markov Decision Process (MDP). The method applies Group Relative Policy Optimization at the trajectory level, optimizing both model parameters and inference schedule parameters cooperatively, avoiding costly multi-step backpropagation.
Experimental results show that Co-GRPO significantly improves generation quality across multiple benchmarks, particularly on ImageReward and HPSv2. The method not only enhances generation efficiency but also demonstrates strong generalization capabilities, providing new directions for future MDM research.
Deep Analysis
Background
Masked Diffusion Models (MDMs) have made significant progress in image generation. Compared to autoregressive models, MDMs offer significant advantages in generation efficiency, producing high-quality images in fewer iterations. However, there is a significant discrepancy between the training and inference processes of MDMs, affecting their generation quality.
Core Problem
MDM inference is a multi-step iterative process, while training typically uses a simplified single-step objective. This discrepancy leaves inference schedules unoptimized during training, affecting generation quality. Addressing this issue is crucial for improving the generation quality of MDMs.
Innovation
Co-GRPO redefines the MDM generation process as a unified Markov Decision Process (MDP), innovatively considering the inference schedule as a trainable action rather than a fixed hyperparameter. The method applies Group Relative Policy Optimization at the trajectory level, optimizing both model parameters and inference schedule parameters cooperatively.
Methodology
- �� Redefine MDM generation process as a unified MDP.
- �� Apply Group Relative Policy Optimization at the trajectory level.
- �� Optimize both model parameters and inference schedule parameters cooperatively.
- �� Avoid costly multi-step backpropagation.
Experiments
Experiments were conducted on four benchmarks: ImageReward, HPSv2, GenEval, and DPG-Bench. A fixed 48-step inference was used to evaluate model performance across different benchmarks. Results show significant improvements in generation quality with Co-GRPO.
Results
On the ImageReward benchmark, Co-GRPO improved the score from 0.942 to 1.122. On the HPSv2 benchmark, the score increased from 28.83 to 29.37. Co-GRPO also demonstrated strong generalization capabilities on GenEval and DPG-Bench benchmarks.
Applications
Co-GRPO can be applied in fields requiring efficient high-quality image generation, such as autonomous driving, virtual reality, and augmented reality. Its optimized inference schedule significantly improves generation efficiency and quality.
Limitations & Outlook
Co-GRPO may require longer training times in complex scenarios to achieve optimal performance. The optimization of the inference schedule may lead to overfitting in some cases. Future research could explore applying Co-GRPO to larger datasets.
Plain Language Accessible to non-experts
Imagine you're cooking in a kitchen. Traditional MDMs are like having a recipe but only doing one step at a time, then stopping to check the result. Co-GRPO is like having an assistant who not only knows the recipe but can adjust the next steps based on each result. This method makes cooking more efficient, and the final dish tastier.
ELI14 Explained like you're 14
Imagine you're playing a puzzle game. Traditional MDMs are like trying one answer at a time and waiting for the result. Co-GRPO is like having a smart helper who adjusts the strategy based on each attempt, helping you solve the puzzle faster. This method makes the game more fun, and you win quicker!
Glossary
Markov Decision Process (MDP)
A mathematical model for decision-making processes involving states, actions, and rewards.
Co-GRPO treats MDM generation as an MDP.
Group Relative Policy Optimization (GRPO)
An optimization strategy that improves models by comparing the performance of different policies.
Co-GRPO applies GRPO at the trajectory level.
Masked Diffusion Model (MDM)
An image generation model that produces images by progressively denoising.
MDMs are widely used in vision, language, and cross-modal generation.
Inference Schedule
A plan that determines how many tokens to decode at each step during generation.
Co-GRPO treats the inference schedule as a trainable action.
ImageReward
A benchmark for evaluating image generation quality.
Co-GRPO shows superior performance on ImageReward.
Open Questions Unanswered questions from this research
- 1 How can Co-GRPO be applied to larger datasets?
- 2 Does optimizing the inference schedule lead to overfitting?
- 3 How can model robustness be further improved?
Applications
Immediate Applications
Autonomous Driving
Enhancing environmental perception in autonomous systems through optimized image generation.
Virtual Reality
Generating high-quality images in virtual reality to enhance user experience.
Long-term Vision
Augmented Reality
Applying Co-GRPO in augmented reality to provide more realistic visual effects.
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/ .