Accelerating Inference of Masked Image Generators via Reinforcement Learning

TL;DR

MGMs are powerful but slow due to many sampling steps. Speed-RL uses reinforcement learning to directly optimize a joint quality-and-speed objective, adapting GRPO with a low-variance cross-entropy KL estimator to stabilize training on discrete MGMs. Empirical results on a 1B-parameter Meissonic base model show up to 3x inference speedup with comparable ImageReward, HPSv2.1, and CLIP metrics across benchmarks, aided by low-quality sample filtering. This work demonstrates that RL can effectively control sampling trajectories in discrete diffusion-style models, enabling practical, real-time capable masked image generation.

Abstract

Masked Generative Models (MGM)s demonstrate strong capabilities in generating high-fidelity images. However, they need many sampling steps to create high-quality generations, resulting in slow inference speed. In this work, we propose Speed-RL, a novel paradigm for accelerating a pretrained MGMs to generate high-quality images in fewer steps. Unlike conventional distillation methods which formulate the acceleration problem as a distribution matching problem, where a few-step student model is trained to match the distribution generated by a many-step teacher model, we consider this problem as a reinforcement learning problem. Since the goal of acceleration is to generate high quality images in fewer steps, we can combine a quality reward with a speed reward and finetune the base model using reinforcement learning with the combined reward as the optimization target. Through extensive experiments, we show that the proposed method was able to accelerate the base model by a factor of 3x while maintaining comparable image quality.

Figures (9)

• Figure 1: Speed-RL is a framework for masked generative models to generate high quality images while using significantly fewer steps. Specifically, Speed-RL uses reinforcement learning to optimize a speed and quality reward.
• Figure 2: Side-by-side qualitative comparison of Speed-RL and Meissonic (baseline). Speed-RL produces higher quality images while using significantly fewer steps.
• Figure 3: Ablation study on our speed reward and KL regularizer. We see a divergence in HPSv2.1 scores after 300 RL steps, where Speed-RL increases faster compared to having no speed reward and no low-variance KL estimator.
• Figure 4: Ablation study on the effect of enabling per-token-likelihood. There is a gap between using token-level-likelihood and sequence-level-likelihood where using token-level-likelihood consistently has a HPSv2.1 score higher than using sequence-level-likelihood.
• Figure 5: Ablation study on varying the reward standard deviation percentile threshold.