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

The paper introduces UpSkill, a training method that adapts Mutual Information Skill Learning within Group Relative Policy Optimization to enhance response diversity and improve pass@k performance in large language models without compromising single-attempt accuracy.

The Gist

The Big Problem: The "Echo Chamber" Effect

Imagine you ask a smart AI to solve a math problem. You ask it once, and it gives you an answer. You ask it again, and it gives you the exact same answer, word for word. You ask it a third time, and it's still the same.

This is a common issue with Large Language Models (LLMs). When we train them to be "correct," they tend to become very confident and very repetitive. They find one path to the answer and stick to it like a dog on a leash.

Why is this bad?
Think of it like a team of detectives trying to solve a mystery.

• The Old Way: You send 5 detectives out, but they all follow the exact same clue, walk the same path, and check the same house. If that house is a dead end, all 5 fail.
• The Goal: You want 5 detectives who split up. One checks the basement, one checks the attic, one interviews the neighbor, and one looks at the security footage. Even if the first three fail, the fourth one might find the key.

In the world of AI, this is called pass@k. It means: "If we let the AI try $k$ times, what are the odds that at least one of those tries is correct?" If the AI just repeats itself, your odds don't get better no matter how many times you ask.

The Solution: UpSkill (The "Strategy Switch")

The researchers at Princeton University created a method called UpSkill. Their goal was to teach the AI to have a "toolbox" of different strategies instead of just one hammer.

Here is how they did it, using a simple analogy:

1. The "Secret Handshake" (The Latent Variable $z$)

Imagine you have a robot chef. Usually, you just say, "Make me a sandwich." The robot makes the same sandwich every time.

With UpSkill, you give the robot a secret code before it starts.

• If you say "Code 1," the robot makes a sandwich with a knife.
• If you say "Code 2," the robot makes a sandwich with a spoon.
• If you say "Code 3," the robot makes a sandwich using a blender.

The robot learns that different codes lead to different styles of making the sandwich. Crucially, the robot learns that all these styles can still result in a delicious sandwich.

2. The "Diversity Reward" (Mutual Information)

How do you teach the robot to actually use these different styles? You can't just tell it "be different." You have to reward it for being distinct.

The researchers invented a special "score" called Mutual Information.

• The Test: If the robot uses "Code 1," does it always make a knife-sandwich? If yes, great! That's a distinct strategy.
• The Punishment: If the robot uses "Code 1" and sometimes makes a knife-sandwich, but other times makes a spoon-sandwich (just like it would have without the code), then the code is useless. The robot gets a low score.

The AI is trained to maximize this score. It learns: "To get a high score, I must make sure that when I use Code 1, I do something totally different than when I use Code 2."

3. The Result: A Team of Specialists

After training, when you ask the AI a hard math problem, you don't just ask it once. You ask it 5 times, each time giving it a different "Code" (Strategy 1, Strategy 2, etc.).

• Strategy 1 might try to solve the problem using algebra.
• Strategy 2 might try to draw a picture.
• Strategy 3 might try to guess and check.

Even if the AI isn't perfect at any single strategy, the chance that at least one of them gets the right answer goes way up.

What Did They Find?

The researchers tested this on famous math datasets (like GSM8K) using different AI models.

• The Good News: For smart models (like Qwen and Llama), UpSkill made them much better at solving problems when given multiple tries. They got about 3% better at getting the right answer at least once, without making them worse at getting the right answer on the very first try.
• The "Aha!" Moment: They proved mathematically that the more "distinct" the strategies are (the higher the Mutual Information), the better the team performs. It's not magic; it's information theory.
• The Caveat: It didn't work perfectly for every model. For a very small, highly specialized model (R1-Distilled), the method actually made things worse. It seems like if a model is already too "stuck" in its ways or too small to learn new tricks, forcing it to be diverse can break it.

The Takeaway

UpSkill is like teaching a student not just how to solve a problem, but how to think about the problem in five different ways.

Instead of training an AI to be a single, confident expert who might be wrong, they trained it to be a diverse team of experts. When you need an answer, you don't just ask one person; you ask the whole team, each using their unique perspective. If one fails, another might succeed.

In short: Don't just ask the AI to be right. Ask it to be right in many different ways.

Technical Summary

1. Problem Statement

Large Language Models (LLMs) have demonstrated significant reasoning capabilities in mathematics and programming through Reinforcement Learning with Verifiable Rewards (RLVR). However, standard RL approaches that optimize for single-attempt accuracy (pass@1) often inadvertently suppress response diversity. When an LLM is asked to generate multiple solutions (e.g., for code generation or formal proofs), standard sampling often yields highly correlated, redundant outputs.

This lack of diversity is detrimental to multi-attempt metrics like pass@k (the probability that at least one of $k$ attempts is correct). If $k$ attempts are identical or highly similar, the effective number of independent attempts is low, reducing the likelihood of finding a correct solution even if the model's single-attempt accuracy is high. Existing methods to increase diversity (e.g., temperature sampling, prompt perturbation) are often brittle, require manual tuning, or fail to produce semantically distinct reasoning strategies.

