LAD: Learning Advantage Distribution for Reasoning

Wendi Li; Sharon Li

cs.LGFeb 23, 2026

LAD: Learning Advantage Distribution for Reasoning

Authors: Wendi Li, Sharon Li

Organizations: Department of Computer Sciences, University of Wisconsin–Madison

Abstract

Current reinforcement learning objectives for large-model reasoning primarily focus on maximizing expected rewards. This paradigm can lead to overfitting to dominant reward signals, while neglecting alternative yet valid reasoning trajectories, thereby limiting diversity and exploration. To address this issue, we introduce Learning Advantage Distributions (LAD), a distribution-matching framework that replaces advantage maximization with learning the advantage-induced distribution. By establishing the equivalence between the optimal policy update and an advantage-based target distribution, we derive a practical LAD objective formulated as minimizing an $f$ -divergence between the policy-induced and advantage-induced distributions. This yields a gradient update that increases likelihood for high-advantage responses while suppressing over-confident probability growth, preventing collapse without requiring auxiliary entropy regularization. LAD incurs no extra training cost compared to GRPO and scales naturally to LLM post-training. In a controlled bandit setting, LAD faithfully recovers the multimodal advantage distribution, validating the theoretical formulation. Experiments on math and code reasoning tasks across several LLM backbones show that LAD reliably improves both accuracy and generative diversity.

Figures & tables

Figure 1: Advantage distribution and policy distribution on 50-arm bandit trained via different training objectives.

Figure 2: Learning dynamics of GRPO, the theoretical LAD loss , and the practical LAD loss . The loss landscape is constructed with the practical LAD loss. Top: LAD loss under Hellinger distance. Bottom: LAD loss under Jensen–Shannon divergence.

Methods	MATH	Olympiad Bench	Minerva	AIME 2024	AIME 2025	AMC	Average
GRPO ( Shao et al., 2024 )	75.57	39.47	36.32	14.03	4.86	53.44	37.28
EntAdv ( Cheng et al., 2025 )	76.60	41.16	35.68	15.00	7.23	52.97	38.11
KLCov ( Cui et al., 2025b )	76.40	40.13	36.28	14.66	9.81	53.45	38.46
ClipCov ( Cui et al., 2025b )	76.03	40.30	37.58	16.29	7.39	52.27	38.31
FlowRL ( Zhu et al., 2025 )	77.17	40.80	36.21	14.67	7.51	47.34	37.28
LAD	77.66	41.25	36.68	19.60	8.84	56.45	40.08

Table 1: Performance of different losses on math reasoning benchmarks. The best two results under Avg@32 are highlighted in bold and underlined .

Method	LiveCodeBench		HumanEval+		CodeForces
Method	Avg@16	Pass@16	Avg@16	Pass@16	Rating	Percentile
GRPO ( Shao et al., 2024 )	32.46	51.25	80.90	93.25	1355.35	70.30 %
EntAdv ( Cheng et al., 2025 )	32.75	50.90	81.05	95.09	1382.78	72.20%
KLCov ( Cui et al., 2025b )	29.50	50.18	80.75	95.09	1373.07	71.60%
ClipCov ( Cui et al., 2025b )	32.10	51.97	81.48	94.49	1473.09	79.20%
FLowRL ( Zhu et al., 2025 )	33.24	51.61	81.29	95.09	1364.70	71.00%

Table 2: Performance of different losses on code reasoning benchmarks. The best two results under a specific metric are highlighted in bold and underlined .

Methods	MATH	Olympiad Bench	Minerva	AIME 2024	AIME 2025	AMC	Average
OpenMath-Nemotron-1.5B: Avg@32
FlowRL	80.43	45.81	21.02	21.36	22.66	55.47	41.13
LAD	86.81	55.63	22.98	32.70	26.47	68.71	48.88
OpenMath-Nemotron-1.5B: Pass@32
FlowRL	94.60	65.93	38.24	56.35	37.60	80.65	66.98
LAD	97.20	78.52	42.28	72.29	54.38	89.08	72.29

Table 3: Performance of LAD and FlowRL on different LLM backbones. The best result is highlighted in bold .

Methods	dist-3( $\uparrow$ )	dist-4( $\uparrow$ )	GPT4-Judge( $\uparrow$ )
GRPO	0.2306	0.2902	2.04
EntAdv	0.2978	0.3782	2.07
KLCov	0.2471	0.3262	2.07
ClipCov	0.3315	0.4171	2.10
FLowRL	0.2878	0.3654	2.38
LAD	0.3498	0.4442	2.58

