UpSkill: Mutual Information Skill Learning for Structured Response Diversity in LLMs

Devan Shah; Owen Yang; Daniel Yang; Chongyi Zheng; Benjamin Eysenbach

cs.LGFeb 25, 2026

UpSkill: Mutual Information Skill Learning for Structured Response Diversity in LLMs

Authors: Devan Shah, Owen Yang, Daniel Yang, Chongyi Zheng, Benjamin Eysenbach

Organizations: Princeton University

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) has improved the reasoning abilities of large language models (LLMs) on mathematics and programming tasks, but standard approaches that optimize single-attempt accuracy can inadvertently suppress response diversity across repeated attempts, narrowing exploration and overlooking underrepresented strategies. We introduce UpSkill, a training time method that adapts Mutual Information Skill Learning (MISL) to LLMs for optimizing pass@k correctness. We propose a novel reward that we implement within Group Relative Policy Optimization (GRPO): a token-level mutual information (MI) reward that encourages trajectory specificity to z. Experiments on GSM8K with three open-weight models, Llama 3.1-8B, Qwen 2.5-7B, and R1-Distilled-Qwen2.5-Math-1.5B, show that UpSkill improves multi-attempt metrics on the stronger base models, yielding mean gains of ~3% in pass@k for both Qwen and Llama without degrading pass@1. Additionally, we find both empirical and theoretical evidence that improvements in pass@k are closely tied to the mutual information objective.

Figures & tables

Figure 1: UpSkill improves mean multi-attempt accuracy without hurting single-attempt accuracy on GSM8K for the Qwen 2.5-7B model (See Sec. 5.2 ).

Figure 2: UpSkill is an unsupervised method for training LLMs to produce diverse responses. After training, different latent vectors $z$ (blue boxes above) correspond to different response strategies. Because of space constraints, the figure shows summarized responses from UpSkill; we report the full responses in Appendix K .

Figure 3: Example illustration of how the MISL reward improves pass@k performance. Before MISL (left), the trajectory distribution is independent of the latents $z$ , so the conditional entropy is close to the marginal. MISL training prevents distribution collapse due to pass@1 training (middle). Adding the token-level MI reward (right) yields well-separated clusters indexed by $z$ , reducing conditional entropy while preserving high marginal entropy. At inference, fixing different $z$ values produces consistent and diverse solution strategies.

Figure 4: Arithmetic environment results. Training curves show that under GRPO alone (blue), pass@1 and pass@5 converge together, indicating that multiple attempts provide little benefit. With MISL (orange; $N{=}5$ ), pass@5 improves substantially while pass@1 remains modest, demonstrating that different latents yield complementary solutions. Operator distributions further highlight this effect: without MISL, they are nearly identical across $z$ , reflecting a lack of specialization, whereas with MISL, distinct latents focus on different operators, producing diverse strategies that drive multi-attempt gains.

Figure 5: Performance on $500$ held-out problems with $N{=}5$ strategies. We observe gains on all metrics for the Qwen model. Base refers to the model before GRPO training, Without MI refers to after GRPO training without token MI, and With MI refers to training with correctness rewards and token MI. We test multiple configurations for With MI and Without MI and plot all successful runs as multiple bars, as elaborated in Appendix G . For entries with multiple bars, the labeled value is the maximum.

Figure 6: GSM8K pass@5 scores over selected checkpoints of two training processes, overlaid on the mutual information rewards during training. The tuples in the legend are of the form $(\alpha_{1},\text{lr})$ .

Model	Method	$\boldsymbol{\Delta}\textbf{pass@k}$	Method	$\boldsymbol{\Delta}\textbf{pass@k}$
Qwen	Correctness	+3.8%	MI	+5.2%
Qwen	MI + correctness	+7.0%	MI + KL	+4.0%
Llama	Correctness	+2.8%	MI	+0.8%
Llama	MI + correctness	+3.6%	MI + KL	+2.4%

Table 1: Improvement in pass@k relative to the base model for MI-based methods. Full results can be found in Table 10 .

Model	Method	$\boldsymbol{\Delta}\textbf{pass@k}$	Method	$\boldsymbol{\Delta}\textbf{pass@k}$
R1	Correctness	+5.2%	MI	-8.4%
R1	MI + correctness	-9.8%	MI + KL	+0.8%

Table 2: Changes in pass@k relative to the base model for the R1 model fine-tuned on correctness only or using our method. Full results can be found in Table 10 .

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

	After Warmup (Step 0)			Step 1000			Step 2000
