Diversity-Aware Policy Optimization for Large Language Model Reasoning
Jian Yao, Ran Cheng, Xingyu Wu, Jibin Wu, Kay Chen Tan
TL;DR
This work systematically examines how diversity in reinforcement learning shapes large-language-model reasoning. By introducing Potential@k to quantify training potential and a token-level diversity objective applied to positive samples, the authors demonstrate that promoting diversity correlates with improved reasoning performance in high-capacity models. They integrate this diversity-aware objective into the R1-zero framework (R1-zero-Div), yielding a 3.5% average improvement across four mathematical benchmarks and more diverse, robust solutions. The approach balances exploration and solution quality, analyzes gradient behavior, and validates generalization to smaller models, offering a practical method for enhancing reasoning via diversity-aware policy optimization. Limitations include computational constraints and the need to validate on larger LLMs, with future work pointing toward semantic, user-driven diversity and broader task applications.
Abstract
The reasoning capabilities of large language models (LLMs) have advanced rapidly, particularly following the release of DeepSeek R1, which has inspired a surge of research into data quality and reinforcement learning (RL) algorithms. Despite the pivotal role diversity plays in RL, its influence on LLM reasoning remains largely underexplored. To bridge this gap, this work presents a systematic investigation into the impact of diversity in RL-based training for LLM reasoning, and proposes a novel diversity-aware policy optimization method. Across evaluations on 12 LLMs, we observe a strong positive correlation between the solution diversity and Potential at k (a novel metric quantifying an LLM's reasoning potential) in high-performing models. This finding motivates our method to explicitly promote diversity during RL training. Specifically, we design a token-level diversity and reformulate it into a practical objective, then we selectively apply it to positive samples. Integrated into the R1-zero training framework, our method achieves a 3.5 percent average improvement across four mathematical reasoning benchmarks, while generating more diverse and robust solutions.