Learning Advantage Distributions (LAD)

• Learning Advantage Distributions (LAD) is a principled reframing of RL objectives that matches the full advantage distribution to preserve multiple valid reasoning paths.
• It minimizes f-divergences between the policy and an exponential advantage-induced target, effectively preventing mode collapse typical of standard policy-gradient methods.
• Empirical results demonstrate LAD’s improved accuracy and diversity in LLM math and code reasoning tasks compared to traditional approaches like GRPO.

Learning Advantage Distributions (LAD) constitutes a principled reframing of reinforcement learning (RL) objectives for large-model reasoning, notably LLMs. Rather than optimizing for maximum expected reward or advantage as in standard policy-gradient approaches, LAD trains the policy to explicitly match the entire distribution over actions induced by their advantages. This approach is designed to preserve alternative valid reasoning trajectories, combating the well-documented tendency of conventional RL algorithms to collapse towards a single dominant solution, thus enhancing both accuracy and diversity in generative reasoning tasks (Li et al., 23 Feb 2026).

1. Theoretical Foundations and Motivation

Traditional RL policy optimization techniques such as REINFORCE and Proximal Policy Optimization (PPO, also known as GRPO in this context) maximize expected reward by focusing the probability mass on actions with the highest estimated advantage:

$$J(\pi_\theta) = \mathbb{E}_{x \sim \rho, y \sim \pi_\theta}[r(x, y)],$$

with advantage $A(x, y) = r(x, y) - b(x)$ and gradient

$$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$

This expected-advantage maximization commonly leads to rapid entropy decline and mode collapse, whereby the policy overfits to a single dominant action trajectory.

By contrast, LAD frames the objective as a distribution-matching problem: for each context $x$, the induced action probability distribution should match the exponential weighting of the advantage,

$$\pi^*_\theta(y|x) \propto \pi_{\text{old}}(y|x) \exp(A(x, y)/\eta),$$

where $\eta$ is a trust-region “temperature” parameter. The result is a formal equivalence between the optimal policy update and an "advantage-induced" target distribution, supporting explicit preservation of multimodal behaviors (Li et al., 23 Feb 2026).

2. Objective Formulation via f-Divergence Minimization

Matching the policy-induced and advantage-induced distributions reduces to minimizing an $f$-divergence. For suitable convex $f$ with $f(1) = 0$, this is expressed as:

$$D_f(P_A \Vert P_\pi) = \sum_y P_A(y|x) f\left( \frac{P_\pi(y|x)}{P_A(y|x)} \right),$$

where $A(x, y) = r(x, y) - b(x)$0 is the normalized exponential advantage distribution and $A(x, y) = r(x, y) - b(x)$1 is the normalized policy ratio. The practical LAD loss, which avoids intractable normalizing terms, becomes:

$A(x, y) = r(x, y) - b(x)$2

This form leverages surrogate weighting so that the unique optimum is attained at the trust-region solution without requiring evaluation of full partition sums (Li et al., 23 Feb 2026).

3. Gradient Dynamics and Implicit Regularization

The LAD gradient for parameter updates is

$A(x, y) = r(x, y) - b(x)$3

or, after change of measure,

$A(x, y) = r(x, y) - b(x)$4

with $A(x, y) = r(x, y) - b(x)$5 and $A(x, y) = r(x, y) - b(x)$6. Importantly, for most strict divergences (Hellinger, Jensen–Shannon, etc.), $A(x, y) = r(x, y) - b(x)$7 as $A(x, y) = r(x, y) - b(x)$8, so LAD actively down-weights gradients for over-confident actions, preventing runaway amplification of already-dominant modes. No auxiliary entropy regularizer is required: diversity preservation is inherent in the divergence-minimization objective itself.

4. Algorithmic Implementation and Hyperparameters

LAD is implemented as an iterative RL post-training procedure with the following workflow:

1. Collect a batch of contexts $A(x, y) = r(x, y) - b(x)$9 from data distribution $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$0.
2. Generate actions $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$1, reward $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$2, and compute advantages $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$3.
3. For each sample, compute $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$4, $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$5, and loss $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$6.
4. The policy gradient is estimated as the batch mean of $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$7.
5. Parameters are updated via standard gradient descent.
6. The behavior policy $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$8 is periodically synchronized with $$\mathbb{E}_{x, y \sim \pi_\theta}\left[ A(x, y) \nabla_\theta \log \pi_\theta(y|x) \right].$$9.

