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

TL;DR

MaskFocus introduces a reinforcement learning framework that enhances masked image generation by identifying and optimizing critical steps, leading to improved efficiency and quality in text-to-image synthesis.

Contribution

The paper presents a novel RL approach that focuses on critical steps in masked generative models, utilizing step-level information gain and dynamic sampling to improve image generation.

Findings

Effective policy optimization on critical steps improves image quality.

Dynamic routing sampling encourages exploration of valuable masking strategies.

Validated on multiple Text-to-Image benchmarks.

Abstract

Reinforcement learning (RL) has demonstrated significant potential for post-training language models and autoregressive visual generative models, but adapting RL to masked generative models remains challenging. The core factor is that policy optimization requires accounting for the probability likelihood of each step due to its multi-step and iterative refinement 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 masked generative models 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 focused policy optimization on them. Furthermore, we design a dynamic routing sampling mechanism based on entropy to encourage the model to explore more valuable masking strategies for samples with low entropy. Extensive experiments on multiple Text-to-Image benchmarks validate the effectiveness of our method.