Table 4: Generative diversity results of on rollouts of AIME 24/25 problems. The best two results under a specific metric are highlighted in bold and underlined .

Figure 5: The distribution of GPU time per token throughout the training process.

Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: The average advantage value and the log-ratio of

|\log\frac{\pi_{\theta}(y|x)}{\pi_{\mathrm{old}}(y|x)}|

throughout the training.

Divergence	$f(x)$	$\mathcal{L}_{\mathrm{LAD}}$
KL divergence	$x\log x$	$\mathbb{E}_{x\sim\rho,y\sim\pi_{\mathrm{old}}(\cdot\|x)}\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}\Big(\log\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}-\frac{A(x,y)}{\eta}\Big)$
reverse KL divergence	$-\log x$	$\mathbb{E}_{x\sim\rho,y\sim\pi_{\mathrm{old}}(\cdot\|x)}-{e^{\frac{A(x,y)}{\eta}}}\log{\pi_{\theta}(y\|x)}$
Jeffreys divergence	$(x-1)\log x$	$\mathbb{E}_{x\sim\rho,y\sim\pi_{\mathrm{old}}(\cdot\|x)}(\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}-e^{\frac{A(x,y)}{\eta}})\log[\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}e^{\frac{-A(x,y)}{\eta}}]$
Total Variance Distance ∗	$\|x-1\|$	$\mathbb{E}_{x\sim\rho,y\sim\pi_{\mathrm{old}}(\cdot\|x)}{e^{\frac{A(x,y)}{\eta}}}\Big\|\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}\cdot e^{\frac{-A(x,y)}{\eta}}-1\Big\|$
Hellinger Distance ∗	$\frac{1}{2}(\sqrt{t}-1)^{2}$	$\mathbb{E}_{x\sim\rho,y\sim\pi_{\mathrm{old}}(\cdot\|x)}{e^{\frac{A(x,y)}{\eta}}}(\sqrt{\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}}\cdot e^{\frac{-A(x,y)}{2\eta}}-1)^{2}$
Jensen–Shannon divergence ∗	$\frac{1}{2}(x\log x-(x+1)\log\frac{x+1}{2})$	$\begin{aligned} \scriptstyle\mathbb{E}_{x\sim\rho,y\sim\pi_{\mathrm{old}}(\cdot\|x)}{e^{\frac{A(x,y)}{\eta}}}\Big(\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}e^{\frac{-A(x,y)}{\eta}}\log(\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}\\ \scriptstyle\cdot e^{\frac{-A(x,y)}{\eta}})-(\frac{\pi_{\theta}(y\|x)}{\pi_{\mathrm{old}}(y\|x)}e^{\frac{A(x,y)}{\eta}}+1)\log(\frac{\pi_{\theta}(y\|x)}{2\pi_{\mathrm{old}}(y\|x)}e^{\frac{A(x,y)}{\eta}})\Big)\end{aligned}$

Appendix

Table 5: LAD losses under different divergences. The

*

represents the corresponding divergence class is a strict divergence that forces exact matching at its minimum.

Methods	$1/\eta$	MATH	Olympiad Bench	Minerva	AIME 2024	AIME 2025	AMC	Average
		Avg@16
GRPO	-	75.71	39.24	36.63	13.89	4.90	53.39	37.29
LAD-JS	1/4	75.24	38.83	35.59	15.68	5.70	50.34	36.90
	1/2	76.20	38.02	38.07	18.15	6.73	53.08	38.38
	1	74.84	37.16	39.84	14.95	11.28	55.98	39.00
	2	76.90	36.50	34.93	18.67	9.13	55.12	38.54

Appendix

Table 6: The detailed math reasoning performances of LAD with different hyperparameters.

Methods	MATH	Olympiad Bench	Minerva	AIME 2024	AIME 2025	AMC	Average
GRPO ( Shao et al., 2024 )	90.60	60.15	55.88	30.52	26.77	80.50	57.40
EntAdv ( Cheng et al., 2025 )	90.60	62.51	55.51	35.83	34.27	81.93	60.10
KLCov ( Cui et al., 2025b )	92.80	64.15	56.99	36.25	41.25	85.99	63.04
ClipCov ( Cui et al., 2025b )	91.40	61.48	60.29	37.40	32.08	80.61	60.54
FlowRL ( Zhu et al., 2025 )	94.20	69.04	59.56	37.50	35.52	82.94	63.13
LAD	93.80	67.70	59.93	47.19	37.08	88.86	65.76

