PowerFlow: Unlocking the Dual Nature of LLMs via Principled Distribution Matching

Probleemstelling

Ongecontroleerd Reinforcement Learning from Internal Feedback (RLIF) is een veelbelovende aanpak om de latente capaciteiten van Large Language Models (LLMs) te activeren zonder externe supervisie of gelabelde data. Echter, bestaande RLIF-methoden lijden onder fundamentele tekortkomingen:

1. Heuristische Beloningen: Ze vertrouwen op handgemaakte, heuristische intrinsieke beloningen (zoals zelfzekerheid of consistentie) zonder een goed gedefinieerd theoretisch optimalisatiedoel.
2. Pathologische Bias: Door het ontbreken van een principieel doel zijn deze methoden vatbaar voor onbedoelde biases, wat leidt tot pathologisch gedrag zoals "mode collapse" (herhaling), oververtrouwen, en vooral lengte-bias.
   • Bij het "scherpen" van de verdeling (voor redenering) neigen modellen naar extreem korte, triviale antwoorden.
   • Bij het "vlakken" (voor creativiteit) neigen ze naar repetitieve, lange sequenties.
3. Dualiteit van LLMs: Er is een spanning tussen het verbeteren van logisch redeneren (vaak vereist door verdelingsscherping) en het behoud van creatieve diversiteit (vaak onderdrukt door alignment-procedures die de verdeling te sterk scherpen). Bestaande methoden kunnen deze twee doelen niet gelijktijdig en gecontroleerd aanpakken.

Methodologie: PowerFlow Framework

PowerFlow lost deze problemen op door ongesuperviseerde fine-tuning te herformuleren als een principieel distributie-matching probleem, in plaats van een zoektocht naar heuristische beloningen.

1. Doelverdeling: De $\alpha$-power Distributie
In plaats van een externe reward-functie te gebruiken, target PowerFlow de $\alpha$-power distributie (ook wel de escort distribution genoemd) van het basismodel:
$$p_\alpha(y|q) \propto p_{base}(y|q)^\alpha$$
Hierbij fungeert de parameter $\alpha$ als een "controleknop" voor de dualiteit van het model:

• $\alpha > 1$ (Scherpen): Verdicht de waarschijnlijkheidsmassa op hoogwaardige redeneerpaden, waardoor logisch redeneren wordt geïntensiveerd.
• $\alpha < 1$ (Vlakken): Verspreidt de waarschijnlijkheidsmassa naar de lange staart, waardoor onderdrukte creatieve expressie wordt vrijgegeven.

2. GFlowNets als Amortized Samplers
Om deze ongenormaliseerde distributie te matchen, gebruikt PowerFlow Generative Flow Networks (GFlowNets). GFlowNets leren een stochastisch beleid dat steekproeven trekt die evenredig zijn met de ongenormaliseerde dichtheid. Dit wordt benaderd via variational inference, waarbij de doelverdeling wordt gematcht door de reverse KL-divergentie te minimaliseren.

3. Length-Aware Trajectory-Balance (LA-TB) Doel
De kerninnovatie is de oplossing voor de structurele lengte-bias inherent aan autoregressieve generatie (waarbij de log-probabiliteit lineair afneemt met de lengte).

• Het Probleem: Standaard distributie-matching objectieven worden gedomineerd door sequentielengte in plaats van semantische dichtheid, wat leidt tot instabiele training en degeneratie.
• De Oplossing: PowerFlow reparametriseert de partitiefunctie $Z$ als een lengte-bewuste energeterm: $Z_\phi(q, y) = (Z'_\phi(q))^{|y|}$.
• Het Objectief: De nieuwe Length-Aware Trajectory-Balance (LA-TB) loss normaliseert de mismatch op token-niveau:
  $$L_{LA-TB} = \left( \log Z'_\phi(q) + \frac{1}{|y|} \log \frac{\pi_\theta(y|q)}{\tilde{p}_{target}(y|q)} \right)^2$$
  Dit zorgt ervoor dat de optimalisatiegradiënt schaal-invariant blijft ten opzichte van de sequentielengte, waardoor het model semantische kwaliteit prioriteert boven lengte of redundantie.

Belangrijkste Bijdragen

1. Principieel Framework: PowerFlow biedt een theoretisch onderbouwde basis voor ongesuperviseerde fine-tuning door het te reduceren tot distributie-matching met een controleerbare parameter ($\alpha$).
2. Length-Aware Objective: De afleiding van de LA-TB loss neutraliseert de exponentiële lengte-bias van autoregressieve modellen, wat stabiele training mogelijk maakt zonder degeneratieve lengte-collaps.
3. Activering van Dualiteit: Het framework demonstreert dat één methode zowel logisch redeneren kan verbeteren ($\alpha > 1$) als creatieve diversiteit kan herstellen ($\alpha < 1$), afhankelijk van de gekozen parameter.

Resultaten

Uitgebreide experimenten op verschillende modellen (Qwen2.5, Llama-3.2) en benchmarks tonen aan:

• Redeneren ($\alpha > 1$):

  • PowerFlow presteert consistent beter dan bestaande RLIF-methoden (zoals Intuitor, EMPO, TTRL).
  • Het bereikt of overtreft zelfs supervised GRPO (Group Relative Policy Optimization) op benchmarks zoals MATH500, OlympiadBench en AIME, zonder gebruik te maken van externe verifiers of gelabelde data.
  • Het behoudt een hoge diversiteit in oplossingsstrategieën, in tegenstelling tot andere methoden die vaak in mode collapse vervallen.

• Creativiteit ($\alpha < 1$):

  • Toegepast op instructie-getrainde modellen, herstelt PowerFlow de creatieve diversiteit die vaak verloren gaat tijdens alignment.
  • Het bereikt een Pareto-improvement: het verbetert zowel de kwaliteit van de output als de semantische diversiteit, terwijl andere methoden (zoals het verhogen van de temperatuur) vaak een afname in kwaliteit zien bij een toename in diversiteit.

Betekenis en Impact

PowerFlow vertegenwoordigt een verschuiving in het paradigma van LLM-optimatie:

• Van Heuristiek naar Principieel: Het vervangt het "tweaken" van beloningsfuncties door een wiskundig onderbouwde aanpak die de intrinsieke structuur van het model respecteert.
• Efficiëntie: Het elimineert de noodzaak voor kostbare, gelabelde datasets of externe reward-modellen voor het activeren van geavanceerde vaardigheden.
• Veiligheid en Controle: Hoewel het de latentie van het model benut, biedt de controleerbare parameter $\alpha$ een mechanisme om de richting van de verbetering (redeneren vs. creativiteit) nauwkeurig te sturen.
• Toekomstperspectief: De methode opent de deur voor een unificatie van ongesuperviseerde alignment, waarbij taakspecifieke vaardigheden worden opgeëist door het model te matchen met geoptimaliseerde doel-geometrieën, in plaats van het trainen van nieuwe kennis.

Kortom, PowerFlow bewijst dat de "dualiteit" van LLMs (logisch vermogen versus creatieve expressie) niet inherent tegenstrijdig is, maar gecontroleerd kan worden via principieel distributie-matching, mits de inherente lengte-bias van autoregressieve generatie correct wordt aangepakt.