ReLaX: Reasoning with Latent Exploration for Large Reasoning Models

1. Probleemstelling

De paper adresseert een fundamentele beperking in het trainen van Grote Redeneringsmodellen (LRMs) via Versterkend Leren met Verifieerbare Beloningen (RLVR). Hoewel RLVR de redeneercapaciteit van modellen (zoals LLMs en Multimodale LLMs) heeft verbeterd, leidt het vaak tot een over-determinisme en een informatie-entropie-instorting (entropy collapse).

• De Kern van het Probleem: Zonder expliciete interventie concentreert de RLVR-beleidspolicy (policy) zich snel op een smal subspace van mogelijke antwoorden. Dit resulteert in een vroege convergentie waarbij het model stopt met exploreren en vastloopt in starre, repetitieve patronen.
• Huidige Oplossingen en Hun Tekortkomingen: Bestaande methoden proberen dit te bestrijden door de token-level entropie te verhogen (bijv. door beloningen te herschappen of entropie-regulering toe te passen). De auteurs betogen echter dat dit een oppervlakkige oplossing is:
  1. Het verhogen van token-entropie botst vaak met de natuurlijke neiging van RL om deterministische, lage-entropie beleidsplannen te vinden.
  2. Bij multimodale modellen (MLLMs) reflecteert token-level feedback de onderliggende, complexe cross-modale berekeningen niet nauwkeurig genoeg.
  3. De echte oorzaak van de prestatie-daling ligt in de latent dynamics (de dynamiek van de verborgen staten), die steeds rigider worden, in plaats van alleen in de uitkomst van de tokens.

2. Methodologie: ReLaX

De auteurs introduceren ReLaX (Reasoning with Latent eXploration), een framework dat de exploratie-exploitatie trade-off reguleert door in te grijpen op de latente dynamiek van het model, in plaats van alleen op de token-uitvoer.

A. Theoretische Basis: Koopman Operator Theorie

Om de niet-lineaire, hoge-dimensionale dynamiek van de verborgen staten van een LRM analyseerbaar te maken, maken de auteurs gebruik van de Koopman-operator theorie.

• Deze theorie stelt dat niet-lineaire dynamische systemen kunnen worden weergegeven als lineaire evolutie in een oneindig-dimensionale ruimte van waarneembare grootheden (observables).
• Ze gebruiken een ResKoopNet (een neurale Koopman-dictionary) om een lineaire representatie te leren van de verborgen staten van het model. Dit maakt het mogelijk om de complexiteit van de interne berekeningen te kwantificeren.

B. Nieuwe Metriek: Dynamic Spectral Dispersion (DSD)

Op basis van de lineaire representatie introduceren ze DSD als een nieuwe metriek voor policy-exploratie.

• Definitie: DSD is de variantie van de grootte van de eigenwaarden van de geschatte Koopman-operator ($K$).
• Interpretatie:
  • Een lage DSD (geconcentreerd spectrum) duidt op degeneratief, repetitief gedrag (rigide interne berekeningen).
  • Een hoge DSD (verspreid spectrum) duidt op een rijkere, meer expressieve dynamische structuur, wat betekent dat het model flexibel blijft en effectief kan exploreren.
• DSD fungeert dus als een directe indicator voor de "rekenflexibiliteit" van het model.

C. Het ReLaX Framework

ReLaX integreert DSD in het Group Relative Policy Optimization (GRPO) trainingsproces via een regularisatieterm.

1. DSD Regularisatie: Een term ($L_{xp}$) wordt toegevoegd aan de loss-functie die de DSD-scores van trajecten maximaliseert. Dit straft het model af als het interne dynamische spectrum te homogeen wordt.
2. Advantage Shaping: Om te voorkomen dat het model nutteloos exploreert, wordt de DSD-regularisatie alleen toegepast op trajecten met een positieve voordeel (advantage). Dit zorgt ervoor dat flexibiliteit alleen wordt beloond als het leidt tot betere antwoorden.
3. Adaptieve KL-regularisatie: Om instabiliteit te voorkomen bij overmatige spreiding, wordt een adaptieve KL-straf toegepast op trajecten met extreme DSD-waarden.

