Reinforcement Learning Meets Masked Generative Models: Mask-GRPO for Text-to-Image Generation

TL;DR

This work addresses how to apply reinforcement learning to Masked Generative Models (MGMs) for text-to-image generation by redefining the transition probability for MGMs and treating unmasking as a multi-step Markov decision process. It introduces Mask-GRPO, a GRPO-based framework with two transition definitions ($p_{\theta1}$ and $p_{\theta2}$) and enhancements such as removing the KL constraint, reduction strategies, and sample filtering to improve training efficiency and sample quality. Empirical results show substantial gains over a strong base model (Show-o) on GenEval and MSCOCO-FID, with robustness across reward models like CLIP and ImageReward, and strong zero-shot generalization. The findings highlight the potential of RL-based fine-tuning for MGMs and suggest directions for scaling to larger models and extending to video generation, with a continued emphasis on developing more robust reward mechanisms.

Abstract

Reinforcement learning (RL) has garnered increasing attention in text-to-image (T2I) generation. However, most existing RL approaches are tailored to either diffusion models or autoregressive models, overlooking an important alternative: masked generative models. In this work, we propose Mask-GRPO, the first method to incorporate Group Relative Policy Optimization (GRPO)-based RL into this overlooked paradigm. Our core insight is to redefine the transition probability, which is different from current approaches, and formulate the unmasking process as a multi-step decision-making problem. To further enhance our method, we explore several useful strategies, including removing the KL constraint, applying the reduction strategy, and filtering out low-quality samples. Using Mask-GRPO, we improve a base model, Show-o, with substantial improvements on standard T2I benchmarks and preference alignment, outperforming existing state-of-the-art approaches. The code is available on https://github.com/xingzhejun/Mask-GRPO

Figures (5)

• Figure 1: The unsatisfactory performance using the AR-style probability.
• Figure 2: The framework of Mask-GRPO. We redefine the transition probability to formulate the unmasking process as a Markov decision-making problem. The reward is got from CLIP radford2021learning.
• Figure 3: The illustration of transition probability. It is observed that different token samplings finally can lead to the same next state.
• Figure 4: Qualitative Comparisons on GenEval.
• Figure 5: Visualization results of Mask-GRPO.