MAR-GRPO: Stabilized GRPO for AR-diffusion Hybrid Image Generation

2026-04-08

AI Summary

Existing reinforcement learning approaches in autoregressive–diffusion hybrid image generation models often suffer from training instability and premature performance saturation due to interleaved inference and noisy log-probability estimation. This work proposes a stabilized reinforcement learning framework that introduces Multi-Trajectory Expectation (MTE) to reduce gradient noise and incorporates an uncertainty-aware top-k% token optimization mechanism together with a consistency-aware token selection strategy. The proposed method substantially enhances training stability, improves visual quality of generated images, and strengthens spatial structure comprehension. Empirical evaluations demonstrate consistent outperformance over baseline GRPO and pretrained reinforcement learning models across multiple benchmarks.

Abstract

Reinforcement learning (RL) has been successfully applied to autoregressive (AR) and diffusion models. However, extending RL to hybrid AR-diffusion frameworks remains challenging due to interleaved inference and noisy log-probability estimation. In this work, we study masked autoregressive models (MAR) and show that the diffusion head plays a critical role in training dynamics, often introducing noisy gradients that lead to instability and early performance saturation. To address this issue, we propose a stabilized RL framework for MAR. We introduce multi-trajectory expectation (MTE), which estimates the optimization direction by averaging over multiple diffusion trajectories, thereby reducing diffusion-induced gradient noise. To avoid over-smoothing, we further estimate token-wise uncertainty from multiple trajectories and apply multi-trajectory optimization only to the top-k% uncertain tokens. In addition, we introduce a consistency-aware token selection strategy that filters out AR tokens that are less aligned with the final generated content. Extensive experiments across multiple benchmarks demonstrate that our method consistently improves visual quality, training stability, and spatial structure understanding over baseline GRPO and pre-RL models. Code is available at: https://github.com/AMAP-ML/mar-grpo.

Problem

• hybrid AR-diffusion
• reinforcement learning
• gradient noise
• training instability
• noisy log-probability estimation

Innovation

• multi-trajectory expectation
• token-wise uncertainty
• consistency-aware token selection
• stabilized RL
• AR-diffusion hybrid