De totale doelstelling is:
$$J(\theta) = J_{GRPO}(\theta) + \alpha \cdot \tilde{L}_{xp} + \beta \cdot \sum DKL(\pi_\theta || \pi_{ref})$$

3. Belangrijkste Bijdragen

1. DSD Metriek: De introductie van Dynamic Spectral Dispersion, een fundamentele maatstaf voor de heterogeniteit van latente redeneerdynamiek, die dieper gaat dan oppervlakkige token-statistieken.
2. ReLaX Framework: Een nieuw trainingsparadigma dat expliciet de interne berekeningsflexibiliteit reguleert om de exploratie-exploitatie trade-off te optimaliseren en prestatie-saturatie te voorkomen.
3. State-of-the-Art Resultaten: Het bewijzen dat het sturen van exploratie via de latente ruimte superieur is aan token-level methoden, met name bij multimodale taken.

4. Resultaten

De auteurs hebben uitgebreide experimenten uitgevoerd op zowel tekst-only LLMs als multimodale VLMs (Vision-Language Models) in maten van 3B en 7B parameters.

• Multimodale Redenering (VLMs):

  • ReLaX-VL-7B bereikte een nieuw state-of-the-art (SOTA) gemiddelde score van 53.2 op 7 multimodale benchmarks (zoals MMMU, MathVista, MathVerse), wat een verbetering is ten opzichte van de vorige beste (VL-Rethinker-7B met 52.5).
  • Zelfs het 3B-model van ReLaX presteerde beter dan veel bestaande 7B-modellen.
  • Training dynamics tonen aan dat ReLaX de entropie en DSD op een hoger, gestabiliseerd niveau houdt, terwijl standaard GRPO snel instort.

• Tekst-only Wiskundige Redenering (LLMs):

  • Op de MATH500, AMC en AIME benchmarks overtrof ReLaX bestaande RLVR-methoden (zoals SimpleRL, DAPO, FR3E).
  • Op het Qwen2.5-7B-Base model verbeterde ReLaX de gemiddelde score met 4.3 punten ten opzichte van de vorige SOTA (FR3E).
  • Op het Qwen2.5-7B-Math model werd een verbetering van 6.3 punten behaald.

• Ablatie Studies:

  • Het verwijderen van de DSD-regularisatie leidt tot een snelle daling in prestaties.
  • Het verwijderen van de "Advantage Shaping" (het filteren op positieve beloningen) leidt tot instabiel trainen en slechtere prestaties dan het basismodel, wat aantoont dat exploratie gericht moet zijn.

5. Betekenis en Conclusie

De paper toont aan dat het oplossen van het "exploration-exploitation dilemma" in RLVR niet alleen gaat om het diverser maken van de output-tokens, maar om het behoud van flexibiliteit in de interne berekeningsdynamiek.

• Paradigmaverschuiving: ReLaX beweegt de focus van "token-level manipulatie" naar "latente ruimte-sturing". Dit is vooral cruciaal voor multimodale modellen, waar de relatie tussen visuele input en tekstuele output complex is en niet goed wordt gevangen door token-entropie alleen.
• Toekomstperspectief: De resultaten suggereren dat het gebruik van wiskundige theorieën (zoals Koopman-operatoren) om de interne dynamiek van neurale netwerken te analyseren en te reguleren, een veelbelovende richting is voor het ontwikkelen van robuustere en schaalbare AI-systemen met betere redeneercapaciteiten.

Kortom, ReLaX biedt een principieel en effectief middel om de "mode collapse" in RLVR te voorkomen, waardoor modellen langer kunnen leren en complexere redeneertaken kunnen oplossen.