Technische Samenvatting: PriFT (Prior-Support Guided Fine-Tuning)

Probleemstelling

Supervised Fine-Tuning (SFT) is een standaard en computationeel efficiënte methode voor het aanpassen van voorgetrainde Large Language Models (LLMs) aan downstream-taken, wat vaak dient als initialisatie voor Reinforcement Learning (RL). SFT lijdt echter aan een generalisatiekloof vergeleken met RL. Als een off-policy methode past SFT vaste demonstratiedata token voor token aan. Wanneer doel-tokens in de demonstraties slecht zijn uitgelijnd met de voorgetrainde distributie van het model, kan dit proces leiden tot overfitting en de verstoring van de voorgetrainde kennis.

Recente benaderingen proberen dit te verzachten door niet-uniforme trainingsgewichten toe te kennen aan tokens op basis van hun uitlijning met de voorspellende distributie van het model (bijv. door gebruik te maken van de doel-token waarschijnlijkheid of entropie). Deze methoden berekenen deze gewichten echter doorgaans met behulp van het online model dat momenteel wordt gefinetuned. Het artikel stelt dat dit een kritiek gebrek creëert: de voorspellende distributie van het online model wijkt tijdens de training snel af van de voorgetrainde distributie. Bijgevolg raakt het wegingssignaal verstrengeld met het optimalisatietraject, wat een zelfversterkende dynamiek induceert waarbij het model de leerprocessen concentreert op tokens die het aanvankelijk al prefereerde, terwijl andere worden onderdrukt. Dit vermindert de steekproefdiversiteit en beperkt het vermogen van het model voor verdere on-policy optimalisatie (RL).

Methodologie: PriFT

Om de instabiliteit van online herweging aan te pakken, stellen de auteurs PriFT (Prior-support guided Fine-Tuning) voor. De kerninzicht is dat token-gewichten afgeleid moeten worden van een bevroren voorgetraind referentiemodel in plaats van van het evoluerende online model. Dit biedt een stabiele schatting van "prior support" — de mate waarin een doel-token wordt ondersteund door de kennis van het model vóór taakspecifieke aanpassing.

Framework

PriFT wijzigt de standaard SFT-doelfunctie door token-gewichten $m_t$ te introduceren die afgeleid zijn van het bevroren voorgetrainde model $\pi_{\theta_{pt}}$:
$$L_{PriFT}(\theta_{on}) = -\sum_{t=1}^{T} u(\pi_{\theta_{pt}}(\cdot | x, y_{<t}), y_t) \log \pi_{\theta_{on}}(y_t | x, y_{<t})$$
waarbij $u(\cdot)$ een token-utiliteitsfunctie is. Door de gewichten te berekenen vanuit het bevroren referentiemodel, ontkoppelt PriFT de gewichtsbepaling van het optimalisatietraject, waardoor de zelfversterkende bias die bij online methoden wordt waargenomen, wordt voorkomen.

Instantiaties

Het artikel introduceert twee specifieke instantiaties van de utiliteitsfunctie $u$:

1. PriFT-prob: Gebruikt direct de voorgetrainde doel-token waarschijnlijkheid.
   $$m_t = \pi_{\theta_{pt}}(y_t | x, y_{<t})$$
   Dit schaalt de gradiëntgrootte op basis van hoe plausibel het token was onder de voorgetrainde distributie. Het behoudt de standaard SFT-update richting maar verlaagt het gewicht van tokens met een lage prior support.

2. PriFT-mass: Gebruikt de cumulatieve waarschijnlijkheidsmassa onder de voorgetrainde distributie om bias naar "makkelijke" tokens met scherpe distributies te verminderen.
   $$m_t = \mathbb{I}[u^{(MASS)}_t \geq 0.5]$$
   waarbij $u^{(MASS)}_t = \sum_{v: \pi_{\theta_{pt}}(v) \leq \pi_{\theta_{pt}}(y_t)} \pi_{\theta_{pt}}(v)$.
   Deze metriek meet het deel van de waarschijnlijkheidsmassa dat is toegewezen aan tokens die minder waarschijnlijk zijn dan de doel-token. Een drempelwaarde van 0.5 selecteert tokens die worden ondersteund door ten minste de helft van de voorgetrainde waarschijnlijkheidsmassa. In tegen tegenstelling tot PriFT-prob resulteert dit in een binaire selectiemasker waarbij alle geselecteerde tokens een gelijk gewicht krijgen, wat voorkomt dat triviale tokens met een hoge zekerheid (bijv. functiewoorden) overmatig worden benadrukt, terwijl informatieve tokens uit moeilijkere contexten behouden blijven.

Computationele Overhead

