# DISCO Balances the Scales: Adaptive Domain- and Difficulty-Aware Reinforcement Learning on Imbalanced Data

arXiv: https://arxiv.org/abs/2505.15074 · HTML: https://mlsanity.com/paper/2505.15074/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in Reinforcement learning for reasoning, one of 17 directions in LLM reasoning.
- Builds on 40 of its 57 references that are in the map.
- First posted on arXiv 21 May 2025.

## Placed in

- LLM reasoning / [Reinforcement learning for reasoning](https://mlsanity.com/lr/reinforcement-learning-for-reasoning/)
- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)

## Builds on (most cited first, 30 of 40)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena](https://mlsanity.com/paper/2306.05685/) (2023)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge](https://mlsanity.com/paper/1803.05457/) (2018)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [Understanding R1-Zero-Like Training: A Critical Perspective](https://mlsanity.com/paper/2503.20783/) (2025)
- [LIMA: Less Is More for Alignment](https://mlsanity.com/paper/2305.11206/) (2023)
- [MetaMath: Bootstrap Your Own Mathematical Questions for Large Language Models](https://mlsanity.com/paper/2309.12284/) (2023)
- [Large-Scale Long-Tailed Recognition in an Open World](https://mlsanity.com/paper/1904.05160/) (2019)
- [Latent Retrieval for Weakly Supervised Open Domain Question Answering](https://mlsanity.com/paper/1906.00300/) (2019)
- [Dice Loss for Data-imbalanced NLP Tasks](https://mlsanity.com/paper/1911.02855/) (2019)
- [Crossing the Reward Bridge: Expanding RL with Verifiable Rewards Across Diverse Domains](https://mlsanity.com/paper/2503.23829/) (2025)
- [SoTA with Less: MCTS-Guided Sample Selection for Data-Efficient Visual Reasoning Self-Improvement](https://mlsanity.com/paper/2504.07934/) (2025)
- [Branch-Train-MiX: Mixing Expert LLMs into a Mixture-of-Experts LLM](https://mlsanity.com/paper/2403.07816/) (2024)
- [Enhancing Visual-Language Modality Alignment in Large Vision Language Models via Self-Improvement](https://mlsanity.com/paper/2405.15973/) (2024)
- [SRPO: A Cross-Domain Implementation of Large-Scale Reinforcement Learning on LLM](https://mlsanity.com/paper/2504.14286/) (2025)
- [Explore Spurious Correlations at the Concept Level in Language Models for Text Classification](https://mlsanity.com/paper/2311.08648/) (2023)
- [Beyond Reward Hacking: Causal Rewards for Large Language Model Alignment](https://mlsanity.com/paper/2501.09620/) (2025)
- [Reward-Robust RLHF in LLMs](https://mlsanity.com/paper/2409.15360/) (2024)
- [GVPO: Group Variance Policy Optimization for Large Language Model Post-Training](https://mlsanity.com/paper/2504.19599/) (2025)
- [Towards Self-Improvement of LLMs via MCTS: Leveraging Stepwise Knowledge with Curriculum Preference Learning](https://mlsanity.com/paper/2410.06508/) (2024)
- [Reflection-Tuning: Data Recycling Improves LLM Instruction-Tuning](https://mlsanity.com/paper/2310.11716/) (2023)
- [Is Model Ensemble Necessary? Model-based RL via a Single Model with Lipschitz Regularized Value Function](https://mlsanity.com/paper/2302.01244/) (2023)

## Built on by (most cited first, 7 of 7)

- [Alignment and Safety in Large Language Models: Safety Mechanisms, Training Paradigms, and Emerging Challenges](https://mlsanity.com/paper/2507.19672/) (2025)
- [SAGE-OPD: Selective Agent-Guided Intervention for Multi-Turn On-Policy Distillation](https://mlsanity.com/paper/2606.19659/) (2026)
- [Semantically-Aware Rewards for Open-Ended R1 Training in Free-Form Generation](https://mlsanity.com/paper/2506.15068/) (2025)
- [Synthetic Sandbox for Training Machine Learning Engineering Agents](https://mlsanity.com/paper/2604.04872/) (2026)
- [G2RPO-A: Guided Group Relative Policy Optimization with Adaptive Guidance](https://mlsanity.com/paper/2508.13023/) (2025)
- [EBPO: Empirical Bayes Shrinkage for Stabilizing Group-Relative Policy Optimization](https://mlsanity.com/paper/2602.05165/) (2026)
- [Injecting Distributional Awareness into MLLMs via Reinforcement Learning for Deep Imbalanced Regression](https://mlsanity.com/paper/2605.01402/) (2026)

## Abstract (arXiv)

Large Language Models (LLMs) are increasingly aligned with human preferences through Reinforcement Learning from Human Feedback (RLHF). Among RLHF methods, Group Relative Policy Optimization (GRPO) has gained attention for its simplicity and strong performance, notably eliminating the need for a learned value function. However, GRPO implicitly assumes a balanced domain distribution and uniform semantic alignment across groups, assumptions that rarely hold in real-world datasets. When applied to multi-domain, imbalanced data, GRPO disproportionately optimizes for dominant domains, neglecting underrepresented ones and resulting in poor generalization and fairness. We propose Domain-Informed Self-Consistency Policy Optimization (DISCO), a principled extension to GRPO that addresses inter-group imbalance with two key innovations. Domain-aware reward scaling counteracts frequency bias by reweighting optimization based on domain prevalence. Difficulty-aware reward scaling leverages prompt-level self-consistency to identify and prioritize uncertain prompts that offer greater learning value. Together, these strategies promote more equitable and effective policy learning across domains. Extensive experiments across multiple LLMs and skewed training distributions show that DISCO improves generalization, outperforms existing GRPO variants by 5% on Qwen3 models, and sets new state-of-the-art results on multi-domain alignment benchmarks. Our code and data are available at https://github.com/Tonyzhou98/disco_grpo.