GTS: Inference-Time Scaling of Latent Reasoning with a Learnable Gaussian Thought Sampler

1. Problemstellung

Inferenzzeit-Skalierung (Inference-Time Scaling, ITS) ist ein zentraler Ansatz zur Verbesserung der reasoning-Fähigkeiten von Large Language Models (LLMs), indem zusätzliche Rechenleistung genutzt wird, um mehrere Denkpfade (Trajektorien) zu generieren und die beste Antwort auszuwählen.

• Diskrete Modelle: Bei herkömmlichen token-basierten Modellen erfolgt dies durch Sampling aus expliziten Wahrscheinlichkeitsverteilungen (z. B. Temperature Scaling).
• Latente Reasoning-Modelle: Neuere Modelle führen das Reasoning direkt im kontinuierlichen Raum der versteckten Zustände (Hidden States) durch, ohne Zwischen-Tokens zu generieren. Dies erhöht die Effizienz, entfernt aber die expliziten token-level Verteilungen.
• Das Kernproblem: Der aktuelle Stand der Technik für ITS in latenten Räumen stützt sich auf heuristische Perturbationen (Störungen), wie z. B. Dropout oder festes Gaußsches Rauschen. Die Autoren zeigen, dass diese Methoden zwei fundamentale Mängel haben:
  1. Sie definieren keine explizite bedingte Sampling-Verteilung.
  2. Eine stärkere Perturbation führt nicht automatisch zu besserer Qualität. Oft führt zu viel Rauschen zu einer übermäßigen Verteilungsverschiebung (Distribution Shift), die den Modellzustand von den für die korrekte Antwort relevanten Regionen wegführt („Over-Exploration"), während zu wenig Rauschen keine ausreichende Vielfalt bietet („Under-Exploration").

2. Methodik: Gaussian Thought Sampler (GTS)

Um diese Limitierungen zu überwinden, schlagen die Autoren den Gaussian Thought Sampler (GTS) vor. Dies ist ein leichter, lernbarer Modul, das die latente Exploration als bedingtes Sampling aus einer expliziten Wahrscheinlichkeitsverteilung reformuliert.

• Architektur: GTS wird an ein eingefrorenes (frozen) latentes Reasoning-Backbone angehängt. Anstatt festes Rauschen hinzuzufügen, sagt GTS für jeden Reasoning-Schritt $k$ eine kontextabhängige Gaußsche Verteilung über die Perturbationen vorher.
  • Der Perturbationsvektor $z_k$ wird aus einer Verteilung $q_\phi(z_k | c_k) = \mathcal{N}(\mu_\phi(c_k), \text{diag}(\sigma^2_\phi(c_k)))$ gesampelt, wobei $c_k$ der Kontext (der deterministische Hidden State) ist.
  • Der gestörte Zustand ist $\tilde{h}_k = h^{det}_k + z_k$.
• Training (Policy Optimization): Da die Endleistung (Accuracy) nicht differenzierbar bezüglich der Sampler-Parameter ist, wird GTS mit Reinforcement Learning (RL) trainiert, spezifisch mit einer GRPO-ähnlichen (Group Relative Policy Optimization) Strategie.
  • Reward-Design: Für jede Eingabe werden $N$ Trajektorien gesampelt. Der Reward basiert auf der Korrektheit der Antwort (symmetrisch: +1 für korrekt, -1 für falsch) und einem zusätzlichen „Shaping"-Term, der die Konfidenz der Antwort berücksichtigt. Dies fördert Trajektorien, die nicht nur korrekt, sondern auch selbstbewusst sind.
  • Optimierung: Die Policy wird so aktualisiert, dass sie Perturbationen bevorzugt, die die Erfolgschance bei der Inferenzzeit-Skalierung erhöhen, unter Beibehaltung einer Referenz-Policy (via KL-Strafe), um ein Kollabieren zu deterministischen Pfaden zu verhindern.

3. Wichtige Beiträge

1. Identifikation einer Schlüssel-Limitierung: Die Autoren zeigen analytisch, dass effektives Scaling im latenten Raum nicht einfach mehr Stochastizität erfordert, sondern eine kontrollierte Sampling-Qualität. Heuristische Methoden scheitern oft daran, den Trade-off zwischen Vielfalt und Entscheidungsverbesserung zu steuern.
2. Reformulierung als bedingtes Sampling: GTS wandelt die heuristische Perturbation in eine explizite, optimierbare probabilistische Policy um. Dies bringt latentes ITS näher an das theoretische Fundament des token-basierten Samplings.
3. Empirische Validierung: Durch Experimente auf zwei verschiedenen latenten Reasoning-Architekturen (COCONUT und CODI) und mehreren Benchmarks wird gezeigt, dass GTS heuristische Baselines (Dropout, festes Rauschen) konsistent übertrifft.

4. Ergebnisse

Die Experimente wurden auf mathematischen Reasoning-Datensätzen (GSM8K, MultiArith, SVAMP, GSM8K-Hard) durchgeführt.

• Skalierungsverhalten: GTS zeigt eine stärkere und zuverlässigere Skalierung (gemessen an pass@N) als heuristische Baselines, insbesondere bei moderaten bis großen Sampling-Budgets ($N$). Während heuristische Methoden bei starkem Rauschen ($p=0.5$) oft sogar schlechter abschneiden als deterministische Inferenz, bleibt GTS stabil und verbessert die Leistung.
• Sampling-Qualität: Metriken wie „Sampling Gain" (Verbesserung der Wahrscheinlichkeit der korrekten Antwort) zeigen, dass GTS die Verteilung der Antworten gezielt in Richtung korrekter Lösungen verschiebt, ohne unnötig große Divergenzen (Distribution Shift) zu erzeugen. Heuristische Methoden neigen hingegen dazu, entweder zu wenig zu explorieren oder die Entscheidungsrelevanz zu zerstören.
• Generalisierung: GTS generalisiert gut auf Out-of-Distribution-Datensätze, was darauf hindeutet, dass die gelernte Perturbations-Policy robust gegenüber unterschiedlichen Problemtypen ist, während heuristische Methoden stark von der spezifischen Einstellung abhängen.
• Verhaltensanalyse: Die Analyse des Signal-zu-Rausch-Verhältnisses (SNR) zeigt, dass GTS lernt, den Grad der Stochastizität dynamisch an die internen Reasoning-Dynamiken des jeweiligen Backbones anzupassen (z. B. saw-tooth-Muster bei COCONUT vs. front-loaded bei CODI).

5. Bedeutung und Fazit

Die Arbeit demonstriert, dass das bloße Hinzufügen von Rauschen für latentes Reasoning nicht ausreicht. Stattdessen ist ein kontrollierter, lernbarer Sampling-Mechanismus notwendig, der den Kontext berücksichtigt und gezielt nach Pfaden sucht, die die Antwortqualität verbessern.

GTS bietet einen prinzipiellen Ansatz, der die Lücke zwischen diskretem Token-Sampling und kontinuierlichem latentem Reasoning schließt. Die Ergebnisse legen nahe, dass die Zukunft des Inference-Time Scaling in latenten Räumen nicht in der bloßen Amplifikation von Stochastizität liegt, sondern in der Entwicklung optimierbarer, kontextsensitiver Sampling-Policies, die die Exploration effizient steuern. Dies ist ein wichtiger Schritt hin zu robusteren und effizienteren Reasoning-Systemen, die Rechenleistung gezielt nutzen, um bessere Entscheidungen zu treffen.