Condition	p@1	p@5	$H_{m}$	p@1	p@5	$H_{m}$	p@1	p@5	$H_{m}$
Control – ( $\alpha_{1}=0$ )	0.313	0.665	0.723	0.668	0.673	0.016	0.793	0.793	0.013
( $\alpha_{1}=0.5$ , cap=1.0)	0.313	0.665	0.723	0.353	0.843	0.755	0.390	0.897	0.768
( $\alpha_{1}=0.5$ , cap=1.5)	0.313	0.665	0.723	0.338	0.833	0.764	0.399	0.830	0.852
( $\alpha_{1}=1.0$ , cap=1.0)	0.313	0.665	0.723	0.281	0.813	0.813	0.373	0.897	0.844

Appendix

Table 3: Arithmetic environment training outcomes. Accuracy is reported as pass@1 and pass@5 ;

H_{m}

is marginal token entropy. MISL prevents entropy collapse and sustains diverse skill-conditioned strategies.

Figure 7: Learned distribution over first response digit with $\alpha_{1}=0.5$ and $\text{cap}=1.0$

Warmup Steps	pass@1	pass@k
50	0.235	0.540
100	0.313	0.665

Appendix

Table 4: Performance of the base model after warmup only (no RL). Warmup 50 yields a weaker base capacity, while warmup 100 yields a stronger base capacity.

kl_coef	warmup	C	$\mathbf{pass@1_{a0.0}}$	$\mathbf{pass@1_{a0.1}}$	$\mathbf{pass@1_{a0.3}}$	$\mathbf{pass@1_{a0.5}}$
0.050	50	5	0.779	0.780	0.561	0.457
0.050	50	10	0.845	0.817	0.655	0.439
0.100	50	5	0.813	0.801	0.557	0.366
0.100	50	10	0.833	0.860	0.629	0.437
0.050	100	5	0.908	0.897	0.749	0.521
0.050	100	10	0.914	0.897	0.767	0.521

Appendix

Table 5: pass@1 after RL with KL penalty in the arithmetic environment.

kl_coef	warmup	C	$\mathbf{pass@5_{a0.0}}$	$\mathbf{pass@5_{a0.1}}$	$\mathbf{pass@5_{a0.3}}$	$\mathbf{pass@5_{a0.5}}$
0.050	50	5	0.793	0.807	0.840	0.870
0.050	50	10	0.867	0.833	0.867	0.817
0.100	50	5	0.840	0.843	0.840	0.847
0.100	50	10	0.857	0.893	0.850	0.820
0.050	100	5	0.940	0.917	0.863	0.883
0.050	100	10	0.927	0.903	0.903	0.887

Appendix

Table 6: pass@5 after RL with KL penalty in the arithmetic environment.

Figure 8: Graph of model rewards during training in which it suffers strategy collapse.

				GSM8K (In-Distribution)				HumanEval
Model	Method	$\boldsymbol{\alpha_{1}}$	lr	pass@1	pass@k	plur@k	cons@k	pass@1	pass@5
Qwen	Base	–	–	53.4%	86.0%	61.0%	55.0%	68.8%	80.0%
	Control	–	1e-5	54.9%	87.2%	65.2%	56.2%	68.9%	80.5%
	Control	–	3e-5	58.5%	89.8%	66.4%	60.2%	70.5%	82.9%
	Experimental	1	1e-5	64.7%	93.0%	75.0%	69.4%	70.1%	82.9%
	Experimental	1	3e-5	63.5%	91.2%	74.6%	67.8%	73.7%	82.3%

Appendix

Table 7: Detailed performance metrics for all individual experimental runs on GSM8K and HumanEval (OOD). Bold values indicate the highest performance within each model group.

Model	$\boldsymbol{\alpha_{1}}$	lr	Acc (Last 100)	Avg MI (Last 100)
Qwen	1	1e-5	91.9%	0.43
	1	3e-5	95.4%	8.05
	5	1e-5	16.9%	-4.26
	5	3e-5	95.4%	8.05
Llama	1	1e-5	68.7%	0.08
	1	3e-5	74.7%	-0.57

Appendix

Table 8: Training dynamics over the final 100 steps.

