Group-Aware Reinforcement Learning for Output Diversity in Large Language Models

TL;DR

GAPO addresses mode collapse in large language models by extending GRPO to compute rewards over groups of completions, enabling learning of distributional properties such as diversity and uniform coverage. A frequency-aware reward drives near-uniform sampling from predefined valid outputs, demonstrated on list-selection tasks, open-ended prompts, and creative writing, while preserving accuracy on GSM8K, MATH, HumanEval, and MMLU-Pro. The approach maintains the existing training pipeline, leveraging LoRA finetuning and synthetic data, and shows improved diversity across creativity metrics with controlled coherence. The work discusses limitations, potential risks, and future directions, including earlier integration and balancing diversity with task-specific accuracy.

Abstract

Large Language Models (LLMs) often suffer from mode collapse, repeatedly generating the same few completions even when many valid answers exist, limiting their diversity across a wide range of tasks. We introduce Group-Aware Policy Optimization (GAPO), a simple extension of the recent and popular Group Relative Policy Optimization (GRPO) that computes rewards over the group as a whole. GAPO enables learning from the group-level properties such as diversity and coverage. We demonstrate GAPO using a frequency-aware reward function that encourages uniform sampling over valid LLM completions, and show that GAPO-trained models produce valid and more diverse model responses. Beyond this setup, GAPO generalizes to open-ended prompts and improves response diversity without compromising accuracy on standard LLM benchmarks (GSM8K, MATH, HumanEval, MMLU-Pro). Our code will be made publicly available.

Figures (10)

• Figure 1: GAPO with a frequency-aware reward promotes output uniformity and diversity. GAPO mitigates over-representation of frequent completions and encourages more balanced generation, both in list selection tasks and open-ended prompts.
• Figure 2: Bias Analysis. (a) Using the same list of countries, we prompt each model 100 times with each of the two instruction prompts. (b) We prompt each model 100 times with the same instruction prompt, each time the list of countries is randomly shuffled. For each model the largest and second-largest probabilities are shown.
• Figure 3: GRPO vs. GAPO. Comparison between GRPO and GAPO (left) and illustration of frequency-aware rewards (right). GRPO assigns per-sample rewards, while GAPO computes rewards over the whole group, enabling distributional signals such as diversity and coverage. Our frequency-aware reward function penalizes overrepresented outputs and boosts underrepresented ones, guiding the model toward uniform sampling over equally valid responses.
• Figure 4: Sampling from a list. Output distributions for different models when prompted to choose from a list of planets (a) or musical instruments (b). NV denotes an invalid response not in the list. Panel (c) shows the average Jensen-Shannon divergence (JS) from uniform distribution across all topics.
• Figure 5: Open-set diversity. (a,b) Cumulative number of unique responses across 500 samples for open-ended prompts. (c) Average number of unique responses across 500 samples, aggregated over all ten categories.