Proximal Supervised Fine-Tuning

Probleemstelling

Supervised Fine-Tuning (SFT) is een cruciale stap in het post-training proces van foundation modellen, vaak gebruikt om modellen te distilleren van reasoning-vaardigheden die zijn geleerd via Reinforcement Learning (RL). Hoewel SFT efficiënt en eenvoudig is, lijdt het aan twee fundamentele beperkingen:

1. Slechte Generalisatie: SFT werkt vaak als "behavior cloning" (gedrag kopiëren). Als het fine-tuning dataset suboptimaal is of niet goed overeenkomt met de pre-training distributie, kan dit leiden tot grote veranderingen in het beleid (policy drift). Dit resulteert in een model dat goed presteert op de trainingsdata maar slecht generaliseert naar out-of-domain taken.
2. Entropie-instorting (Entropy Collapse): Tijdens standaard SFT neigt de output-distributie van het model om te instorten (over-concentratie op specifieke tokens). Dit vermindert de diversiteit van de output en beperkt de exploratiecapaciteit van het model. Dit maakt het moeilijk om het model later effectief te optimaliseren met RL-algoritmen zoals PPO of GRPO, omdat het model te "stijf" is geworden.

Methodologie: Proximal Supervised Fine-Tuning (PSFT)

De auteurs inspireren zich op Trust Region Policy Optimization (TRPO) en Proximal Policy Optimization (PPO) uit het domein van Reinforcement Learning en passen deze concepten toe op het SFT-proces.

• Theoretische Basis: SFT wordt gezien als een speciaal geval van policy gradient waarbij de "advantage" ($\hat{A}_t$) voor alle ground-truth tokens constant 1 is.
• De PSFT-Doelfunctie: In plaats van de standaard cross-entropy loss te minimaliseren, introduceert PSFT een geclipte surrogate doelfunctie (gebaseerd op PPO) die een "trust region" creëert rondom het oude beleid.
  De loss functie wordt gedefinieerd als:
  $$L_{PSFT}(\theta) = \mathbb{E}_{(s_t, a_t) \sim D} \left[ \min \left( r_t(\theta), \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \right) \right]$$
  Waarbij $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$ de ratio is tussen de nieuwe en de oude policy.
• Dynamische Update van $\pi_{\theta_{old}}$: Een belangrijk onderscheid met traditionele PPO is dat in PSFT het "oude" beleid ($\pi_{\theta_{old}}$) niet statisch is (zoals een referentiemodel), maar dynamisch wordt bijgewerkt tijdens het trainingstraject. Dit zorgt voor een geleidelijke en stabiele aanpassing aan de dataset, in plaats van het beperken van de optimalisatie tot een trust region rond een verouderd referentiepunt.
• Warm-up Fase: Om bias in de beginfase te verminderen (waarbij de oude policy nog niet goed afgestemd is op de dataset), wordt een optionele warm-up SFT-fase voorgesteld om de initiële policy beter te aligneren.

Kernbijdragen

1. Nieuwe Optimalisatiestrategie: PSFT is een innovatieve SFT-methode die trust-region beperkingen gebruikt om overmatige updates van het beleid te voorkomen, zonder de voordelen van RL te vereisen tijdens de SFT-fase.
2. Preventie van Entropie-instorting: Door de ratio tussen nieuwe en oude kansen te begrenzen, voorkomt PSFT dat het model te snel overtuigd raakt van specifieke tokens. Dit behoudt de entropie (diversiteit) van de output, wat essentieel is voor latere RL-fasen.
3. Verbeterde Generalisatie: PSFT behoudt beter de algemene vaardigheden van het model (zoals instructievolging en redeneren in andere domeinen) terwijl het toch prestaties op de doeltaak verbetert.
4. Sterke Basis voor RL: Modellen getraind met PSFT bieden een betere "cold start" voor subsequent RL-training (zoals GRPO of DPO), omdat ze meer ruimte hebben voor exploratie en minder overfitting vertonen.

Resultaten

De auteurs hebben PSFT uitgebreid getest op wiskundige redenering, menselijke waarde-alignement en multimodale taken (met Qwen en Llama modellen).

• In-Domain Prestaties: PSFT bereikt prestaties die vergelijkbaar zijn met, of soms zelfs beter zijn dan, standaard SFT op de doeltaak (bijv. wiskundeproblemen zoals AIME en MATH).
• Out-of-Domain Generalisatie: In tegenstelling tot standaard SFT, dat vaak prestaties op andere taken (zoals IFEval of GPQA) ziet dalen door overfitting, behoudt PSFT deze vaardigheden beter. In sommige gevallen verbetert PSFT zelfs de generalisatie.
• Entropie en Trainingstabiliteit:
  • Standaard SFT toont een scherpe daling in entropie (instorting) na elke epoch.
  • PSFT behoudt een soepelere entropiecurve, wat aangeeft dat het model diversiteit behoudt.
• RL Potentieel: Wanneer PSFT-modellen worden gebruikt als startpunt voor RL (GRPO/DPO), laten ze een snellere en steilere verbetering zien dan modellen die met standaard SFT zijn gestart. Ze bereiken uiteindelijk hogere scores op zowel in-domain als out-of-domain benchmarks.
• Multimodale Toepassing: Bij Vision Language Models (VLMs) voorkomt PSFT dat het toevoegen van tekst-only reasoning-data de visuele redeneercapaciteit verstoort, terwijl standaard SFT hier wel last van heeft.

Betekenis en Impact

Dit paper biedt een fundamentele verbetering voor het post-training ecosysteem van Large Language Models (LLMs).

• Alternatief voor RL: Het toont aan dat de voordelen van trust-region optimalisatie (stabiliteit en generalisatie) al in de SFT-fase kunnen worden toegepast, zonder de complexiteit en kosten van on-policy RL.
• Oplossing voor het "Alignment Tax" Probleem: PSFT vermindert de "alignment tax" (het verlies aan algemene vaardigheden tijdens het aligneren van een model), wat een groot probleem is in de huidige LLM-ontwikkeling.
• Schaalbaarheid: De methode is reproduceerbaar met bestaande frameworks en biedt een robuuste basis voor de volgende generatie van reasoning-modellen, waarbij het de overgang tussen SFT en RL soepeler maakt.

Kortom, PSFT transformeert SFT van een methode die vaak leidt tot overfitting en verlies van flexibiliteit, naar een gestructureerd proces dat zowel specifieke taken leert als de algemene intelligentie van het model behoudt en versterkt.