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

TL;DR

This paper introduces GAPO, a reinforcement learning method that enhances output diversity in large language models by optimizing group-level rewards, leading to more varied and valid responses without sacrificing accuracy.

Contribution

GAPO extends existing policy optimization techniques to incorporate group-level rewards, improving diversity in LLM outputs while maintaining performance on benchmarks.

Findings

GAPO increases response diversity in LLMs.

GAPO maintains accuracy on standard benchmarks.

GAPO generalizes to open-ended prompts.

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.