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

AI summary

Overview

• Research area: Natural Language Processing / Large Language Model training, specifically reinforcement learning for diversity in LLM outputs.
• Technical level: Advanced (requires familiarity with reinforcement learning, policy optimization, and LLM fine-tuning concepts; the summary below is written to be accessible).
• Scope: The paper introduces Group-Aware Policy Optimization (GAPO), a reinforcement learning method that encourages LLMs to produce more varied outputs by rewarding diversity at the group level rather than per sample.

What This Paper Is About

Large Language Models often repeat the same answers to open-ended prompts even when many valid alternatives exist — a problem known as "mode collapse." For example, asking a model for a joke frequently produces the same "Why don't scientists trust atoms?" punchline. This paper introduces GAPO, a training method that modifies the reinforcement learning objective so that models learn to spread their responses more evenly across valid options, without sacrificing accuracy on reasoning and coding tasks.

Key Contributions

1. GAPO (Group-Aware Policy Optimization): A simple extension of the popular GRPO reinforcement learning algorithm that computes rewards over an entire group of model completions rather than per individual sample, enabling learning from distributional properties like diversity and coverage.
2. Frequency-aware reward function: A reward design that penalizes over-represented answers and boosts under-represented ones, directly countering mode collapse without changing the model architecture or decoding strategy.
3. Demonstrated near-uniform sampling: GAPO-trained models achieve near-uniform distributions when selecting items from lists, and generate substantially more diverse responses to open-ended prompts — including categories unseen during training.
4. Diversity without sacrificing coherence: GAPO improves creativity in writing tasks (poetry, dialogue, songs) while maintaining comparable performance on standard benchmarks such as GSM8K, MATH, HumanEval, and MMLU-Pro.

Main Findings

• Baseline models show severe selection bias: ChatGPT-4o, Claude Sonnet 3.5, and Gemini 2.5 overwhelmingly favor specific list items (e.g., "Japan" selected 75–87% of the time regardless of list order), demonstrating item-specific, positional, and contextual biases.
• GAPO achieves near-uniform list selection: GAPO-trained Qwen2.5 models show Jensen-Shannon divergence below 0.1 from a uniform distribution, versus greater than 0.3 for all baseline models on unseen list categories.
• Open-ended diversity improves dramatically: On prompts like "Name a city," GAPO-finetuned Qwen2.5 32B produced an average of 147 unique responses across 500 samples, compared to just 24 for the base model.
• Creative writing becomes far more varied: GAPO nearly doubled semantic diversity (average embedding distance rose from 0.20 to 0.41) and increased lexical uniqueness (1-Self-BLEU from 0.54 to 0.82) across eight creative tasks.
• Reasoning accuracy is preserved: On GSM8K, MATH, HumanEval, and MMLU-Pro, GAPO performs comparably to the baseline — sometimes slightly better (e.g., flexible MATH scoring improved from 0.865 to 0.905).
• Creativity-coherence tradeoff is favorable: Across sampling temperatures, GAPO achieves higher creativity at every coherence level compared to the baseline.

Methodology in Plain English

The researchers started with GRPO, an existing reinforcement learning method that generates a group of completions for each prompt and scores each one individually. Their key insight was to flip this: instead of scoring each answer on its own, score the whole group together based on how diverse and well-balanced it is.

They designed a frequency-aware reward function that works like a fairness mechanism. Within a group of, say, 32 generated answers, if a particular valid item appears too often, it gets a lower reward; if an item is under-represented, it gets a higher reward. Invalid answers get penalized outright. This pushes the model to spread probability mass evenly across all valid options.

The team then fine-tuned Qwen2.5 Instruct models (7B and 32B parameters) using LoRA — a lightweight fine-tuning technique — on synthetic datasets of random lists spanning topics like animals, emotions, numbers, and foods. They tested on both list-selection tasks (where a fixed set of valid answers exists) and open-ended prompts (where the space of valid answers is much larger), along with creativity benchmarks and standard reasoning/coding/ knowledge benchmarks.

Why This Matters

Research impact: The paper offers a generalizable extension to GRPO that reframes reward functions as ways to shape output distributions, not just individual responses. This opens a new axis for reinforcement learning research beyond accuracy-based rewards — potentially useful for alignment, safety, and creative applications alike.

Real-world applications:

• Creative writing assistants: More varied brainstorming, storytelling, and poetry generation without repetitive phrasing.
• Data augmentation for AI training: Generating diverse synthetic datasets that cover more of the response space.
• Chatbot and dialog systems: Reducing the "bot feel" of repetitive responses and making conversations feel more natural.
• Content moderation and sampling: Producing balanced coverage across categories in tasks like moderation summaries or classification explanations.

Industry relevance: Any company deploying LLMs at scale for customer-facing or creative applications suffers from output repetition, which degrades user experience. GAPO offers an architecture-agnostic intervention that costs only a LoRA fine-tuning pass — a comparatively cheap fix for a widely felt problem. Companies with strong RLHF pipelines (OpenAI, Anthropic, Google, Meta, Amazon) could integrate these ideas into existing training loops.

Future Directions

1. Earlier pipeline integration: Currently GAPO is applied as a post-hoc LoRA fine-tune. Applying it during or before initial instruction tuning could yield deeper, more robust diversity.
2. Balancing diversity with accuracy: The current reward assumes all completions are equally valid. Extending it to handle accuracy-diversity tradeoffs would broaden applicability to reasoning tasks.
3. Unbounded response spaces: Generalizing beyond predefined valid sets (lists, known categories) to fully open-ended tasks where "valid" is implicit or unbounded.
4. Understanding generalization limits: The paper shows GAPO generalizes to unseen categories, but the boundaries of that generalization — and whether it might reduce safety guardrails — remain unclear.

Target Audience

The primary audience is machine learning researchers and engineers working on LLM training, particularly those familiar with reinforcement learning from human feedback, GRPO, and policy optimization techniques. Secondary audiences include NLP practitioners frustrated by output repetition in deployed systems, safety and alignment researchers interested in how training objectives shape model behavior, and creative-AI researchers seeking methods for more diverse generation. Beginners may find the reinforcement learning math challenging, but the high-level motivation and case studies are accessible to anyone familiar with LLMs.

Authors’ 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.