Key hyperparameters include temperature $x$0 (typical range $x$1, default $x$2), choice of divergence $x$3 {KL, reverse KL, Jeffreys, total variation*, Hellinger*, Jensen–Shannon*}, learning rate $x$4, batch size $x$5, and update frequency $x$6 ($x$7 denotes strict divergences) (Li et al., 23 Feb 2026).

5. Empirical Evaluation and Mode Diversity

LAD demonstrates a capacity to recover and maintain multimodal action distributions in settings where expected-advantage optimizers collapse to single modes:

• In tri-modal bandit environments, LAD accurately reconstructs all high-advantage modes, whereas GRPO collapses to a single arm.
• On LLM math reasoning tasks (Qwen2.5-7B, six benchmarks), LAD achieves Avg@32 accuracy of $x$8 versus $x$9 for KLCov and $$\pi^*_\theta(y|x) \propto \pi_{\text{old}}(y|x) \exp(A(x, y)/\eta),$$0 for GRPO.
• For code reasoning (DeepSeek-R1-7B on LiveCodeBench, HumanEval+, CodeForces), LAD outperforms GRPO in both Pass@16 and CodeForces percentile.
• Generative diversity metrics (distinct-3, distinct-4, GPT-4 Judge) reveal systematic gains with LAD: e.g., distinct-3/4 values of $$\pi^*_\theta(y|x) \propto \pi_{\text{old}}(y|x) \exp(A(x, y)/\eta),$$1 for LAD compared to $$\pi^*_\theta(y|x) \propto \pi_{\text{old}}(y|x) \exp(A(x, y)/\eta),$$2 for GRPO.

Ablation studies show that strict divergences yield better mode preservation, and performance is robust to the choice of $$\pi^*_\theta(y|x) \propto \pi_{\text{old}}(y|x) \exp(A(x, y)/\eta),$$3 within the tested range. LAD’s training cost (GPU-time per token) matches that of GRPO, reflecting no additional computational overhead (Li et al., 23 Feb 2026).

Benchmark	GRPO	KLCov	LAD
Math Reasoning Acc.	37.28	38.46	40.08
Code Pass@16	51.25	—	51.97
CodeForces Percent.	70.30%	—	82.50%
Div. (dist-3/4)	0.2306/0.2902	—	0.3498/0.4442

6. Relation to Distributional Critic Methods

Distributional RL is also explored in the context of value distribution critics, notably in the Distributional Advantage Actor–Critic (DA2C/QR-A2C) framework (Li et al., 2018). While DA2C models the full return distribution $$\pi^*_\theta(y|x) \propto \pi_{\text{old}}(y|x) \exp(A(x, y)/\eta),$$4 using quantile regression and computes scalar advantages from it, the actual actor update remains focused on maximizing expected advantage, supplemented by entropy bonuses for exploration. By contrast, LAD applies distribution-matching directly at the policy level via advantage-induced distributions and $$\pi^*_\theta(y|x) \propto \pi_{\text{old}}(y|x) \exp(A(x, y)/\eta),$$5-divergence minimization, introducing inherent regularization and multimodality preservation not present in DA2C or canonical expected-advantage frameworks.

7. Limitations and Prospects

LAD is validated in post-training RL scenarios with verifiable (automatic) rewards in math and code benchmarks for models within the 1.5B–7B parameter range. Open questions remain for:

• Integration with learned reward models in RLHF settings or when rewards are human-provided.
• Scalability and hyperparameter robustness for 30B+ models.
• Theoretical convergence guarantees for arbitrary $$\pi^*_\theta(y|x) \propto \pi_{\text{old}}(y|x) \exp(A(x, y)/\eta),$$6-divergences and in highly nonconvex neural policy landscapes.

Future research directions include adapting LAD for RLHF, fully on-policy online RL, rigorous convergence analysis, and applications in complex multi-modal RL environments (Li et al., 23 Feb 2026).

---

References

• "LAD: Learning Advantage Distribution for Reasoning" (Li et al., 23 Feb 2026)
• "Distributional Advantage Actor-Critic" (Li et al., 2018)