Work

Efficient and Stable Reinforcement Learning for Diffusion Language Models

Year: 2026

Type: preprint

Abstract: Reinforcement Learning (RL) is crucial for unlocking the complex reasoning capabilities of Diffusion-based Large Language Models (dLLMs). However, applying RL to dLLMs faces unique challenges in efficiency and stability. To address these challenges, we propose Spatio-Temporal Pruning (STP), a framework designed to simultaneously improve the efficiency and stability of RL for dLLMs. STP compresses the redundancy in the generative process through: (1) \textit{spatial pruning}, which constrains the exploration space using static priors; and (2) \textit{temporal pruning}, which bypasses redundant late-stage refinement steps. Our theoretical analysis demonstrates that STP strictly reduces the variance of the log-likelihood estimation, thereby ensuring more stable policy updates. Extensive experiments demonstrate that STP surpasses state-of-the-art baselines in both efficiency and accuracy. Our code is available at https://github.com/Lolo1222/STP.

Source: arXiv (Cornell University)

DOI: 10.48550/arxiv.2602.08905

Authors: Liu, Jiawei, Wang, Xiting, Zhong, Yuanyuan, Lian, Defu, Yang, Yu

Topic: Reinforcement Learning in Robotics

Subfield: Artificial Intelligence

Field: Computer Science

Domain: Physical Sciences

Open Access status: green

OpenAlex ID: W7128470381

Cites:

Cited by:

Related to: