TLDR: A new reinforcement learning algorithm, Stratified GRPO, has been developed to improve how Large Language Model (LLM) search agents learn. It tackles ‘cross-stratum bias,’ a problem arising from the diverse ways LLM agents use search tools, which can distort learning. Stratified GRPO’s core component, Stratified Advantage Normalization (SAN), partitions agent experiences into homogeneous groups and calculates advantages locally, ensuring fair comparisons. This method significantly outperforms existing techniques on question-answering benchmarks, leading to higher rewards, greater training stability, and more effective multi-step search policies.
Large Language Models (LLMs) are becoming increasingly sophisticated, often acting as agents that use external tools like search engines to tackle complex, multi-step problems. Training these agents effectively often involves reinforcement learning (RL), a powerful method where models learn from trial and error based on rewards. However, a significant challenge arises from the diverse nature of how these LLM search agents operate.
Imagine an LLM agent trying to answer a question. Sometimes it might find the answer with a single search query, while other times it might need multiple searches, or even no searches at all if the information is already within its knowledge. These different approaches lead to what researchers call ‘structurally heterogeneous trajectories.’ In simpler terms, the paths an agent takes to solve a problem can look very different, leading to varied outcomes and reward distributions. Standard RL methods, which typically use a single, global benchmark to evaluate all these diverse paths, struggle with this heterogeneity. This leads to what the researchers identify as ‘cross-stratum bias’ – essentially, comparing apples to oranges. This bias can distort how the agent learns, making it harder to assign credit correctly for successful actions and hindering its ability to explore more complex search strategies.
Introducing Stratified GRPO: A Principled Solution
To address this fundamental issue, a new approach called Stratified GRPO has been proposed. Its core innovation is Stratified Advantage Normalization (SAN). Instead of a single global benchmark, SAN divides the agent’s experiences (trajectories) into homogeneous groups, or ‘strata,’ based on their structural properties – for example, how many search calls were made. Within each of these homogeneous groups, advantages (how much better or worse an action was compared to the average) are calculated locally. This ensures that trajectories are only compared against their true peers, eliminating the ‘apples-to-oranges’ problem.
The theoretical analysis of SAN shows several key benefits. It effectively removes cross-stratum bias, ensuring fair credit assignment. It also provides learning signals that are conditionally unbiased and have consistent variance within each stratum, leading to a purer and more stable learning process. Importantly, while achieving these superior local properties, SAN maintains the global unbiasedness and unit-variance properties of standard normalization methods.
For practical stability, especially when some groups have very few examples, Stratified GRPO further incorporates a ‘Blended Advantage.’ This linearly combines SAN with a global estimator, balancing SAN’s local purity with the global signal’s stabilizing effect on smaller strata.
Also Read:
- Unlocking Advanced Reasoning in Diffusion Language Models Through Amortized GRPO
- DeSA: A Two-Stage Approach to Enhance LLM Agent Search and Answering Capabilities
Experimental Validation and Impact
The effectiveness of Stratified GRPO was rigorously tested across seven diverse question-answering benchmarks, including both single-hop (requiring one piece of information) and multi-hop (requiring multiple pieces of information) tasks. The results were compelling: Stratified GRPO consistently and substantially outperformed existing methods, including the standard GRPO, by up to 11.3 points in average performance. The improvements were particularly significant on multi-hop benchmarks, where the method achieved an average performance gain of up to 14.5 points over the strongest baseline.
Beyond just higher scores, Stratified GRPO demonstrated greater training stability, avoiding issues like ‘training collapse’ seen in some baseline models. Crucially, it also learned more effective search policies. While standard GRPO often stagnated at around one search call per question, Stratified GRPO successfully learned to perform iterative searches, converging to approximately 2.5 search calls. This ability to effectively explore and utilize multi-step search strategies directly translated to its superior performance on complex tasks.
These findings underscore that addressing structural heterogeneity is a critical factor in developing robust and effective reinforcement learning algorithms for LLM search agents. Stratified GRPO offers a principled solution to this challenge, paving the way for more capable and stable LLM agents. You can read the full research paper here.
