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

Die Arbeit stellt UpSkill vor, eine Trainingsmethode, die mittels gegenseitiger Information die Antwortvielfalt von Large Language Models erhöht und so die Pass@k-Leistung bei mathematischen Aufgaben verbessert, ohne die Genauigkeit des ersten Versuchs zu beeinträchtigen.

Das Wesentliche

Das Problem: Der „Einheitsbrei" der KI

Stell dir vor, du fragst einen sehr klugen, aber etwas sturen Freund (eine KI) immer wieder nach demselben mathematischen Rätsel.

• Frage: „Wie viel Wäsche macht David, wenn Sarah 400 Pfund macht?"
• Versuch 1: Der Freund rechnet es aus und sagt: „100 Pfund."
• Versuch 2: Er denkt kurz nach, rechnet es genau gleich aus und sagt wieder: „100 Pfund."
• Versuch 3: Er sagt es noch einmal: „100 Pfund."

Das Problem ist: Wenn der Freund dreimal hintereinander exakt denselben Weg geht, hast du eigentlich nur einen Versuch gemacht, nicht drei. Wenn er sich bei diesem einen Weg vertippt, hast du keine Chance, das Rätsel zu lösen.

Bisherige KI-Training-Methoden haben genau das gefördert: Sie haben die KI darauf trainiert, den „sichersten" Weg zu finden. Das ist gut für die Genauigkeit beim ersten Versuch, aber es tötet die Kreativität. Die KI wird wie ein Roboter, der immer denselben Pfad entlangläuft, statt verschiedene Wege auszuprobieren.

Die Lösung: UpSkill – „Die Magie der versteckten Karten"

Die Forscher von der Princeton University haben eine Methode namens UpSkill entwickelt. Stell dir das wie folgt vor:

