Technische Samenvatting: SMOPD

1. Probleemstelling

Het artikel behandelt een fundamentele beperking in multi-reward Reinforcement Learning from Human Feedback (RLHF) voor Large Language Models (LLMs). Hoewel bestaande methoden zoals Group Reward-Decoupled Normalization Policy Optimization (GDPO) de standaard Group Relative Policy Optimization (GRPO) hebben verbeterd door beloningsdimensies afzonderlijk te normaliseren vóór aggregatie, worstelen zij nog steeds met belonings-densiteitsimbalans (reward-density imbalance).

In veel praktische scenario's worden modellen getraind met een mix van dense rewards (die fijnmazige, continue feedback bieden, bijvoorbeeld nauwkeurigheidsscores van 0,1 tot 1,0) en sparse rewards (die binaire, alles-of-niets feedback bieden, bijvoorbeeld format-compliance van 0 of 1).

• Het kernprobleem: In groep-gebaseerde advantage-estimatie bieden dense rewards consistente variatie binnen de groep, waardoor het model effectief kan leren. Sparse rewards kennen echter vaak identieke waarden toe (bijv. 0) aan de meerderheid van de responsen binnen een rollout-groep. Dit resulteert in een nul of bijna nul binnen-groep variantie voor de sparse dimensie.
• Gevolg: Zelfs met per-reward normalisatie wordt het sparse reward-signaal overstemd door het dense reward-signaal tijdens gebalanceerde optimalisatie. Het model slaagt er niet in de capaciteit te leren die bij de sparse reward hoort (bijv. strikte output-formattering), omdat het optimalisatiesignaal onvoldoende is om policy-updates aan te drijven, ongeacht de prioriteitsgewichten die in een single-policy setting worden toegewezen. Het verhogen van de prioriteit van de sparse reward om deze leerbaar te maken, verstoort vaak de uiteindelijke policy, wat de prestaties op andere doelstellingen verslechtert.

2. Methodologie: SMOPD

Om deze spanning op te lossen, stellen de auteurs Specialize-and-Merge Online Policy Distillation (SMOPD) voor, een tweestaps trainingsframework dat de acquisitie van capaciteiten ontkoppelt van de beloningsbalancering.

Fase 1: Specialiseer (Reward-gespecialiseerde Teachers)

In plaats van één enkele policy te trainen om alle beloningen te balanceren, traint SMOPD meerdere "teacher" policies, die elk gespecialiseerd zijn voor een specifieke beloningsdimensie.

• Mechanisme: De base policy wordt getraind met onderscheidende reward-prioriteitsconfiguraties. Voor een doelbeloningsdimensie $k$ wordt het prioriteitsgewicht $w_k$ aanzienlijk hoger ingesteld dan de andere gewichten ($w_k \gg w_j$).
• Effect: Dit versterkt het gradiëntsignaal voor de doelbeloning, waardoor de teacher effectief de sparse of moeilijk te leren distributie kan leren (bijv. een "Format Teacher" die prioriteit geeft aan binaire format-compliance) zonder te worden overstemd door dense signalen.
• Resultaat: Er wordt een set complementaire teachers geproduceerd, waarbij elke teacher uitblinkt in een specifieke beloningsdimensie, potentieel ten koste van anderen, maar collectief de volledige reward-ruimte dekt.

Fase 2: Merge (Online Policy Distillation)

De gespecialiseerde capaciteiten van de teachers worden samengevoegd in een enkele "student" policy via On-Policy Distillation (OPD).

• Teacher Mixture: Op elke token-positie $t$ in een door de student gegenereerde sequentie, is de doelverdeling een mengeling van de next-token distributies van de teachers. SMOPD gebruikt een uniforme mengeling ($\alpha_m = 1/M$) en benadert de volledige vocabulaire-mengeling met behulp van een top-$\kappa$ constructie (waarbij de top-$\kappa$ tokens van elke teacher worden behouden) om de computationele efficiëntie te waarborgen.
• Distillatieverlies: De student minimaliseert de Forward KL-divergentie tussen de eigen distributie en de teacher-mengeling. Forward KL is gekozen omdat het "mode-covering" is, wat de student aanmoedigt om waarschijnlijkheidsmassa te plaatsen op alle tokens met een hoge waarschijnlijkheid van elke teacher (cruciaal wanneer verschillende teachers overtuigend zijn over verschillende tokens, bijv. format-wrappers versus functie-argumenten).
• GDPO Anchor: Om te voorkomen dat de student enkel de teachers imiteert en de taak-niveau optimalisatie verliest, wordt een GDPO anchor loss toegevoegd. Deze loss gebruikt gebalanceerde gewichten voor alle beloningen, wat een sequentie-niveau signaal biedt dat de student begeleidt bij het optimaliseren van de oorspronkelijke multi-reward doelstelling.
• Totale Doelstelling: De student optimaliseert een combinatie van de distillatie-loss (token-niveau begeleiding) en de anchor loss (sequentie-niveau taakrichting).

3. Belangrijkste Bijdragen

