UpSkill: Mutual Information Skill Learning for Structured Response Diversity in LLMs
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
| Model | Method | Method | ||
|---|---|---|---|---|
| Qwen | Correctness | +3.8% | MI | +5.2% |
| MI + correctness | +7.0% | MI + KL | +4.0% | |
| Llama | Correctness | +2.8% | MI | +0.8% |
| MI + correctness | +3.6% | MI + KL | +2.4% |
| Model | Method | Method | ||
|---|---|---|---|---|
| R1 | Correctness | +5.2% | MI | -8.4% |
| MI + correctness | -9.8% | MI + KL | +0.8% |
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 | p@1 | p@5 | p@1 | p@5 | |||
| Control – ( ) | 0.313 | 0.665 | 0.723 | 0.668 | 0.673 | 0.016 | 0.793 | 0.793 | 0.013 |
| ( , cap=1.0) | 0.313 | 0.665 | 0.723 | 0.353 | 0.843 | 0.755 | 0.390 | 0.897 | 0.768 |
| ( , cap=1.5) | 0.313 | 0.665 | 0.723 | 0.338 | 0.833 | 0.764 | 0.399 | 0.830 | 0.852 |
| ( , cap=1.0) | 0.313 | 0.665 | 0.723 | 0.281 | 0.813 | 0.813 | 0.373 | 0.897 | 0.844 |
| Warmup Steps | pass@1 | pass@k |
|---|---|---|
| 50 | 0.235 | 0.540 |
| 100 | 0.313 | 0.665 |
| kl_coef | warmup | C | ||||
|---|---|---|---|---|---|---|
| 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 |
| kl_coef | warmup | C | ||||
|---|---|---|---|---|---|---|
| 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 |
| GSM8K (In-Distribution) | HumanEval | ||||||||
| Model | Method | 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% | |
| Model | 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 |
| Model | pass@1 | pass@k | plurality@k | consensus@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% |
| Model | Rewards used | pass@1 | pass@k | plurality@k | consensus@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% |
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