Appendix

Table 7: Pass@32 performance of different losses on math reasoning benchmarks. The best two results under a specific metric are highlighted in bold and underlined .

Methods	$k$	MATH	Olympiad Bench	Minerva	AIME 2024	AIME 2025	AMC	Average
		Avg@ $k$
LAD-KL	@1	76.80	37.33	38.97	20.00	5.83	51.13	38.01
	@16	76.35	36.93	39.18	20.03	4.64	51.74	38.15
	@32	76.20	36.83	38.73	19.91	4.68	51.63	37.99
LAD-JF	@1	72.60	32.74	37.13	15.52	7.50	53.69	36.20
	@16	72.65	32.54	34.92	15.96	7.23	53.47	36.13

Appendix

Table 8: Performance of different losses on math reasoning benchmarks. The best two results under a specific metric are highlighted in bold and underlined .

RL for Language Model Reasoning

Reinforcement learning for language model reasoning uses reward-driven training to improve multi-step problem solving, tool use, and reliable answers.

See all 1,094 papers

Reinforcement-aware Knowledge Distillation for LLM Reasoning
Feb 26, 2026Zhaoyang Zhang, Shuli Jiang, Yantao Shen +6Knowledge Distillation RL for Knowledge Distillation
AWS Agentic AI
The Easy, the Hard, and the Learnable: Confidence and Difficulty-Adaptive Policy Optimization for LLM Reasoning
Jun 6, 2026Zhanke Zhou, Xiangyu Lu, Chentao Cao +4RL for Language Model Reasoning Policy Optimization
1TMLR Group, Department of Computer Science, Hong Kong Baptist University · 2Stanford University · 3Sydney AI Centre, The University of Sydney
DSDR: Dual-Scale Diversity Regularization for Exploration in LLM Reasoning
Feb 23, 2026Zhongwei Wan, Yun Shen, Zhihao Dou +9Reinforcement Learning RL for Language Model Reasoning
The Ohio State University, USA · Case Western Reserve University, USA · The Chinese University of Hong Kong, Hong Kong +4

Feb 26, 2026J/K move · Enter open · S save

Reinforcement-aware Knowledge Distillation for LLM Reasoning

Zhaoyang Zhang, Shuli Jiang, Yantao Shen, Yuting Zhang+5

AWS Agentic AI

Knowledge Distillation RL for Knowledge Distillation On-Policy Distillation Selective Knowledge Distillation RL for Language Model Reasoning

Open paper AskPDF arXiv

Reinforcement learning (RL) post-training has recently driven major gains in long chain-of-thought reasoning large language models (LLMs), but the high inference cost of such models motivates distillation into smaller students. Most existing knowledge distillation (KD) methods are designed for supervised fine-tuning (SFT), relying on fixed teacher traces or teacher-student Kullback-Leibler (KL) divergence-based regularization. When combined with RL, these approaches often suffer from distribution mismatch and objective interference: teacher supervision may not align with the student's evolving rollout distribution, and the KL regularizer can compete with reward maximization and require careful loss balancing. To address these issues, we propose RL-aware distillation (RLAD), which performs selective imitation during RL -- guiding the student toward the teacher only when it improves the current policy update. Our core component, Trust Region Ratio Distillation (TRRD), replaces the teacher-student KL regularizer with a PPO/GRPO-style likelihood-ratio objective anchored to a teacher--old-policy mixture, yielding advantage-aware, trust-region-bounded distillation on student rollouts and naturally balancing exploration, exploitation, and imitation. Across diverse logic reasoning and math benchmarks, RLAD consistently outperforms offline distillation, standard GRPO, and KL-based on-policy teacher-student knowledge distillation.

Abstract

RL for Language Model Reasoning

Explore similar work

Reinforcement-aware Knowledge Distillation for LLM Reasoning

The Easy, the Hard, and the Learnable: Confidence and Difficulty-Adaptive Policy Optimization for LLM Reasoning

DSDR: Dual-Scale Diversity Regularization for Exploration in LLM Reasoning