Consolidating Reinforcement Learning for Multimodal Discrete Diffusion Models

2025-10-03

AI Summary

This work addresses the challenge of integrating discrete diffusion models (DDMs) into reinforcement learning (RL) under the non-autoregressive paradigm. We propose MaskGRPO, the first RL framework with a scalable theoretical foundation for DDMs. It introduces a modality-aware importance sampling mechanism, a vision-specific rollout strategy, and token-level volatility modeling to improve gradient estimation quality. Unlike existing RL methods incompatible with DDM architectures (e.g., GRPO), MaskGRPO enables joint policy optimization over both textual and visual sequences. Experiments demonstrate significant performance gains on mathematical reasoning, code generation, and multimodal visual synthesis tasks, alongside enhanced training stability and superior output quality.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Optimizing discrete diffusion model (DDM) with rewards remains a challenge: the non-autoregressive paradigm makes importance sampling intractable and rollout complex, puzzling reinforcement learning methods such as Group Relative Policy Optimization (GRPO). In this study, we introduce MaskGRPO, the first viable approach to enable scalable multimodal reinforcement learning in discrete diffusion with effective importance sampling and modality-specific adaptations. To this end, we first clarify the theoretical foundation for DDMs, which facilitates building an importance estimator that captures valuable token fluctuation for gradient updates. We then delicately tailored the rollout method for visual sequences, which yields diverse completions and reliable optimization gradients. Upon math reasoning, coding, and visual generation benchmarks, MaskGRPO brings more stable and efficient updates, leading to stronger reasoning performance and better generation quality. This study establishes MaskGRPO as a systematic policy optimization approach and the first practical way for discretized visual diffusion.

Problem

• Optimizing discrete diffusion models with reinforcement learning rewards
• Enabling scalable multimodal reinforcement learning for discrete diffusion
• Improving reasoning performance and generation quality across benchmarks

Innovation

• MaskGRPO enables scalable multimodal reinforcement learning
• Importance estimator captures token fluctuation for gradients
• Tailored rollout method yields diverse visual completions