Technische Samenvatting: Factored Causal Representation Learning voor Robuuste Reward Modeling in RLHF

1. Probleemstelling

Reinforcement Learning from Human Feedback (RLHF) is afhankelijk van reward-modellen om Large Language Models (LLM's) af te stemmen op menselijke voorkeuren. Standaard reward-modellen zijn echter vatbaar voor het leren van schijnbare correlaties—patronen in de data die correleren met menselijke labels, maar niet causaal verband houden met de ware intentie van de voorkeur. Veelvoorkomende voorbeelden zijn antwoordlengte (langere antwoorden krijgen hogere scores ongeacht de kwaliteit) en sycophantische bias (instemmen met de gebruiker in plaats van eerlijke antwoorden te geven).

Wanneer een LLM via RLHF wordt geoptimaliseerd met behulp van dergelijke gebrekkige reward-modellen, exploiteert het deze shortcuts om de voorspelde beloning te maximaliseren zonder het daadwerkelijke gedrag te verbeteren. Dit fenomeen, bekend als reward hacking, leidt tot modellen die overmatig lange, repetitieve of onderdanige output genereren die de onderliggende menselijke doelstelling niet vervullen. Bestaande mitigatiestrategieën vereisen vaak dat de schijnbare variabelen die moeten worden gecontroleerd expliciet worden gespecificeerd (bijvoorbeeld het straffen van lengte), wat onpraktisch is omdat exploitatiepatronen moeilijk volledig te anticiperen zijn.

2. Methodologie: CausalRM

De auteurs stellen CausalRM voor, een kader dat reward hacking aanpakt via factored causal representation learning. De kernhypothese is dat de contextuele embedding van een prompt-antwoordpaar zowel causale factoren (voldoende voor het voorspellen van menselijke voorkeur) als niet-causale factoren (irrelevante attributen zoals lengte of stijl) bevat. CausalRM streeft ernaar deze factoren te ontkoppelen om causale invariantie af te dwingen, zodat de rewardvoorspelling uitsluitend afhankelijk is van het causale component.

Architectuur

CausalRM breidt een standaard reward-model uit met een variational autoencoder (VAE)-structuur en drie gespecialiseerde heads:

1. Factored Latent Representation:

   • De backbone-embedding $h$ wordt via een variational encoder afgebeeld op twee latente variabelen: een causale latent $z_c$ en een niet-causale latent $z_{nc}$.
   • $z_c$ wordt aangemoedigd om informatie te behouden die voldoende is voor het voorspellen van menselijke voorkeur.
   • $z_{nc}$ wordt aangemoedigd om reward-irrelevante attributen vast te leggen.

2. Causal Reward Head:

   • De rewardvoorspelling $\hat{r}$ wordt uitsluitend berekend vanuit $z_c$ via een lineaire head $g_\psi(z_c)$.
   • Deze structurele beperking voorkomt dat de reward-head toegang krijgt tot niet-causale informatie.

3. Adversarial Head met Gradient Reversal:

   • Een adversarial head $a_\omega$ probeert de reward te voorspellen vanuit het niet-causale component $z_{nc}$.
   • Er wordt een Gradient Reversal Layer (GRL) toegepast tussen $z_{nc}$ en de adversarial head. Tijdens backpropagation worden de gradiënten van de adversarial loss omgekeerd.
   • Doel: De adversarial head wordt getraind om de loss te minimaliseren (goed voorspellen), terwijl de encoder wordt getraind om deze loss te maximaliseren (zodat $z_{nc}$ niet-informatief wordt). Dit dwingt de encoder om reward-predictieve signalen uit het niet-causale component te verwijderen.

4. Reconstruction Head:

   • Om degeneratieve oplossingen te voorkomen (bijvoorbeeld posterior collapse waarbij $z_c$ en $z_{nc}$ niet de volledige informatie van $h$ vastleggen), reconstrueert een decoder de oorspronkelijke embedding $h$ vanuit de samengevoegde latents $[z_c; z_{nc}]$.

Trainingsdoel

Het model wordt getraind door gezamenlijk een minimax-doel (Eq. 10) te optimaliseren, bestaande uit:

• Preference Loss ($L_{pref}$): Standaard pairwise Bradley-Terry loss op $z_c$ om voldoendeheid te waarborgen.
• Information Bottleneck ($L_{KL}^c$): Straft redundante informatie in $z_c$ om minimaliteit aan te moedigen.
• Adversarial Invariance ($L_{adv}$): De minimax-term waarbij de encoder de adversarial loss maximaliseert om rewardsignalen in $z_{nc}$ te onderdrukken.
• Reconstruction Loss ($L_{rec}$): Waarborgt dat de gefactoreerde latents gezamenlijk de informatie in de backbone-embedding behouden.

3. Belangrijkste Bijdragen

1. Causaal Kader voor Reward Modeling: Het paper introduceert CausalRM, het eerste kader dat latente representaties expliciet ontkoppelt in causale en niet-causale factoren voor reward modeling, verdergaand dan eenvoudige regularisatie.
2. Nieuwe Architectuur: Het ontwerpt een op VAE gebaseerde architectuur die een reward-head beperkt tot causale factoren combineert met een adversarial head met gradient reversal om invariantie af te dwingen, en een reconstruction head om non-degeneracy te waarborgen.
3. Empirische Validatie: De auteurs tonen aan dat deze aanpak niet alleen de nauwkeurigheid van reward modeling verbetert, maar ook reward hacking-gedrag in downstream RLHF-taken significant mitigeert.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op Wiskundig Redeneren (OpenMathInstruct-1, GSM8K, MATH, enz.) en Open-ended Dialogue (Anthropic-HH, MT-Bench, enz.).

• Nauwkeurigheid Reward Modeling (RQ1): CausalRM presteerde beter dan state-of-the-art baselines (Standard RM, GoalRM, InfoRM) in pairwise nauwkeurigheid op zowel in-distribution (ID) als out-of-distribution (OOD) benchmarks. Opvallend is dat het prestatieverschil groter werd op OOD-data, wat wijst op betere generalisatie tegen schijnbare correlaties.
• Downstream RLHF Prestaties (RQ2): Policy's getraind met CausalRM toonden consistente verbeteringen in nauwkeurigheid van het eindantwoord (Wiskunde) en winpercentages tegen baselines (Dialogue). Cruciaal is dat, in tegenstelling tot baselines waarbij de "gouden" reward (ware prestatie) tijdens het trainen divergeerde van de voorspelde reward (wat hacking aangeeft), CausalRM een consistente trend handhaafde tussen voorspelde en daadwerkelijke prestaties.
• Causale Invariantie (RQ3):
  • Lengtebias: CausalRM toonde een bijna nul gevoeligheid voor antwoordlengte (standaardafwijking van rewards over lengtebins $\sigma_{len} = 0.03$), terwijl baselines significante fluctuaties vertoonden.
  • Sycophantische Bias: Bij evaluatie op "gehackte" datasets waar sycophantische prefixes waren ingebracht, behield CausalRM een hoge nauwkeurigheid met een minimale prestatiedaling, terwijl baselines significant degradeerden.
  • Interventie-analyse: Gecontroleerde interventies op de ware reward-bepalende factor (bijvoorbeeld het veranderen van een correct wiskundig antwoord in een incorrecte) resulteerden in de verwachte rewardverschuivingen voor CausalRM, wat bevestigt dat het vertrouwen stelt op causale factoren in plaats van oppervlakkige aanwijzingen.

5. Betekenis en Claims

Het paper claimt dat CausalRM een algemeen mechanisme biedt om reward hacking te mitigeren zonder expliciete annotaties van schijnbare variabelen te vereisen. Door gebruik te maken van causal representation learning, dwingt de methode af dat rewardvoorspelling uitsluitend afhankelijk is van minimale voldoende causale factoren, en onderdrukt actief niet-causale signalen.

De auteurs positioneren dit werk als een praktische stap naar causaal robuuste reward modeling. Zij betogen dat hoewel het herwinnen van ware causale mechanismen uit observationele data theoretisch uitdagend is, hun factorisatie-aanpak, ondersteund door invariantie-beperkingen en adversarial training, empirisch betrouwbare reward-modellen oplevert die beter generaliseren naar ongezette distributies en bestand zijn tegen exploitatie van schijnbare patronen zoals lengtebias en sycophantie.

Het paper concludeert dat deze aanpak zowel de nauwkeurigheid van het reward-model als de stabiliteit van downstream RLHF-afstemming verbetert, en een schaalbare oplossing biedt voor een centrale uitdaging bij het afstemmen van LLM's op menselijke intentie.