Table of Contents
Fetching ...

DRA-GRPO: Your GRPO Needs to Know Diverse Reasoning Paths for Mathematical Reasoning

Xiwen Chen, Wenhui Zhu, Peijie Qiu, Xuanzhao Dong, Hao Wang, Haiyu Wu, Huayu Li, Aristeidis Sotiras, Yalin Wang, Abolfazl Razi

TL;DR

The paper addresses the limitation of scalar reward signals in post-training RLHF for mathematical reasoning, which promotes mode collapse and underexplores diverse reasoning strategies. It introduces Diversity-aware Reward Adjustment (DRA) for GRPO, using Submodular Mutual Information to measure semantic diversity and applying inverse propensity scoring to reweight samples. The approach is plug-and-play with GRPO variants and yields data-efficient improvements; on five math benchmarks with only 7000 training samples, it achieves an average accuracy of 58.2% on a 1.5B model. These results demonstrate that explicitly modeling reasoning diversity is critical for robust alignment of LLMs in data-constrained settings and can guide future work toward diversity-aware RL for reasoning tasks.

Abstract

Post-training LLMs with Reinforcement Learning, specifically Group Relative Policy Optimization (GRPO), has emerged as a paradigm for enhancing mathematical reasoning. However, standard GRPO relies on scalar correctness rewards that are often non-injective with respect to semantic content: distinct reasoning paths receive identical rewards. This leads to a Diversity-Quality Inconsistency, where the policy collapses into a narrow set of dominant modes while ignoring equally valid but structurally novel strategies. To bridge this gap, we propose Diversity-aware Reward Adjustment (DRA), a theoretically grounded framework that calibrates the reward signal using the semantic density of sampled groups. By leveraging Submodular Mutual Information (SMI), DRA implements an Inverse Propensity Scoring (IPS) mechanism that effectively de-biases the gradient estimation. This creates a repulsive force against redundancy, driving the policy to achieve better coverage of the high-reward landscape. Our method is plug-and-play and integrates seamlessly with GRPO variants. Empirical evaluations on five math benchmarks demonstrate that DRA-GRPO consistently outperforms strong baselines, achieving an average accuracy of 58.2% on DeepSeek-R1-Distill-Qwen-1.5B with only 7,000 training samples and $55 cost, highlighting the critical role of diversity calibration in data-efficient alignment.

DRA-GRPO: Your GRPO Needs to Know Diverse Reasoning Paths for Mathematical Reasoning

TL;DR

The paper addresses the limitation of scalar reward signals in post-training RLHF for mathematical reasoning, which promotes mode collapse and underexplores diverse reasoning strategies. It introduces Diversity-aware Reward Adjustment (DRA) for GRPO, using Submodular Mutual Information to measure semantic diversity and applying inverse propensity scoring to reweight samples. The approach is plug-and-play with GRPO variants and yields data-efficient improvements; on five math benchmarks with only 7000 training samples, it achieves an average accuracy of 58.2% on a 1.5B model. These results demonstrate that explicitly modeling reasoning diversity is critical for robust alignment of LLMs in data-constrained settings and can guide future work toward diversity-aware RL for reasoning tasks.

Abstract

Post-training LLMs with Reinforcement Learning, specifically Group Relative Policy Optimization (GRPO), has emerged as a paradigm for enhancing mathematical reasoning. However, standard GRPO relies on scalar correctness rewards that are often non-injective with respect to semantic content: distinct reasoning paths receive identical rewards. This leads to a Diversity-Quality Inconsistency, where the policy collapses into a narrow set of dominant modes while ignoring equally valid but structurally novel strategies. To bridge this gap, we propose Diversity-aware Reward Adjustment (DRA), a theoretically grounded framework that calibrates the reward signal using the semantic density of sampled groups. By leveraging Submodular Mutual Information (SMI), DRA implements an Inverse Propensity Scoring (IPS) mechanism that effectively de-biases the gradient estimation. This creates a repulsive force against redundancy, driving the policy to achieve better coverage of the high-reward landscape. Our method is plug-and-play and integrates seamlessly with GRPO variants. Empirical evaluations on five math benchmarks demonstrate that DRA-GRPO consistently outperforms strong baselines, achieving an average accuracy of 58.2% on DeepSeek-R1-Distill-Qwen-1.5B with only 7,000 training samples and $55 cost, highlighting the critical role of diversity calibration in data-efficient alignment.
Paper Structure (21 sections, 71 equations, 17 figures, 6 tables, 1 algorithm)

This paper contains 21 sections, 71 equations, 17 figures, 6 tables, 1 algorithm.

Figures (17)

  • Figure 1: Illustration of the Exploration-Exploitation trade-off in GRPO. The grey dots represent the landscape of potential high-reward reasoning paths, distributed across a common dominant mode (center) and novel but sparser modes (sides). (a) Vanilla GRPO suffers from Mode Collapse: relying solely on scalar rewards, the policy may collapse into the dominant mode, ignoring equally valid but semantically distinct strategies. (b)DRA-GRPO (Ours) achieves Diverse Exploration: by penalizing semantic redundancy, our method effectively disperses probability mass to uncover and reinforce novel reasoning paths, aligning the policy with the full spectrum of correct solutions.
  • Figure 2: Case study illustrating the Diversity-Quality Inconsistency. We present two correct completions for the same sequence problem. Completion 1 ($o_1$) adopts a concise, formula-driven strategy, whereas Completion 2 ($o_2$) exhibits an exploratory, "thinking-out-loud" reasoning style with step-by-step verification. Despite their profound semantic disparity, Vanilla GRPO assigns them nearly indistinguishable scalar rewards (2.782 vs. 2.855), failing to capture the structural diversity of the reasoning paths.
  • Figure 3: Distribution of $p$-values from Spearman’s rank correlation between completion quality and semantic diversity. The test is conducted for every prompt.
  • Figure S4: Distribution of $p$-values from Spearman’s rank correlation between completion quality and semantic diversity. Embedding model is nomic-ai/nomic-embed-text-v1.5.
  • Figure S5: Prompt used for Example Question 1.
  • ...and 12 more figures