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

Abstract: 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.

## introduction aligning large language models (llms) with human preferences is a central challenge in modern *equal contribution. ai systems @xcite @xcite . reinforcement learning from human feedback (rlhf) has become the dominant approach for fine-tuning llms toward desirable behavior, enabling alignment with nuanced human intent @xcite @xcite @xcite @xcite . within this framework, group relative policy optimization (grpo) @xcite offers a promising alternative to value-based methods, simplifying training while achieving strong performance. despite its advantages, grpo faces a significant challenge when applied to multi-domain datasets, a common scenario in real-world applications spanning math, question answering, and commonsense reasoning. while grpo effectively removes the need for a value network and mitigates withingroup variance, it implicitly assumes that prompt groups are sampled uniformly and that reward signals are semantically aligned across domains. however, this assumption often fails in practice. realworld datasets are typically imbalanced, with a few dominant domains and many underrepresented ones @xcite . optimization gradients become skewed toward high-frequency domains, starving rare domains of learning signal. this results in models that generalize poorly to critical low-resource domains and amplify existing data biases. data augmentation offers one workaround but introduces substantial overhead in generating highquality synthetic prompts @xcite . to address the inter-group imbalance in grpo, we propose domain-informed self-consistency policy optimization (disco), an enhanced framework designed to promote equitable learning across imbalanced multi-domain datasets. disco introduces two complementary strategies to improve generalization under distributional skew: domainaware and difficulty-aware reward scaling. domain-aware scaling reweights prompt groups inversely by their frequency, reducing overoptimization on dominant domains while amplifying learning signals from underrepresented ones. difficulty-aware scaling leverages prompt-level self-consistency, an intrinsic signal in grpo, to identify and upweight prompts where the policy exhibits high uncertainty or inconsistent responses. since not all prompts are equally challenging, treating them uniformly can lead the policy to overfit on easy examples while neglecting harder, more informative ones. by prioritizing uncertain prompts, this strategy guides the model to focus its learning on cases that offer a greater signal for improvement. by integrating these two forms of adaptive scaling, disco enables all domains, regardless of their prevalence, to meaningfully contribute to policy optimization. as a result, it effectively mitigates grpo's inter-group imbalance and achieves state-of-the-art performance across diverse llm architectures and training distributions. our contributions are summarized as follows: ## background and motivation this section first provides the necessary background on the grpo algorithm. we then present an empirical study to concretely demonstrate how its optimization behavior is affected by domain imbalance in training data, thereby establishing the specific motivation for the enhancements proposed in subsequent sections. ## domain-and difficulty-aware scaling building on the observation that grpo lacks mechanisms for addressing domain-level imbalance, our disco, as depicted in figure 1 , proposes two modifications in multi-domain settings: domain-aware and difficulty-aware reward scaling. first, we introduce domain-aware scaling, which reweights prompt groups based on domain frequency. this strengthens the learning signal for underrepresented ones. secondly, our difficulty-aware scaling enhances learning efficiency by focusing on challenging examples. it uses prompt-level self-consistency, the average reward from g candidate outputs per prompt, as a signal of policy uncertainty within grpo. a lower self-consistency indicates higher uncertainty and thus greater difficulty. group rewards are then scaled by a weight inversely proportional to this self-consistency score, thereby prioritizing more ambiguous prompts. together, these two mechanisms allow grpo to adaptively rescale group-level advantages based on both domain rarity and prompt difficulty, offering a more globally informed optimization process while preserving its core group-relative framework. ## experiment setup we now evaluate the effectiveness of our disco. our experiments aim to assess whether these methods improve alignment performance on underrepresented or difficult domains, and whether the improvements are consistent across different models. dataset setup. we evaluate across four task domains: imdb (text classification), gsm8k and math (math problem solving), natural ques-tions (open-domain qa), and arc (reasoning qa). our training dataset consists of 4,000 examples, as detailed in baseline methods. we compare our method against the following baselines: (1) base model, the pretrained model without any fine-tuning; (2) naive grpo, which applies the original grouprelative optimization without any reward reweighting; and (3) dr. grpo @xcite , which removes the length and standard deviation normalization. while other grpo variants address issues like length bias or training instability @xcite , they are not included as baselines as our work targets grpo's vulnerability to domain imbalance. in addition, we conduct ablation studies to isolate the contributions of domainaware and difficulty-aware components. proposed methods. we evaluate three variants of disco, each utilizing one of the domain-aware weighting strategies in section 3.1. disco-log uses log-scaled domain weights (v1), disco-logsq uses squared-log weights (v2), and disco-inv uses inverse-frequency weights (v3). all variants incorporate the same difficultyaware component based on self-consistency. this setup allows us to assess how different strengths of domain correction interact with difficulty-aware scaling in multi-domain alignment. model setup. we evaluate our method across a diverse set of language models. these include qwen2.5-0.5b, qwen2.5-1.5b, qwen2-0.5b, qwen3-0.6b, qwen2.5-7b, and qwen1.5-moe-a2.7b (14b total parameters, 2.7b activated) @xcite , as well as llama3.2-1b @xcite , olmo2-1b (olmo et al., 2024), and gemma2-2b-it @xcite . this selection spans both dense and mixture-of-experts (moe) architectures and covers a range of model capacities from 0.5b to 14b parameters. for the grpo group size, we use @xmath0, 4, 8, 16 depending on the experiment. additional implementation details are provided in appendix a. ## results we begin our analysis by identifying the most effective variant of our proposed scaling strategy. after selecting a default method, we evaluate its impact through comparisons with baseline methods and targeted ablation studies. ## related work reinforcement learning from human feedback (rlhf) has emerged as a powerful paradigm for aligning llms with human preferences and values. the most popular approach is proximal policy optimization (ppo) @xcite . building upon ppo, grpo eliminates the need for value functions through estimation from group scores and has emerged as a popular training diagram @xcite . rule-based grpo further streamlines training by using predefined rules to evaluate outputs, avoiding the need for a separate reward model @xcite . vari-2 4 8 16 group size 20 25 30 35 avg. accuracy math-heavy 2 4 8 16 group size 20 25 30 35 avg. accuracy imdb-heavy 2 4 8 16 group size 20 25 30 35 avg. accuracy nq-heavy 2 4 8 16 group size 20 25 30 35 avg. accuracy ## conclusion in this work, we investigate the vulnerability of grpo to domain imbalance in multi-domain alignment and introduce disco that combines domain and difficulty-aware reward scaling to enable grpo to better handle skewed training distributions. through experiments across diverse model scales, group sizes, and alignment settings, we demonstrate that disco consistently improves generalization and outperforms existing baselines. ablation studies further confirm that both scaling components are essential. overall, our method provides a robust enhancement to grpo for aligning llms with real-world, imbalanced data.