Statt die KI einfach zu fragen, geben wir ihr vor jedem Versuch eine geheime Karte (im Papier eine „latente Variable" namens z).

• Karte A (z=1): „Löse das Problem mit einer algebraischen Formel."
• Karte B (z=2): „Löse das Problem, indem du eine Geschichte erzählst."
• Karte C (z=3): „Löse das Problem, indem du eine Tabelle zeichnest."

Das Ziel von UpSkill ist es, die KI so zu trainieren, dass sie auf jede dieser Karten unterschiedlich reagiert.

Wie funktioniert das Training? (Das „Belohnungsspiel")

Normalerweise bekommt eine KI nur einen Punkt, wenn die Antwort richtig ist. UpSkill gibt ihr aber einen zweiten Bonus-Punkt:

1. Der Richtigkeits-Punkt: Die Antwort muss stimmen (z. B. 100 Pfund).
2. Der Unterschieds-Punkt: Die KI bekommt extra Punkte, wenn ihre Antwort auf Karte A anders aussieht als auf Karte B.

Die KI lernt also: „Hey, wenn ich Karte A ziehe, muss ich einen anderen Weg gehen als bei Karte B, sonst kriege ich keine Bonuspunkte!"

Das Ergebnis ist, dass die KI eine ganze Werkzeugkiste voller verschiedener Lösungsstrategien entwickelt. Sie lernt nicht nur eine Art zu rechnen, sondern viele verschiedene.

Warum ist das super? (Der „Lotterie-Effekt")

Stell dir vor, du musst ein Schloss öffnen.

• Ohne UpSkill: Du hast 5 Schlüssel, aber alle sehen gleich aus und passen nur in eine Art Schloss. Wenn das Schloss kaputt ist, hast du Pech.
• Mit UpSkill: Du hast 5 Schlüssel, die alle unterschiedlich geformt sind (einer ist lang, einer dick, einer gebogen). Selbst wenn der erste Schlüssel nicht passt, hast du eine viel höhere Chance, dass einer der anderen Schlüssel das Schloss öffnet.

In der Welt der KI bedeutet das: Wenn man die KI 5 Mal fragt (mit 5 verschiedenen Karten), ist die Wahrscheinlichkeit, dass mindestens eine Antwort richtig ist, viel höher, weil die Antworten so unterschiedlich sind.

Was haben die Forscher herausgefunden?

Sie haben die Methode an drei verschiedenen KI-Modellen getestet (Llama, Qwen und R1).

• Das Ergebnis: Bei den stärkeren Modellen (Llama und Qwen) hat sich die Erfolgsrate bei mehreren Versuchen (pass@k) um etwa 3 % bis 10 % verbessert, ohne dass die Genauigkeit beim ersten Versuch schlechter wurde.
• Die Theorie: Die Forscher haben mathematisch bewiesen, dass je mehr „Unterschied" (gegenseitige Information) zwischen den verschiedenen Karten und den Antworten besteht, desto besser die Chancen sind, dass die KI das Problem löst.

Zusammenfassung in einem Satz

UpSkill ist wie ein Trainer, der einer KI nicht nur sagt: „Mach es richtig!", sondern auch: „Versuche es auf 5 verschiedene, völlig unterschiedliche Arten!", damit sie nicht in einer einzigen Denkweise stecken bleibt und bei schwierigen Aufgaben mehr Chancen hat, zu gewinnen.

Technische Zusammenfassung

1. Problemstellung

Große Sprachmodelle (LLMs) haben durch Reinforcement Learning mit verifizierbaren Belohnungen (RLVR) erhebliche Fortschritte bei Aufgaben wie Mathematik und Programmierung gemacht. Ein zentrales Problem besteht jedoch darin, dass Standardansätze, die die Genauigkeit bei einem einzigen Versuch ($pass@1$) optimieren, unbeabsichtigt die Vielfalt der Antworten über wiederholte Versuche unterdrücken.

• Redundanz: Wenn ein Modell mehrfach denselben Prompt erhält, erzeugt es oft fast identische Ausgaben. Dies ist in Szenarien schädlich, bei denen nur ein korrekter Versuch aus einer Menge von $k$ Versuchen benötigt wird (z. B. Code-Generierung mit Tests oder formale Beweise).
• Effektive Versuche: Hohe Ähnlichkeit zwischen den Ausgaben reduziert die „effektive Anzahl" unabhängiger Versuche. Selbst wenn die Genauigkeit bei einem Versuch ($pass@1$) steigt, kann die Wahrscheinlichkeit, dass mindestens einer von $k$ Versuchen korrekt ist ($pass@k$), stagnieren oder sinken, da die Exploration neuer Lösungsstrategien fehlt.
• Lücke: Bisherige Methoden zur Erhöhung der Vielfalt (z. B. Temperatur-Sampling, Prompt-Perturbation) sind oft manuell zu justieren, instabil oder führen zu qualitativ unterschiedlichen, aber inkohärenten Lösungen. Es fehlt an einem Trainingsmechanismus, der strukturierte, reproduzierbare und semantisch verschiedene Strategien lernt, ohne die Einzelversuchs-Accuracy zu beeinträchtigen.

2. Methodik: UpSkill

Die Autoren stellen UpSkill vor, eine Trainingsmethode, die das Konzept des Mutual Information Skill Learning (MISL) auf LLMs überträgt, um die $pass@k$-Leistung zu optimieren.

Kernidee:
Das Modell wird trainiert, um seine Antwort auf eine diskrete latente Variable $z \in \{1, \dots, N\}$ zu konditionieren. Jede $z$-Variable soll einer spezifischen, reproduzierbaren Lösungsstrategie entsprechen.

Technische Umsetzung:

1. Token-Level Mutual Information (MI) Reward:
   Das Ziel ist die Maximierung der bedingten gegenseitigen Information $I(\tau; z | x)$ zwischen dem Trajektorium (der Antwort) $\tau$ und der Strategie $z$ gegeben den Prompt $x$.
   Dies wird durch einen neuen Belohnungsterm im Rahmen von Group Relative Policy Optimization (GRPO) implementiert:
   $$r_{TMI}(\tau_i; x, z) = \sum_{t=1}^{|\tau_i|} \left[ \log p_\pi(y_t | x, z, y_{<t}) - \log p_\pi(y_t | x, y_{<t}) \right]$$
   Der erste Term misst die Wahrscheinlichkeit des Tokens unter der spezifischen Strategie $z$, der zweite Term unter einer gemischten Verteilung aller Strategien. Dies belohnt das Modell dafür, dass seine Ausgabe spezifisch für die gewählte $z$ ist (niedrige bedingte Entropie), während die Gesamtvielfalt der Ausgaben hoch bleibt (hohe marginale Entropie).

2. Trainingsziel:
   Die Gesamtreward-Funktion kombiniert:

   • Verifizierbare Korrektheit ($r_{corr}$): Binärer Reward für richtige Antworten.
   • KL-Regularisierung: Strafe für zu große Abweichung vom Basis-Modell.
   • MI-Reward ($r_{TMI}$): Fördert die Diversität der Strategien.
     $$r(\tau_i) = r_{corr}(\tau_i) - \beta \Delta_{KL}(\tau_i) + \alpha_1 r_{TMI}(\tau_i)$$

3. Inferenz:
   Bei der Inferenz werden $k$ verschiedene Werte für $z$ ausgewählt, und für jeden Wert wird eine Antwort generiert. Da jede $z$-Strategie eine andere Lösungsmethode repräsentiert, erhöht sich die Wahrscheinlichkeit, dass mindestens eine Antwort korrekt ist.

3. Theoretische Grundlage

Die Autoren beweisen einen theoretischen Zusammenhang zwischen der Verbesserung von $pass@k$ und der gegenseitigen Information $I(\tau; z | x)$.

• Sie zeigen, dass die Verbesserung von $pass@k$ durch das MI-Objektiv nach unten beschränkt ist.
• Unter bestimmten Annahmen (z. B. dass die gemischte Verteilung der Strategien der Basisverteilung entspricht und Strategien unterschiedliche Erfolgswahrscheinlichkeiten haben) gilt: Eine Maximierung von $I(\tau; z | x)$ führt garantiert zu einer Steigerung von $pass@k$ gegenüber dem Basis-Modell.

4. Ergebnisse

Die Methode wurde auf dem GSM8K-Datensatz (Grundschul-Mathematik) und in einer kontrollierten arithmetischen Umgebung getestet.

• GSM8K (Open-Weight Modelle):

  • Qwen 2.5-7B & Llama 3.1-8B: UpSkill erzielte signifikante Verbesserungen bei $pass@k$ (durchschnittlich ca. +3,4% für Qwen und +3% für Llama) und bei $plurality@k$, ohne die $pass@1$-Genauigkeit zu verschlechtern.
  • R1-Distilled-Qwen2.5-Math-1.5B: Bei diesem kleineren, stark vortrainierten Modell führte UpSkill zu einem Leistungsabfall, was auf die geringere Kapazität und die Sensitivität gegenüber dem MI-Objektiv hinweist.
  • Ohne Ground-Truth: Interessanterweise konnte UpSkill auch ohne explizite Korrektheits-Rewards (nur mit MI-Reward) die $pass@k$-Werte verbessern, was zeigt, dass das Lernen diverser Strategien auch ohne direkte Feedback-Signale für die Richtigkeit funktioniert.

• Arithmetische Umgebung:

  • In einer kontrollierten Umgebung zeigte sich, dass UpSkill verhindert, dass das Modell in eine deterministische Strategie kollabiert (was bei reinem GRPO ohne MI-Reward passiert).
  • Verschiedene $z$-Werte entwickelten sich zu unterschiedlichen Operatoren-Strategien (z. B. Division vs. Multiplikation), was die $pass@5$-Werte von 0,793 auf 0,897 steigerte, während $pass@1$ moderat blieb.

• Ablationsstudien:

  • Die Ergebnisse zeigen, dass die Verbesserung von $pass@k$ eng mit der Höhe der gegenseitigen Information korreliert.
  • Eine zu starke Regularisierung (KL) oder falsche Hyperparameter können die Diversität wieder unterdrücken.

5. Bedeutung und Beiträge

• Strukturierte Vielfalt: UpSkill bietet einen Trainingsansatz, der nicht nur zufällige Variationen erzeugt, sondern reproduzierbare, semantisch unterschiedliche Denkstrategien lernt.
• Effizienz bei Multi-Attempt: Es löst das Problem der redundanten Versuche in verifizierbaren Aufgaben, indem es die effektive Anzahl unabhängiger Versuche erhöht, ohne die Qualität einzelner Versuche zu opfern.
• Theoretische Verbindung: Der Paper stellt einen wichtigen theoretischen Link her, der zeigt, dass die Maximierung von Mutual Information eine untere Schranke für die Verbesserung von $pass@k$ darstellt.
• Praktische Anwendbarkeit: Die Methode ist einfach zu implementieren (als Zusatzreward zu GRPO) und funktioniert mit bestehenden Open-Weight-Modellen, was sie für Anwendungen wie Code-Generierung, formale Verifikation und komplexe Problemlösung attraktiv macht.

Zusammenfassend demonstriert UpSkill, dass die gezielte Förderung von Diversität durch Mutual Information während des Trainings ein effektiver Weg ist, um die Zuverlässigkeit von LLMs in Szenarien zu erhöhen, bei denen mehrere Versuche erlaubt sind.