Towards Bridging the Reward-Generation Gap in Direct Alignment Algorithms

Titel: Towards Bridging the Reward-Generation Gap in Direct Alignment Algorithms

Auteurs: Zeguan Xiao, Yun Chen, Jian Yang, Guanhua Chen, Ke Tang (Shanghai University of Finance and Economics, Beihang University, Southern University of Science and Technology)

1. Het Probleem: De Reward-Generation Gap

Directe Uitlijning Algoritmen (Direct Alignment Algorithms - DAAs), zoals Direct Preference Optimization (DPO) en Simple Preference Optimization (SimPO), zijn populair geworden als efficiëntere alternatieven voor Reinforcement Learning from Human Feedback (RLHF). Ze omzeilen de noodzaak van een expliciete reward-modelleringsstap.

De auteurs identificeren echter een fundamentele beperking, de "Reward-Generation Gap":

• Misalignement: Er is een discrepantie tussen de trainingsdoelen van DAAs en de autoregressieve decodingsdynamiek van Large Language Models (LLMs).
• Oorzaak: DAAs optimaliseren op sequentie-niveau (de hele respons), maar LLM's genereren token voor token. De impliciete reward-functies van DAAs geven elke token in een sequentie gelijke weging.
• Het Kernprobleem: In autoregressieve generatie is de kwaliteit van de prefix (de eerste tokens) cruciaal. Fouten in de vroege tokens leiden tot "exposure bias", waarbij fouten zich versterken en de kwaliteit van de hele respons degradëren. Omdat DAAs de vroege tokens niet zwaarder wegen dan latere tokens, wordt de kritieke invloed van de prefix ondergewaardeerd tijdens het trainingsproces. Dit leidt ertoe dat een hogere totale reward niet noodzakelijk leidt tot betere generatiekwaliteit.

2. Methodologie: Prefix-Oriented Equal-length Training (POET)

Om deze kloof te dichten, stellen de auteurs POET voor, een eenvoudige maar effectieve data-augmentatiestrategie.

• Theoretische Basis: De auteurs analyseren DAAs vanuit het perspectief van een Token-level Markov Decision Process (MDP). Ze bewijzen (Theorema 1) dat het optimaliseren van een "equal-length sub-trajectory" (subtrajecten van gelijke lengte) leidt tot dezelfde optimale policy als sequentie-niveau optimalisatie, maar wel betere supervisie biedt voor de cruciale vroege tokens.
• Empirische Validatie: Experimenten tonen aan dat het kwaliteitsverschil tussen een "preferred" (gewenste) en "dispreferred" (ongewenste) respons zich zeer vroeg in de generatie manifesteert (in de prefix) en dat dit verschil later in de sequentie nauwelijks meer toeneemt.
• De POET-strategie:
  1. Voor elk paar (gewenste $y_w$, ongewenste $y_l$) wordt de lengte bepaald van het kortste antwoord: $k = \min(|y_w|, |y_l|)$.
  2. Beide antwoorden worden getruncated (afgeknipt) tot deze lengte $k$.
  3. Het trainingsproces gebruikt deze getruncated, gelijke-lengte paren.
• Voordeelpunten:
  • Universeel: Werkt met elke DAA (DPO, SimPO, etc.) zonder de optimalisatie-objektieven te wijzigen.
  • Hyperparameter-vrij: Geen extra parameters nodig; de truncatie gebeurt automatisch op basis van de kortste respons.
  • Minimale Ruis: Omdat het kwaliteitsverschil zich in de prefix bevindt, is het risico dat de truncatie de voorkeur omkeert (dat de ongewenste respons beter wordt dan de gewenste na truncatie) zeer klein.

3. Belangrijkste Bijdragen

1. Identificatie van de Gap: De paper definieert en analyseert de "Reward-Generation Gap" als een fundamenteel misalignement tussen trainingsobjectieven en autoregressieve dynamiek, veroorzaakt door het negeren van de toenemende belangrijkheid van prefix-tokens.
2. Theoretisch Bewijs: Levering van een theoretisch bewijs dat optimalisatie op gelijke-lengte subtrajecten equivalent is aan sequentie-niveau optimalisatie, maar superieure convergentie-eigenschappen biedt voor vroege tokens.
3. POET Methode: Introductie van een hyperparameter-vrije, plug-and-play methode die de training focust op de kwaliteit van de prefix.
4. Uitgebreide Validatie: Demonstratie dat POET werkt over verschillende modelarchitecturen (Mistral, Llama-3, Gemma) en verschillende DAA's (DPO, SimPO).

4. Resultaten

De auteurs hebben uitgebreide experimenten uitgevoerd met DPO en SimPO op diverse benchmarks:

• Instruction Following:
  • Op AlpacaEval 2 (een leidende benchmark voor instructievolging) boekte POET aanzienlijke verbeteringen.
  • De beste prestatie was een stijging van 11.8 punten in de Length-Controlled (LC) win-rate voor Mistral-Base-7B met DPO.
  • POET presteerde consistent beter dan de standaardimplementaties van DPO en SimPO.
• Downstream Taken:
  • Er was geen sprake van een "alignment tax" (verlies van algemene capaciteiten). Integendeel, er waren bescheiden maar consistente verbeteringen in taken zoals wiskunde (GSM8K) en algemene kennis (MMLU).
• Veiligheid (Safety Alignment):
  • Op de PKU-SafeRLHF dataset verbeterde POET de veiligheidsrate aanzienlijk (van 45.2% naar 82.1% voor Mistral-7B). Dit wordt toegeschreven aan het feit dat veiligheidsafwijzingen vaak al in de prefix gebeuren.
• Vergelijking met andere methoden:
  • POET presteerde beter dan andere token-level methoden zoals SamPO (random subset) en D2PO (tijdsgewicht), wat aantoont dat het focussen op de prefix via gelijke lengte effectiever is dan het willekeurig selecteren of exponentieel afwegen van tokens.

5. Betekenis en Conclusie

De paper biedt een cruciaal inzicht in waarom directe uitlijning soms suboptimaal presteert: het negeert de dynamiek van hoe LLM's eigenlijk genereren (token voor token, waarbij de start cruciaal is).

• Praktische Impact: POET biedt een eenvoudige, goedkope en effectieve manier om bestaande uitlijningsmethodes te verbeteren zonder complexe architectuurveranderingen of extra rekenkracht voor reward-modellen.
• Toekomstperspectief: Het werk benadrukt de noodzaak om trainingsdoelen beter af te stemmen op de autoregressieve aard van generatieve modellen. Het suggereert dat de kwaliteit van de "start" van een antwoord vaak belangrijker is dan de totale lengte of de gemiddelde kwaliteit van de hele tekst.

Kortom, POET sluit de kloof tussen wat we trainen (sequentiële voorkeur) en hoe het model werkt (token-voor-token generatie) door de training te focussen op het gedeelte van de tekst dat het meest bepalend is voor de uiteindelijke kwaliteit: de prefix.