2. Methodology: UpSkill

The authors propose UpSkill, a training-time method that induces structured response diversity by adapting Mutual Information Skill Learning (MISL) to LLMs.

Core Concept

The method introduces a discrete latent variable $z \in \{1, \dots, N\}$ (a "skill" or "strategy") into the prompt. The model is trained to condition its response $\tau$ on both the input $x$ and the latent $z$. The goal is to maximize the conditional mutual information $I(\tau; z | x)$, ensuring that:

1. Different values of $z$ lead to distinct, reproducible reasoning strategies.
2. The set of strategies covers a broad range of the solution space.

Technical Implementation

The method is implemented within the Group Relative Policy Optimization (GRPO) framework, a variant of PPO commonly used for RLVR.

• Token-Level Mutual Information Reward: Instead of estimating MI over full trajectories (which is computationally expensive and noisy), UpSkill defines a token-level reward $r_{TMI}$. For a trajectory $\tau$ generated under skill $z$, the reward is the log-ratio of the probability of the token under the specific skill versus the probability under the uniform mixture of all skills:
  $$r_{TMI}(\tau; x, z) = \sum_{t} \left[ \log p_\pi(y_t | x, z, y_{<t}) - \log \left( \frac{1}{N} \sum_{z'} p_\pi(y_t | x, z', y_{<t}) \right) \right]$$
  This encourages the model to make token choices that are specific to the chosen $z$, thereby increasing the distinguishability of strategies.

• Combined Objective: The total reward for a trajectory is a weighted sum of:

  1. Verifiable Correctness ($r_{corr}$): Binary reward (1 if correct, 0 otherwise).
  2. KL Regularization: Keeps the policy close to the base model to prevent catastrophic forgetting.
  3. MI Reward ($r_{TMI}$): Encourages diversity.
     $$r(\tau) = r_{corr}(\tau) - \beta \Delta_{KL}(\tau) + \alpha_1 r_{TMI}(\tau)$$

• Inference: At test time, the system samples $k$ distinct latent values $z_1, \dots, z_k$ and generates one completion for each. Because the strategies are trained to be distinct, these $k$ attempts are less correlated, increasing the probability that at least one is correct.

3. Theoretical Contributions

The paper provides a theoretical link between the mutual information objective and the improvement in pass@k.

• Lower Bound on Improvement: The authors prove that the improvement in pass@k is lower-bounded by a function of the mutual information $I(\tau; z | x)$.
• Implication: Maximizing mutual information during training theoretically guarantees an increase in the multi-attempt success rate, provided the strategies induce different probabilities of success (distributional impact). This establishes a principled justification for using MI to boost pass@k.

4. Key Results

Experiments were conducted on the GSM8K (grade-school math) benchmark and a controlled arithmetic environment using three open-weight models: Llama 3.1-8B, Qwen 2.5-7B, and R1-Distilled-Qwen2.5-Math-1.5B.

• Performance Gains:
  • Qwen 2.5-7B: Achieved a mean gain of ~3.4% in pass@k and ~9.1% in plurality@k without degrading pass@1.
  • Llama 3.1-8B: Showed similar improvements in multi-attempt metrics.
  • Arithmetic Environment: In a controlled setting, UpSkill maintained diverse trajectories (pass@5 = 0.897) while standard GRPO collapsed to a single strategy (pass@5 = pass@1 = 0.793).
• Diversity without Ground Truth: Ablation studies showed that UpSkill can improve pass@k even when trained only with the MI reward and no correctness labels, demonstrating that the method learns diverse strategies intrinsically.
• Model Sensitivity:
  • The method worked well on larger models (Qwen, Llama).
  • It caused performance degradation on the smaller R1-Distilled model, likely due to its limited capacity and heavy prior optimization, suggesting that UpSkill requires sufficient model capacity to learn distinct modes.
• Interpretability: Analysis of the learned strategies revealed that different $z$ values corresponded to distinct reasoning patterns (e.g., algebraic vs. narrative approaches, or specific operator choices in the arithmetic task).

5. Significance and Conclusion

UpSkill addresses a critical gap in LLM reasoning: the trade-off between accuracy and diversity.

• Practical Impact: It offers a training-time solution to improve reliability in multi-attempt scenarios (e.g., code generation, formal verification) without requiring complex decoding-time heuristics or prompt engineering.
• Theoretical Insight: It bridges information theory and RL, proving that maximizing mutual information between latent skills and trajectories is a sufficient condition for improving multi-attempt success rates.
• Future Directions: The authors note that while the method is effective for larger models, it requires careful tuning for smaller models and future work could explore more robust semantic MI estimators to ensure strategies are not just token-diverse but semantically distinct.

In summary, UpSkill provides a principled, scalable framework for inducing structured, reproducible diversity in LLMs, significantly boosting their reliability in tasks where multiple attempts are feasible.