Figure 9: Llama results for learning rate of $1\times 10^{-6}$ and $\alpha_{1}=1$ (green bars) compared to $5\times 10^{-6}$ and $\alpha_{1}=5$ (blue bars). We simultaneously ablate the result of the strategy prompting. In each pair of 4 pairs, the leftmost represents the $5\times 10^{-6}$ without the strategy prompt, the second furthest to the left represents $5\times 10^{-6}$ with the strategy prompt, followed by the models trained with learning rate $1\times 10^{-6}$ and evaluated without and with the strategy prompt, respectively.

Panel of Figure 10 — Figure 10: Rewards obtained throughout training.

Model	pass@1	pass@k	plurality@k	consensus@k	$\boldsymbol{\Delta}\textbf{pass@k}$
Llama (checkpoint 200)	75.2%	91.2%	80.8%	77.4%	+0.8%
Qwen (checkpoint 800)	61.6%	91.2%	72.0%	65.6%	+4.2%

Appendix

Table 9: GSM8K performance with MI-only fine-tuning. The

\Delta\text{pass@k}

column is the change relative to the base model with no fine-tuning.

Model	Rewards used	pass@1	pass@k	plurality@k	consensus@k	$\boldsymbol{\Delta}\textbf{pass@k}$
Qwen	Base model	53.4%	86.0%	61.0%	55.0%	0.0%
	Correctness	58.5%	89.8%	66.4%	60.2%	+3.8%
	MI + correctness	64.7%	93.0%	75.0%	69.4%	+7.0%
	MI	61.6%	91.2%	72.0%	65.6%	+5.2%
	MI + KL	55.0%	90.0%	67.4%	45.6%	+4.0%
Llama	Base model	77.0%	90.4%	81.8%	79.0%	0.0%

Appendix

Table 10: Performance comparison on GSM8K across Qwen, Llama, and R1 models. The correctness and MI + correctness rows taken are those with the highest pass@k values over different configurations. Change in pass@k is calculated relative to the base model.

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

Prompt Interference in LLM Post-Training: Why Pass@k Optimization Can Degrade Pass@1
Feb 24, 2026Anas Barakat, Souradip Chakraborty, Khushbu Pahwa +1RL for Language Model Reasoning
Singapore University of Technology and Design · University of Maryland, College Park · University of Central Florida
ExpRL: Exploratory RL for LLM Mid-Training
Jun 15, 2026Violet Xiang, Amrith Setlur, Chase Blagden +2Reinforcement Learning RL for Language Model Reasoning
1Stanford University · 2Carnegie Mellon University · 3OpenAI +1
Skill-R1: Agent Skill Evolution via Reinforcement Learning
May 10, 2026Yash Vishe, Rohan Surana, Xunyi Jiang +8Agent Skill Learning RL for LLM Agents
UC San Diego · Adobe Research

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

Prompt Interference in LLM Post-Training: Why Pass@k Optimization Can Degrade Pass@1

Anas Barakat, Souradip Chakraborty, Khushbu Pahwa, Amrit Singh Bedi

Singapore University of Technology and Design · University of Maryland, College Park · University of Central Florida

RL for Language Model Reasoning

Open paper AskPDF arXiv

Pass@

k

is a widely used performance metric for verifiable large language model tasks, including mathematical reasoning, code generation, and short-answer reasoning. It defines success if any of

k

independently sampled solutions passes a verifier. This multi-sample inference metric has motivated inference-aware fine-tuning methods that directly optimize pass@

k

. However, prior work reports a trade-off: pass@

k

improves while pass@1 degrades under such methods. This trade-off is practically important because pass@1 often remains a hard operational constraint due to latency and cost budgets, imperfect verifier coverage, and the need for a reliable single-shot fallback. We study the origin of this trade-off and provide a theoretical characterization of when pass@

k

policy optimization can reduce pass@1 through gradient conflict induced by prompt interference. This interference arises because shared model parameters couple prompts, so improving performance on some prompts can adversely affect others. We show that pass@

k

policy gradients can conflict with pass@1 gradients because pass@

k

optimization implicitly reweights prompts toward low-success prompts; when these prompts are negatively interfering, their upweighting can rotate the pass@

k

update direction away from the pass@1 direction. We illustrate our findings with large language model experiments on verifiable mathematical reasoning and coding tasks.

Abstract

RL for Language Model Reasoning

Explore similar work

Prompt Interference in LLM Post-Training: Why Pass@k Optimization Can Degrade Pass@1

ExpRL: Exploratory RL for LLM Mid-Training

Skill-R1: Agent Skill Evolution via Reinforcement Learning