Complementing reinforcement learning with SFT through logit averaging in the post training of LLMs

Technische Samenvatting: Aanvulling van Versterkingsleren met SFT via Logit-gemiddelde

1. Probleemstelling

De post-training van Large Language Models (LLM's) berust doorgaans op twee paradigma's: Supervised Fine-Tuning (SFT) en Versterkingsleren (RL) met verifieerbare beloningen (RLVR). Hoewel SFT dichte supervisie per token biedt om opmaak en stijl te leren, wordt het beperkt door de verdeling van trainingsdemonstraties. Omgekeerd kunnen RL-methoden zoals Proximal Policy Optimization (PPO) en Group Relative Policy Optimization (GRPO) deze distributielimieten overstijgen, maar lijden ze onder schaarse, ruisende beloningssignalen.

De standaardbenadering combineert deze sequentieel: eerst wordt SFT toegepast, gevolgd door RL met een Kullback-Leibler (KL)-regularisatieterm, $\beta \text{KL}(\pi_\theta \| \pi_{\text{sft}})$, om het trainabele beleid $\pi_\theta$ te verankeren aan de SFT-referentie $\pi_{\text{sft}}$. De auteurs identificeren een fundamentele spanning in dit "KL-verankerde" recept: de KL-boete trekt $\pi_\theta$ naar $\pi_{\text{sft}}$, waardoor het beleid mogelijk wordt beperkt tot de beperkingen van het SFT-checkpoint (bijvoorbeeld algebraïsche fouten, herhalingslussen) en het niet volledig kan profiteren van de redeneercapaciteiten die tijdens RL zijn ontwikkeld. Beoefenaars mitigeren dit vaak door een kleine $\beta$ te kiezen, maar de auteurs stellen een structureel alternatief voor regularisatie voor.

2. Methodologie

Het artikel introduceert een methode die KL-regularisatie vervangt door logit-gemiddelde. In plaats van divergentie te bestraffen, worden het trainabele beleid $\pi_\theta$ en een bevroren referentiebeleid $\pi_{\text{ref}}$ (bijvoorbeeld SFT) gekoppeld door hun logits te middelen om een gemengd beleid $\pi^{\text{mix}}_{\theta, \alpha}$ te vormen.

Kernmechanisme:
Gegeven een gewicht $\alpha \in [0, 1]$ en toestand $s$, worden de gemengde logits $z^{\text{mix}}_{\theta, \alpha}(s)$ gedefinieerd als:
$$z^{\text{mix}}_{\theta, \alpha}(s) = (1 - \alpha) z_\theta(s) + \alpha z_{\text{ref}}(s)$$
Het gemengde beleid is vervolgens $\pi^{\text{mix}}_{\theta, \alpha}(a | s) = \text{softmax}(z^{\text{mix}}_{\theta, \alpha}(s))$.

Trainingsdoel:
De auteurs passen het GRPO-doel aan op dit gemengde beleid:

• Rollouts: Trajecten worden bemonsterd vanuit het gemengde beleid $\pi^{\text{mix}}_{\theta, \alpha}$ in plaats van het pure trainabele beleid $\pi_\theta$.
• Voorsprongsschatting: Groepsrelatieve voordelen worden berekend op basis van beloningen van deze gemengde rollouts.
• Optimalisatie: De trainbare parameters $\theta$ worden bijgewerkt om de beloning onder het gemengde beleid te maximaliseren. De referentielogits $z_{\text{ref}}$ blijven gedurende het hele proces bevroren.
• Weegschema's:
  1. Vaste Gewichten: $\alpha$ blijft constant (bijvoorbeeld $\alpha=0.5$) gedurende de training.
  2. Adaptieve Gewichten: $\alpha$ wordt dynamisch bijgewerkt op basis van een terughoudende validatiestandaard. De update-regel gebruikt een "netto-winst"-metriek die de juistheid van het pure beleid vergelijkt met die van het gemengde beleid, en duwt $\alpha$ naar 1 als mixen een netto-winst oplevert, en anders naar 0.

Vergelijking met RLVR:
In tegenstelling tot standaard RLVR-methoden (PPO, GRPO, DAPO) die een criticus en KL-regularisatie gebruiken, heeft deze benadering geen KL-term of criticus. Het referentiebeleid participeert in verkenning via logit-gemiddelde, en de belangverhoudingen in de verliesfunctie worden berekend tussen het huidige en het oude gemengde beleid.

3. Belangrijkste Bijdragen en Bevindingen

Complementaire Vaardigheden (Patroon 1 & 2):
Empirische analyse over de MATH-, cn-k12- en MMLU-datasets onthult dat SFT en het RL-getrainde beleid complementaire vaardigheden bezitten:

• SFT blinkt uit in opmaak (het naleven van outputsjablonen) maar kan falen bij complex redeneren.
• RL-beleid ($\pi_\theta$) blinkt uit in redeneren (het afleiden van correcte numerieke antwoorden) maar verliest vaak de discipline in opmaak.
• Complementariteit: Een aanzienlijk deel van de problemen (22%–38%) wordt door precies één van de twee beleidslijnen correct opgelost. Logit-gemiddelde fungeert als een "Product of Experts", waarbij de opmaaksterkte van SFT wordt gecombineerd met de redeneersterkte van het RL-beleid.

Prestatieresultaten:
Geëvalueerd op Qwen2.5-Instruct-modellen (1.5B, 3B, 7B) over drie datasets:

• Nauwkeurigheid: De methode met logit-gemiddelde bereikt een hogere of vergelijkbare antwoordnauwkeurigheid ten opzichte van het canonieke KL-geregulariseerde GRPO in alle negen (dataset $\times$ modelgrootte) configuraties.
• Ablatie: Het vervangen van de SFT-anker door een basis Instruct-model resulteert in aanzienlijk lagere prestaties, wat bevestigt dat de winst voortkomt uit het mixen met een formatbewuste referentie.
• Logit versus Kansgemiddelde: Logit-gemiddelde (geometrisch gemiddelde van verdelingen) presteert beter dan kansgemiddelde (aritmatisch gemiddelde). Logit-gemiddelde bereikt hogere oplossingspercentages met kleinere menggewichten en piekt bij $\alpha=0.5$, terwijl kansgemiddelde hogere gewichten vereist en piekt bij $\alpha=0.6$.

Theoretische Connectie:
Het artikel stelt vast dat logit-gemiddelde wiskundelijk equivalent is aan een Product of Experts (POE). Het resulterende beleid minimaliseert een gewogen som van KL-divergenties naar de RL- en referentiebeleidslijnen. Deze structuur stelt het beleid in staat om waarschijnlijkheid te concentreren op tokens die door beide experts worden bevoordeeld, waardoor algebraïsche fouten van SFT effectief worden gecorrigeerd terwijl de opmaakbeperkingen behouden blijven.

4. Betekenis en Claims

De auteurs beweren dat deze methode een duidelijk voordeel biedt ten opzichte van het standaard KL-geregulariseerde recept door de "opmaakanker" te ontkoppelen van de "gradiëntbeperking".

• Overcoming Limitations: Waar KL-regularisatie $\pi_\theta$ dwingt dicht bij $\pi_{\text{sft}}$ te blijven (en mogelijk zijn redeneerfouten erft), stelt logit-gemiddelde $\pi_\theta$ in staat de beperkingen van de referentie te overwinnen terwijl het opmaakvoordelen behoudt via de structurele koppeling van het gemengde beleid.
• Impliciete Vertrouwensgebied: Het mixen fungeert als een impliciet vertrouwensgebied, wat distributieverschuiving weg van opmaakstijlen voorkomt, wat pure RL op grote schaal vaak parten speelt.
• Dichte Priors: In tegenstelling tot RL, dat signalen alleen doorgeeft vanaf het laatste beloningstoken, integreert logit-gemiddelde een dichte, token-niveau prior van de SFT-anker bij elke decodestap.
• Schaalbaarheid: De benadering breidt zich natuurlijk uit tot ensembles van Small Language Models (SLM's), waardoor een systeem gespecialiseerde experts kan benutten (bijvoorbeeld voor code, veiligheid of domeinterminologie) zonder de kosten van het trainen van één enorm beleid.

Het artikel concludeert dat logit-gemiddelde een robuust alternatief biedt voor KL-regularisatie, waarbij de complementariteit tussen SFT en RL succesvol wordt benut om de nauwkeurigheid van het redeneren te verbeteren zonder in te leveren op de kwaliteit van de opmaak.