Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models

Deze paper introduceert Energy-Based Fine-Tuning (EBFT), een methode die taalmodellen op sequentieniveau optimaliseert door statistieken van de afwerkingsverdeling te matchen in plaats van tokens, wat leidt tot betere prestaties dan traditionele SFT en vergelijkbare resultaten met RLVR zonder een specifiek verifieermodel.

De kern

Stel je voor dat je een jonge schrijver traint om verhalen te schrijven. De traditionele manier (die we nu al jaren gebruiken) is als een strenge leraar die elke zin van het verhaal controleert. Als de schrijver een woord kiest dat niet precies overeenkomt met wat de leraar had verwacht, krijgt hij een straf.

Dit werkt goed voor korte zinnen, maar er zit een groot probleem aan vast: de schrijver leert alleen maar om één woord perfect te kiezen, gebaseerd op wat de leraar al heeft geschreven. Hij leert niet hoe het hele verhaal gevoeld moet worden. Als hij in het begin van het verhaal een klein foutje maakt, raakt hij de draad kwijt. De rest van het verhaal wordt dan raar, onlogisch of herhaalt zich, omdat hij alleen maar gekeken heeft naar het volgende woord, niet naar het grote plaatje.

Dit artikel introduceert een nieuwe manier om deze schrijvers (in dit geval AI-modellen) te trainen, genaamd EBFT (Energy-Based Fine-Tuning).

Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het oude probleem: "De volgende steen leggen"

Stel je voor dat je een muur moet bouwen. De oude methode (Cross-Entropy) zegt: "Kies de perfecte steen voor nu, gebaseerd op de muur die al staat."

• Het nadeel: De AI leert alleen maar de volgende steen perfect te kiezen. Als ze een keer een verkeerde steen kiest, wordt de muur scheef. De AI weet niet dat de hele muur er uiteindelijk raar uit moet zien. Ze is een perfecte baksteen-legger, maar een slechte architect.

2. De nieuwe methode: "De hele foto vergelijken"

De auteurs van dit paper zeggen: "Laten we stoppen met kijken naar één steen. Laten we kijken naar de hele muur die de AI bouwt, en die vergelijken met een foto van een perfecte muur."

Ze noemen dit Feature Matching (Kenmerk-matchen).

• Hoe werkt het? In plaats van te kijken naar woorden, kijken ze naar de "vibe" of de "essentie" van de tekst.
• De analogie: Stel je voor dat je twee schilderijen hebt. Je kijkt niet naar elke penseelstreek apart, maar je kijkt naar de kleuren, de sfeer en de compositie. Als de sfeer van het schilderij van de AI overeenkomt met de sfeer van het meesterwerk, dan is het goed.

3. Hoe trainen ze de AI? (De "Rollouts")

De AI krijgt een opdracht (bijvoorbeeld: "Schrijf een verhaal over een kat").

1. De AI schrijft een heel verhaal (een "rollout").
2. Een "frozen" (bevroren) expert-systeem kijkt naar dit verhaal en zegt: "Hoe goed past de sfeer van dit verhaal bij het echte verhaal dat we wilden?"
3. De AI krijgt een score. Niet op basis van "was dit woord correct?", maar op basis van "voelt dit verhaal als het juiste verhaal?".
4. De AI past zichzelf aan om die "sfeer" beter te raken.

4. Waarom is dit zo slim?

• Geen "antwoordenboekje" nodig: Bij andere methoden (zoals Reinforcement Learning) heb je vaak iemand nodig die zegt: "Ja, dit antwoord is goed" of "Nee, dit is fout". Bij EBFT hoeft dat niet. De AI vergelijkt gewoon de structuur en betekenis met het voorbeeld. Dit werkt zelfs als je geen duidelijk "goed/fout" antwoord hebt (bijvoorbeeld bij het schrijven van poëzie of het vertalen van een grappige grap).
• Geen "hallucinaties": Omdat de AI leert naar het hele plaatje te kijken, raakt ze minder snel de draad kwijt. Ze blijft consistent, net als een goede schrijver die weet waar het verhaal naartoe moet.
• Beter dan de leraar: Het verrassende resultaat in het artikel is dat deze methode niet alleen betere verhalen schrijft, maar ook minder fouten maakt in de basisgrammatica dan de traditionele methode. Het is alsof de AI door te leren "sfeer te voelen", ook automatisch beter leert "woorden kiezen".