PriFT vereist één forward pass van het bevroren voorgetrainde model om de token-gewichten of maskers vooraf te berekenen. Dit voegt een relatieve overhead toe van ongeveer $1/(3E)$ (waarbij $E$ het aantal epochs is), wat offline kan worden uitgevoerd. In tegen tegenstelling tot KL-regularisatiemethoden die toegang vereisen tot de referentiedistributie tijdens elke optimalisatiestap, vereist PriFT alleen het opslaan van scalaire gewichten of binaire maskers, waardoor aanzienlijke geheugenoverhead wordt vermeden.

Belangrijkste Bijdragen

1. Identificatie van Signaaldrift: De auteurs demonstreren dat online token-statistieken snel afwijken van de voorgetrainde distributie, wat een zelfversterkende lus creëert die het leren concentreert op aanvankelijk geprefereerde tokens en de generalisatie verslechtert.
2. Stabiele Herweging via Bevroren Referentie: Ze tonen aan dat het afleiden van token-gewichten uit een bevroren voorgetraind model consistent beter presteert dan het gebruik van het online model over diverse herwegingsregels (waarschijnlijkheidsgebaseerd en entropiegebaseerd).
3. PriFT Framework: De voorstel van een eenvoudig, algemeen framework dat gewichten vooraf berekent vanuit een bevroren referentie, waardoor de herwegingssignaal wordt ontkoppeld van het optimalisatietraject.
4. Empirische Validatie: Uitgebreide experimenten tonen aan dat PriFT state-of-the-art resultaten behaalt onder SFT-baselines en een superieure initialisatie biedt voor daaropvolgende RL-training.

Experimentele Resultaten

De auteurs evalueerden PriFT op wiskundige redenering, codegeneratie en medische vraagbeantwoording met modellen zoals Qwen2.5-Math, Qwen3 en LLaMA-2.

• Wiskundige Redenering: Op Qwen2.5-Math-7B behaalde PriFT-mass een gemiddelde nauwkeurigheid (Avg@16) van 40.67 en een Pass@16 van 69.60, waarmee het de sterkste eerdere baseline (TALR) respectievelijk met 3.97 en 8.75 punten versloeg. Opvallend genoeg, terwijl online herwegingsmethoden (zoals DFT) vaak de Avg@16 verbeterden ten koste van de Pass@16, verbeterde PriFT beide metrieken simultaan.
• Codegeneratie: De PriFT-varianten presteerden op Qwen3-4B-Instruct beter dan standaard SFT en andere token-herwegingsbaselines, met name op de uitdagende LiveCodeBench splits. PriFT-mass vertoonde de kleinste prestatieverslechtering op HumanEval+ vergeleken met het originele checkpoint, wat suggereert dat bestaande capaciteiten beter behouden blijven.
• Medische QA: PriFT-mass behaalde de beste algehele prestaties op de MMLU-Medical en MedMCQA benchmarks, waarbij het andere token-herwegingsmethoden versloeg.
• RL Initialisatie: Modellen geïnitialiseerd met PriFT vertoonden een sterkere prestatie na daaropvolgende RL-training (DAPO) vergeleken met modellen geïnitialiseerd met standaard SFT of online herwegingsmethoden (DFT). Bijvoorbeeld, op Qwen3-8B-Base leverde PriFT-mass een verbetering van 9.57 punten in Avg@16 op na RL vergeleken met de baseline SFT-initialisatie.

Betekenis en Claims

Het artikel beweert dat PriFT een fundamentele beperking in token-hergewogen SFT aanpakt: de verstrengeling van het wegingssignaal met het optimalisatietraject. Door het voorgetrainde model te gebruiken als een stabiele bron van "prior support", bereikt PriFT het volgende:

• Behoudt Steekproefdiversiteit: Het voorkomt de zelfversterkende concentratie van de waarschijnlijkheidsmassa op een smalle set redeneerpaden, waardoor bredere entropie-distributies en betere Pass@k-prestaties worden behouden.
• Vermindert Distributieverschuiving: Het minimaliseert de KL-divergentie van het voorgetrainde model tijdens het finetunen in vergelijking met online herwegingsmethoden.
• Verbetert RL-gereedheid: Het biedt een sterkere initialisatie voor RL door het vermogen van het model voor verdere on-policy optimalisatie te behouden, wat suggereert dat "goede SFT" niet alleen gaat over directe prestaties, maar over het behouden van de diversiteit die nodig is voor toekomstig leren.

De auteurs merken beperkingen op, waaronder de focus op middelgrote open modellen en specifieke taakdomeinen, wat ruimte laat voor toekomstig onderzoek naar grotere modellen en bredere utiliteitsfuncties. Ze beweren niet dat PriFT alle generalisatieproblemen oplost, maar positioneren het als een robuuste, generaliseerbare verbetering van het SFT-paradigma.