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

arXiv: https://arxiv.org/abs/2512.18766 · HTML: https://mlsanity.com/paper/2512.18766/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Placed in RL with diffusion and flow models, one of 15 directions in Deep reinforcement learning methods.
- First posted on arXiv 21 Dec 2025.

## Placed in

- Deep reinforcement learning methods / [RL with diffusion and flow models](https://mlsanity.com/deeprl/rl-with-diffusion-and-flow-models/)

## Built on by (most cited first, 6 of 6)

- [Flow-OPD: On-Policy Distillation for Flow Matching Models](https://mlsanity.com/paper/2605.08063/) (2026)
- [OmniNFT: Modality-wise Omni Diffusion Reinforcement for Joint Audio-Video Generation](https://mlsanity.com/paper/2605.12480/) (2026)
- [Embedding-perturbed Exploration Preference Optimization for Flow Models](https://mlsanity.com/paper/2605.15803/) (2026)
- [Power Reinforcement Post-Training of Text-to-Image Models with Super-Linear Advantage Shaping](https://mlsanity.com/paper/2605.10937/) (2026)
- [MAR-GRPO: Stabilized GRPO for AR-diffusion Hybrid Image Generation](https://mlsanity.com/paper/2604.06966/) (2026)
- [MaskSense: Confronting the Visual Exploration Trap in Masked Image Generation](https://mlsanity.com/paper/2610.17009/) (2026)

## Abstract (arXiv)

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.