Samenvattend in één zin:

In plaats van de AI te straffen voor elke verkeerde letter, leren we haar om te luisteren naar de muziek van de tekst, zodat ze niet alleen de juiste noten speelt, maar ook een mooi melodie creëert.

Dit maakt de AI slimmer, consistenter en minder afhankelijk van menselijke correctors, waardoor ze beter kan omgaan met complexe taken zoals programmeren, vertalen en creatief schrijven.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het paper "Matching Features, Not Tokens: Energy-Based Fine-Tuning of Language Models" in het Nederlands.

Probleemstelling

De huidige standaard voor het trainen en fine-tunen van Large Language Models (LLM's) is Cross-Entropy (CE) training onder "teacher forcing". Hierbij leert het model het volgende token te voorspellen op basis van de ware (ground-truth) voorgaande tokens. Hoewel dit een stabiele en schaalbare methode is, introduceert het een fundamenteel probleem: distributieverplaatsing (distribution shift).

Tijdens de inferentie (deployments) moet het model genereren op basis van zijn eigen eerdere voorspellingen. Als het model in het begin van een sequentie een fout maakt, verandert de context voor de daaropvolgende tokens. Het model genereert dan vanuit een verdeling waar het tijdens het trainen zelden of nooit op is getraind. Dit leidt tot een degradatie van de prestaties op lange termijn, zelfs als het model een lage perplexiteit (hoge token-accuraatheid) heeft op de trainingsdata.

Bestaande oplossingen zoals Reinforcement Learning met Verifieerbare Beloningen (RLVR) loss dit op door sequentie-niveau beloningen te optimaliseren, maar vereisen een betrouwbare verifier (bijv. een unit test voor code of een menselijke beoordelaar). Voor veel open-ended taken (zoals creatief schrijven of niet-verifieerbare code) ontbreekt deze verifier, waardoor RLVR niet toepasbaar is. Bovendien kan RLVR, zelfs als het beschikbaar is, de taalmodelleringskwaliteit (cross-entropy) van het model verslechteren ten gunste van specifieke taakprestaties.

Methodologie: Energy-Based Fine-Tuning (EBFT)

De auteurs introduceren Energy-Based Fine-Tuning (EBFT), een methode die het fine-tunen richt op het matchen van sequentie-niveau statistieken in plaats van individuele tokens.

1. Het Doel: Feature Matching Loss
In plaats van het next-token verlies te minimaliseren, definieren de auteurs een Feature-Matching Loss ($L_{FM}$).

• Het model genereert meerdere "rollouts" (completions) voor een gegeven prompt.
• Een bevroren feature-netwerk $\phi$ (geïnitialiseerd vanuit het pre-trained model) projecteert zowel de gegenereerde sequenties als de ground-truth sequenties naar een vectoriële feature-ruimte.
• De loss functie meet de kwadratische fout tussen de verwachte feature-embeddings van de gegenereerde rollouts en de feature-embeddings van de ground-truth.
• Als de loss nul is, is het model "gecalibreerd": de verdeling van zijn generaties matcht de statistische momenten van de ware data in de feature-ruimte.

2. Optimalisatie via REINFORCE
Omdat de loss functie afhankelijk is van de verwachting over alle mogelijke generaties, gebruiken de auteurs een REINFORCE-achtige schatter (policy gradient).

• Voor elke prompt worden $n$ rollouts gegenereerd.
• Een reward wordt berekend die bestaat uit twee termen:
  • Alignement-term: Hoe goed de feature van een gegenereerde sample overeenkomt met de ground-truth feature.
  • Diversiteit-term: Hoe goed de sample verschilt van andere gegenereerde samples (om te voorkomen dat het model collapse naar één mode).
• De update wordt gedaan via een RLOO (REINFORCE Leave-One-Out) baseline om de variantie van de gradient te verlagen.

3. Technische Innovaties

• Strided Block-Parallel Sampling: Om efficiënt veel rollouts te genereren, gebruiken de auteurs een aangepaste attention-mask techniek (geïnspireerd door Quiet-STaR). Hiermee kunnen ze meerdere geneste prompts uit dezelfde sequentie parallel verwerken in één forward pass, wat de rekentijd aanzienlijk verlaagt.
• Whitening: Om te voorkomen dat bepaalde dimensies in de feature-ruimte de loss domineren, passen de auteurs "whitening" toe. Dit transformeert de features zodat ze een eenheidsvariantie hebben, wat de loss benadert als een $\chi^2$-divergentie, wat lokaal equivalent is aan KL-divergentie.
• Geen Task-Specifieke Reward: EBFT vereist geen externe verifier of reward model. De "reward" wordt intern gegenereerd op basis van de feature-match met de ground-truth.

Kernbijdragen

1. Feature-Matching als Trainingsdoel: Het introduceren van een trainingsdoel dat direct de statistische momenten van de rollout-verdeling matcht met die van de data, zonder afhankelijk te zijn van token-level supervisie of externe reward modellen.
2. Theoretisch Kader: Het verbinden van EBFT met KL-geregulariseerde Energy-Based Models. De auteurs tonen aan dat het optimaliseren van feature matching onder KL-regularisatie leidt tot een optimale policy die een "exponential tilt" is van het basismodel.
3. Efficiënte Implementatie: De ontwikkeling van een praktische trainingsschema met strided block-parallel sampling en whitening, waardoor EBFT toepasbaar is op grote schaal.
4. Universele Toepasbaarheid: De methode werkt zowel in verifieerbare settings (waar RLVR werkt) als in niet-verifieerbare settings (waar RLVR faalt).

Resultaten

De auteurs evalueren EBFT op drie domeinen: Q&A Coding, Unstructured Coding (ruwe code zonder instructies), en Vertaling. Ze vergelijken EBFT met SFT (Supervised Fine-Tuning) en RLVR.

• Downstream Prestaties: EBFT presteert consistent beter dan SFT en is concurrerend met (of beter dan) RLVR op benchmarks zoals HumanEval (code) en COMET (vertaling).
• Cross-Entropy (CE) en Calibratie:
  • EBFT vs. SFT: EBFT bereikt een lagere validatie cross-entropy dan SFT, ondanks dat SFT expliciet deze loss minimaliseert. Dit suggereert dat feature matching een efficiëntere manier is om de verdeling te calibreren.
  • EBFT vs. RLVR: RLVR verbetert vaak de taakprestaties, maar verslechtert de cross-entropy en de feature-matching loss aanzienlijk (het model wordt "overfit" op de reward en verliest zijn taalvaardigheid). EBFT behoudt de taalmodelleringskwaliteit terwijl het de taakprestaties verbetert.
• Niet-Verifieerbare Settings: Op "Unstructured Coding" (waar geen unit tests zijn) is RLVR onmogelijk. Hier presteert EBFT aanzienlijk beter dan SFT, wat aantoont dat het een krachtige methode is waar geen reward signalen beschikbaar zijn.
• Generalisatie: EBFT generaliseert beter naar out-of-distribution benchmarks (bijv. MultiPL-E voor code in andere talen, of ruwe tekst in vertaling) dan SFT.

Significantie en Conclusie

Dit paper biedt een paradigmaverschuiving in het fine-tunen van taalmodellen. Het toont aan dat het optimaliseren van sequentiële statistieken via feature-matching superieur is aan zowel token-level supervisie (SFT) als reward-geoptimaliseerde RL (RLVR) in termen van een balans tussen taakprestaties en algemene taalvaardigheid.

De belangrijkste implicaties zijn:

1. Onafhankelijkheid van Verifiers: EBFT maakt geavanceerde fine-tuning mogelijk voor taken waar geen automatische correctie of menselijke feedback beschikbaar is.
2. Vermijden van Trade-offs: Het lost het klassieke dilemma op waarbij het verbeteren van specifieke taken vaak ten koste gaat van de algemene taalmodelleringskwaliteit (perplexiteit).
3. Robuustheid: Het is robuuster tegen initiatie-variabiliteit dan RLVR en vereist geen warm-start om goed te presteren.

Kortom, EBFT biedt een schaalbare, theoretisch onderbouwde en praktische methode om taalmodellen te trainen die niet alleen "woorden" voorspellen, maar coherentere en statistisch correctere "sequenties" genereren.