Efficient reinforcement learning improves language model reasoning accuracy
ForkGRPO: Efficient Exploration and Fine-Grained Credit Assignment for Diffusion Language Models
Computation and Language
Summary
Language models that generate text can sometimes struggle to explore different ways of answering questions, which hurts how well they reason. The paper introduces ForkGRPO, a new method that helps the model try multiple possible answer paths from key points and better figures out which parts are most important to improve. This approach uses fewer model runs and boosts performance on several tests compared to earlier methods. It can also train much faster without losing accuracy, making smarter text generation more efficient.
What this means in practice
- •For natural language processing teams: Improve language model training to better explore reasoning paths and assign credit to effective answer segments.
- •For software developers building ai assistants: Reduce training time and costs while enhancing reasoning accuracy of language models that generate complex text.
Authors
Zhaoyan Jia, Chengwei Zhou, Gourav Datta, Massoud Pedram
Abstract
Diffusion large language models (dLLMs) have advanced rapidly, yet strengthening their reasoning through reinforcement learning (RL) remains an open challenge. Any-order generation can suffer from entropy degradation, which narrows the reasoning paths available for exploration. Recent work mitigates this issue by applying standard GRPO with left-to-right generation. However, independently generating every rollout from the input prompt can leave alternative continuations from intermediate states underexplored, while trajectory-level advantages assign uniform credit to all tokens without distinguishing the contributions of individual reasoning segments. We propose ForkGRPO, an RL framework that couples efficient exploration with fine-grained segment-level credit assignment. It samples multiple continuations from high-entropy intermediate states along a reference trajectory, prioritizing textual boundaries and reusing shared prefixes. The continuations provide complementary learning signals: within-state group-relative advantages guide which continuations to reinforce, while Monte Carlo state-value differences indicate which reasoning segments improve expected outcomes. Among the evaluated architecture-preserving RL methods, ForkGRPO achieves the best performance on three of four benchmarks, while improving over the directly comparable JustGRPO by up to 1.2 percentage points and using up to 27.8% fewer generated tokens. ForkGRPO-Fast further achieves a $2.49\times$ training wall-clock speedup with negligible accuracy loss.