MaskFocus: Focusing Policy Optimization on Critical Steps for Masked Image Generation

摘要

Reinforcement learning (RL) has demonstrated significant potential for post-training language models and autoregressive visual generative models, but adapting RL to masked generative models (MGMs) remains challenging. The core factor is that policy optimization requires the probability likelihood of each step due to its multi-step iteration process. This reliance on entire sampling trajectories introduces high computational cost, whereas natively optimizing random steps often yields suboptimal results. In this paper, we present MaskFocus, a novel RL framework that achieves effective policy optimization for MGMs by focusing on critical steps. Specifically, we determine the step-level information gain by measuring the similarity between the intermediate images at each sampling step and the final generated image. Crucially, we leverage this to identify the most critical and valuable steps and execute policy optimization on them. Furthermore, we design a dynamic routing sampling based on entropy to encourage the model to explore more valuable masking strategies. Extensive experiments on multiple Text-to-Image benchmarks validate the effectiveness of our method.