LAD: Learning Advantage Distribution for Reasoning

Paper: 2602.20132 Authors: Wendi Li, Sharon Li Categories: cs.LG

Abstract

Current reinforcement learning approaches for large language model reasoning focus on maximizing expected rewards, which can cause overfitting to dominant reward signals and neglect valid alternative reasoning paths. This paper introduces Learning Advantage Distributions (LAD), a distribution-matching framework that learns the advantage-induced distribution rather than simply maximizing advantages. By formulating the objective as minimizing an $f$-divergence between policy-induced and advantage-induced distributions, LAD increases likelihood for high-advantage responses while preventing over-confident probability growth. The method requires no additional training cost compared to GRPO and demonstrates improved accuracy and generative diversity across math and code reasoning benchmarks.

Key Contributions

• Introduces a distribution-matching framework that learns advantage distributions instead of maximizing expected rewards
• Establishes theoretical equivalence between optimal policy updates and advantage-based target distributions
• Derives a practical objective using $f$-divergence minimization that naturally prevents mode collapse without entropy regularization
• Demonstrates zero additional training cost compared to existing methods like GRPO
• Validates the approach in controlled bandit settings showing faithful recovery of multimodal advantage distributions
• Shows consistent improvements in both accuracy and diversity across multiple LLM backbones on math and code tasks

Methodology and Theoretical Framework

The core innovation of LAD lies in reframing the reinforcement learning objective from advantage maximization to distribution matching. Traditional RL objectives for LLM reasoning optimize:

$\max_\pi \mathbb{E}_{y \sim \pi(\cdot|x)}[A(x, y)]$

where $A(x, y)$ represents the advantage function. This formulation tends to concentrate probability mass on the highest-reward trajectories, potentially ignoring valid alternatives.

LAD instead defines a target distribution proportional to the advantage:

$p^*(y|x) \propto \exp(\beta \cdot A(x, y))$

The learning objective becomes minimizing an $f$-divergence between the policy distribution $\pi(y|x)$ and this advantage-induced distribution $p^*(y|x)$. This formulation yields gradient updates that:

1. Increase likelihood for responses with positive advantages
2. Suppress over-confident probability growth through the divergence structure
3. Naturally maintain exploration of the advantage landscape

The practical implementation uses a reverse KL divergence formulation, resulting in gradient updates that scale with advantage values while incorporating a regularization term that prevents probability collapse. Critically, this regularization emerges from the divergence structure itself rather than requiring auxiliary entropy terms.

Experimental Results and Analysis

The authors validate LAD across multiple dimensions. In a controlled multi-armed bandit setting with known multimodal advantage distributions, LAD successfully recovers the true distribution shape, demonstrating that the theoretical framework translates to practice.

On practical reasoning tasks, LAD shows consistent improvements:

• Math reasoning: Tested on mathematical problem-solving benchmarks across different LLM backbones, LAD improves both pass rates and solution diversity
• Code generation: On coding tasks, the method maintains or improves functional correctness while generating more diverse solution approaches
• Scaling behavior: The improvements hold across model sizes, suggesting the approach scales naturally with model capacity

Importantly, LAD achieves these gains without additional computational cost during training compared to baseline methods like GRPO. The distribution-matching objective integrates seamlessly into standard policy gradient frameworks used for LLM post-training.

The diversity improvements are particularly notable. By learning the full advantage distribution rather than collapsing to single high-reward modes, LAD-trained models generate varied reasoning paths that remain valid. This suggests the method successfully addresses the exploration-exploitation tradeoff that plagues reward-maximization approaches.

Implications for LLM Reasoning

LAD represents a shift in how we think about optimizing reasoning in large language models. Rather than treating reasoning as a single-path optimization problem, the distribution-matching perspective acknowledges that multiple valid reasoning trajectories exist and should be preserved during training.

This has several practical implications:

Robustness: Models that maintain diverse reasoning strategies may be more robust to distribution shift and edge cases where the dominant strategy fails.

Interpretability: Preserving multiple reasoning paths provides richer insight into model behavior and decision-making processes.

Composability: Diverse reasoning capabilities may compose more effectively with other techniques like ensemble methods or multi-step reasoning frameworks.

The theoretical grounding in distribution matching also opens avenues for future work. The $f$-divergence framework is flexible, allowing different divergence choices to encode different inductive biases about the desired policy distribution. The connection to advantage distributions suggests potential synergies with other RL techniques that reason about value distributions rather than expected values.

Takeaways

1. Learning advantage distributions rather than maximizing expected rewards prevents overfitting to dominant reward signals while maintaining valid alternative reasoning paths
2. The distribution-matching framework naturally prevents mode collapse without requiring auxiliary entropy regularization terms
3. LAD incurs zero additional training cost compared to standard policy gradient methods like GRPO
4. Experimental validation shows consistent improvements in both accuracy and diversity across math and code reasoning tasks
5. The theoretical framework establishes a principled connection between optimal policy updates and advantage-based target distributions
6. The approach scales naturally to large language model post-training and works across different model architectures