EGSPO: A New Paradigm for Infusing Reinforcement Learning into Diffusion Language Models

The Texas A&M University team proposes the EGSPO-SA (Entropy-Guided Stepwise Policy Optimization with Stepwise Advantages) framework, which addresses core challenges in RL fine-tuning of diffusion language models through entropy-guided step selection and lightweight advantage estimators. This framework has achieved significant breakthroughs in core benchmark tests such as code generation, logical reasoning, and mathematical reasoning, and has open-sourced the implementation code and model checkpoints.

Diffusion language models (dLLMs) generate sequences through iterative denoising, which is significantly different from the generation method of autoregressive models (such as the GPT series). Traditional sequence-level RL methods assume that the complete output is generated at once, making them difficult to directly apply to the multi-step denoising process of dLLMs, facing three major challenges:

1. State Space Explosion: Denoising trajectories form high-dimensional state sequences, leading to the curse of dimensionality for traditional RL methods;
2. Credit Assignment Difficulty: The quality of the final output depends on the collaboration of all steps, making it hard to determine the contribution of a single step;
3. High Computational Cost: Training a separate value model for each step is infeasible. These issues restrict the performance improvement of dLLMs.

The EGSPO-SA framework addresses the pain points of RL fine-tuning for diffusion models and proposes three major innovations:

1. Diffusion MDP Formalization: Convert the denoising process into a Finite-Horizon Markov Decision Process (Finite-Horizon MDP), derive a policy gradient objective that can be decomposed across steps, and focus on key steps;
2. Entropy-Guided Step Selection: Identify high-information steps (decision points with high model uncertainty) based on entropy, concentrating computational resources and learning signals;
3. Lightweight Step-Level Advantage Estimator: Calculate single-step advantage values without the need for an additional value model, significantly reducing training costs.

EGSPO-SA has been validated for effectiveness in multiple challenging tasks:

• Code Generation: Generate syntactically correct and fully functional code snippets;
• Logical Reasoning: Excel at constructing and verifying complex logical chains;
• Mathematical Reasoning: Demonstrate step-by-step reasoning and precise calculation capabilities on benchmarks such as GSM8K. The team has open-sourced the model checkpoint (fatemehdoudi97/egspo-llada-8b) and detailed usage instructions on HuggingFace.

The project code has a clear structure and supports multi-node distributed training:

• Core training logic: egspo/train.sh;
• Evaluation process: First generate completions via eval/eval_checkpoints.sh, then calculate metrics using eval/get_and_save_metrics.py;
• Environment configuration: Provide environment.yml to manage dependencies, and the README explains key variables (such as WANDB_API_KEY, HF_HOME);
• Based on open-source libraries: Implemented based on the dllm-reasoning/d1 codebase, reflecting the tradition of academic collaboration.

EGSPO-SA marks an important progress in the field of RL fine-tuning for diffusion language models. Its technical ideas (entropy-guided step selection, lightweight advantage estimation) may inspire research in iterative process fields such as multi-modal generation and video generation. For practitioners, this framework provides a ready-to-use RL fine-tuning tool and is expected to become one of the standard tools for RL fine-tuning of diffusion LLMs.