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

TL;DR

UpSkill introduces a novel training method for large language models that enhances response diversity and correctness across multiple attempts by optimizing mutual information, leading to improved performance on mathematical reasoning tasks.

Contribution

The paper proposes UpSkill, a new training approach that applies mutual information-based rewards within reinforcement learning to improve multi-attempt correctness in LLMs.

Findings

UpSkill improves pass@k metrics by approximately 3% on multiple models.

Mutual information optimization is empirically and theoretically linked to performance gains.

The method enhances response diversity without degrading single-attempt accuracy.

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.