1. Oplossing voor de Reward-Balancing Tension: SMOPD adresseert het inherente conflict in single-policy multi-reward RL waarbij sparse rewards onleerbaar zijn onder gebalanceerde prioriteiten. Door de teachers te specialiseren, maakt het sparse distributies leerbaar zonder de uiteindelijke policy permanent te verstoren.
2. Tweestaps Specialize-and-Merge Framework: De methode introduceert een nieuwe workflow waarbij de acquisitie van capaciteiten wordt gescheiden van de beloningsbalancering. Teachers leren in gunstige regimes (hoge prioriteit), en de student voegt deze capaciteiten samen via token-niveau distillatie.
3. Superieur in Complementaire en Conflicterende Settings: Het artikel toont aan dat SMOPD beter presteert dan scalarized baselines (GRPO, GDPO, GD2PO) over verschillende modelmaten (1.5B, 3B, 7B) en beloningsstructuren (complementair: tool accuracy/format; conflicterend: helpful/harmless).
4. Merging Overtreft de Teachers: De samengevoegde student policy balanceert niet alleen de doelstellingen, maar presteert vaak ook beter dan de individuele gespecialiseerde teachers en de scalarized baselines, wat aangeeft dat het distillatieproces complementaire kennis herstelt die niet werd gebruikt door single-objective training.

4. Experimentele Resultaten

De auteurs evalueerden SMOPD in twee primaire settings:

A. Complementaire Beloningen (Tool Calling)

• Setup: Training op RLLA-4K met nauwkeurigheid (dense) en format (sparse binary) beloningen.
• Bevindingen:
  • GDPO Falen: Onder gebalanceerde prioriteiten bereikte GDPO slechts 8,8% format-compliance op de 1.5B backbone, omdat het sparse signaal werd overstemd.
  • Teacher Succes: De op format gespecialiseerde teacher (hoge prioriteit op format) bereikte 97,5% compliance maar lagere nauwkeurigheid.
  • SMOPD Succes: De samengevoegde student bereikte 97,5% format-compliance terwijl een hoge nauwkeurigheid werd behouden (RLLA Mean score van 2.740 vs. GDPO's 1.849). Dit vertegenwoordigt een verbetering van +48% ten opzien van de GDPO baseline op het 1.5B model.
  • Generalisatie: SMOPD vertoonde ook superieure prestaties op held-out benchmarks zoals API-Bank (LLM-Judge) en BFCL.

B. Conflicterende Beloningen (Helpful vs. Harmless)

• Setup: Training op Safe-RLHF benchmarks (HH-RLHF, PKU-SafeRLHF, Alpaca) met helpfulness en harmlessness beloningen.
• Bevindingen:
  • SMOPD presteerde consequent beter dan GDPO en GD2PO over 3B en 7B backbones (Qwen2.5 en Llama-3.2).
  • Op Qwen2.5-7B behaalde SMOPD een Overall score van 5.646, waarmee het de beste scalarized baseline (GD2PO: 5.571) en beide individuele teachers overtrof.
  • Zelfs wanneer individuele teachers achteruitgingen (bijv. op Llama-3.2-3B), herstelde het merge-proces de prestaties om de scalarized baselines te benaderen of te overtreffen.

C. Ablatie-studies

• Noodzaak van de Anchor: Het verwijderen van de GDPO anchor (gebruik van enkel OPD) resulteerde in een student die plateau bereikte op het prestatieniveau van de teacher-mengeling. De anchor is essentieel om de student voorbij de "imitation ceiling" te duwen om de werkelijke taakbeloning te optimaliseren.
• Forward vs. Reverse KL: Forward KL bleek superieur aan Reverse KL in deze multi-teacher setting omdat het de diverse modi van de teacher-mengeling behoudt, terwijl Reverse KL de neiging heeft om mode-seeking te zijn en mogelijk gespecialiseerde capaciteiten verliest.
• Top-$\kappa$ Sensitiviteit: Een kleine top-$\kappa$ (bijv. 16) was voldoende, aangezien het ~99,4% van de teacher waarschijnlijkheidsmassa behield, wat de efficiëntie van de benadering valideert.

5. Betekenis en Claims

Het artikel claimt dat SMOPD een robuuste oplossing biedt voor het belonings-densiteitsimbalans probleem in multi-reward RL. De betekenis ligt in:

• Ontkoppeling van Leren en Balanceren: Het daagt het paradigma uit om alle beloningen tegelijkertijd te balanceren binnen een enkele optimalisatiestap. In plaats daarvan maakt het gebruik van specialisatie om ervoor te zorgen dat alle beloningssignalen effectief worden geleerd voordat ze worden samengevoegd.
• Token-niveau Granulariteit: Door gebruik te maken van on-policy distillatie op token-niveau, biedt SMOPD een dichte supervisie voor sparse rewards (zoals formattering) die groep-gebaseerde advantage-estimatie niet kan bieden.
• Schaalbaarheid en Generaliteit: De methode is effectief over variërende modelgroottes (1.5B tot 7B) en beloningsvormen (complementair en conflicterend), wat suggereert dat het een generaliseerbare aanpak is voor multi-objective alignment in LLM's.

De auteurs concluderen dat SMOPD een verenigde policy produceert die de sterke punten van gespecialiseerde teachers behoudt terwijl het een gebalanceerde taak-niveau optimalisatie bereikt, waarmee het de beperkingen van huidige scalarized multi-reward methoden effectief overwint.