Technische Samenvatting: Reward Shaping voor (Inference-Time) Alignment via een Stackelberg-spel Perspectief

1. Probleemstelling

Huidige alignment-pipelines voor Large Language Models (LLM's) omvatten doorgaans het leren van een beloningsmodel ($r_U$) uit gebruikersvoorkeuren en het gebruik hiervan om het LLM-beleid (policy) te optimaliseren. Deze optimalisatie wordt beperkt door een Kullback-Leibler (KL) divergentie-regularisator om het geleerde beleid dicht bij het basismodel ($\rho_{base}$) te houden, wat de vloeiendheid en taalcompetentie waarborgt.

De auteurs identificeren een fundamentele beperking in deze standaardpipeline: het direct optimaliseren voor de geleerde beloning $r_U$ onder een KL-constraint is suboptimaal voor het maximaliseren van de bruikersutiliteit. Wanneer het basismodel sterke biases vertoont (bijv. politieke voorkeur) die conflicteren met de voorkeuren van de gebruiker, voorkomt de KL-constraint dat het beleid voldoende verschuift naar de geprefereerde distributie, zelfs als het beloningsmodel de voorkeur correct identificeert.

Omgekeerd brengt het simpelweg vergroten van beloningen om een verschuiving te forceren het risico op reward hacking met zich mee, waarbij het model incoherente of degeneratieve outputs genereert om de beloningssignalen te maximaliseren terwijl de intentie van de KL-constraint wordt geschonden. Het kernprobleem dat wordt aangepakt, is hoe men de beloningsfunctie die gebruikt wordt voor alignment optimaal kan ontwerpen (shapen) om de balans te vinden tussen het maximaliseren van de bruikersutiliteit en het risico op reward hacking onder KL-regularisatie.

2. Methodologie: Stackelberg-spel Formulering

De auteurs formaliseren het probleem van beloningsontwerp als een Stackelberg-spel, een tweespeler sequentieel spel:

• Leider (Reward Model Provider): Selecteert een beloningsmodel $r$ om de verwachte utiliteit van de gebruiker (gebaseerd op de ware utiliteitsfunctie $r_U$) te maximaliseren, anticiperend op de reactie van de LLM. De leider is niet verplicht om $r_U$ direct te gebruiken, maar kan elke beloningsfunctie $r: \mathcal{X} \times \mathcal{Y} \to \mathbb{R}$ kiezen die begrensd is door $B$.
• Volger (LLM): Observeert het beloningsmodel $r$ en optimaliseert zijn beleid $\rho_r$ om de verwachte beloning minus de KL-divergentie van het basisbeleid te maximaliseren:
  $$\rho_r(\cdot|x) = \arg\max_{\rho} \left( \mathbb{E}_{y \sim \rho}[r(x, y)] - \beta \cdot D_{KL}(\rho(\cdot|x) \parallel \rho_{base}(\cdot|x)) \right)$$

Theoretische Karakterisering

De auteurs leiden een gesloten vorm karakterisering af van het optimale beloningsmodel $r^*$:

1. Drempelstructuur (Threshold Structure): Het optimale beloningsmodel is een threshold reward. Voor een gegeven prompt $x$ worden reacties gepartitioneerd op basis van of hun ware utiliteit $r_U(x, y)$ boven of onder een specifieke drempelwaarde $m^*(x)$ ligt.
   • Als $r_U(x, y) > m^*(x)$, wordt de beloning ingesteld op de bovengrens $B$.
   • Als $r_U(x, y) < m^*(x)$, wordt de beloning ingesteld op $0$.
   • Als $r_U(x, y) = m^*(x)$, ligt de beloning in $[0, B]$.
2. Zelfconsistentie-conditie: De optimale drempelwaarde $m^*(x)$ voldoet aan de conditie dat de verwachte utiliteit van het resulterende beleid gelijk is aan de drempelwaarde zelf:
   $$m^*(x) = \mathbb{E}_{y \sim \rho_{r^*}(\cdot|x)} [r_U(x, y)]$$
3. Berekening: Omdat de exacte verwachting onhandelbaar is, stellen de auteurs voor de drempelwaarde te berekenen via Monte Carlo-sampling vanuit het basisbeleid. Ze definiëren een hulpfunctie $F_x(m)$ waarvan de unieke wortel overeenkomt met de optimale drempelwaarde. Deze wortel kan efficiënt worden gevonden via een bisection-algoritme.

Praktische Implementatie: Soft Thresholding (SRS)

Hoewel de harde drempel analytisch optimaal is, is deze discontinu en gevoelig voor kleine perturbaties. De auteurs introduceren Stackelberg Reward Shaping (SRS), een soft threshold variant die gebruikmaakt van een sigmoidfunctie:
$$r_{\hat{m}^*, \alpha}(x, y) = B \cdot \sigma(\alpha \cdot (r_U(x, y) - \hat{m}^*(x)))$$
waarbij $\hat{m}^*(x)$ de geschatte drempelwaarde is en $\alpha$ de scherpte van de overgang regelt. Dit stelt de methode in staat om te interpoleren tussen geen alignment-effect ($\alpha=0$) en het analytische optimum ($\alpha \to \infty$).

Integratie met Inference-Time Alignment

Het framework wordt geïntegreerd in bestaande inference-time alignment methoden (specifiek Controlled Decoding (CD) en Alignment as Reward-Guided Search (ARGS)):

1. Offline Fase: Sample $M$ reacties uit het basismodel voor een dataset van prompts. Bereken $r_U$ voor deze samples. Schat de drempelwaarde $\hat{m}^*(x)$ via bisection. Pas de SRS-transformatie toe om een shaped reward dataset te genereren. Train een Q-waarde estimator op deze shaped data.
2. Inference Fase: Gebruik de getrainde Q-waarde estimator (die de verwachte beloning onder de shaped reward bevat) om de token-waarschijnlijkheden van het basismodel tijdens het decoderen te herwegen.

3. Belangrijkste Bijdragen

1. Suboptimaliteit van Standaard Pipelines: Aangetoond dat het direct gebruiken van een beloningsmodel geleerd uit voorkeuren suboptimaal is onder KL-geregulariseerde doelstellingen, met name wanneer de biases van het basismodel conflicteren met de voorkeuren van de gebruiker.
2. Stackelberg-spel Formulering: Het optimaliseren van het beloningsmodel geformuleerd als een Stackelberg-spel, waarbij een gesloten oplossing is afgeleid die aantoont dat het optimale beloningsmodel een drempelstructuur volgt.
3. Efficiënte Approximatie: Een algoritme ontwikkeld om de optimale drempelwaarde te benaderen met behulp van Monte Carlo-sampling vanuit het basisbeleid, wat integratie in bestaande methoden met minimale overhead mogelijk maakt.
4. Empirische Validatie: Getoond dat de voorgestelde SRS-methode consistent de gemiddelde beloning verbetert en een hoge win-tie ratio behaalt ten opzichte van baselines in inference-time settings.

4. Experimentele Resultaten

De methode werd geëvalueerd op HH-RLHF en SHP datasets met Qwen3-8B en Llama3-8B backbones.

• Maximale Beloning: SRS-geïntegreerde methoden (SRS-CD en SRS-ARGS) behaalden consistent de hoogste gemiddelde beloning over alle evaluatiesettings vergeleken met vanilla baselines en heuristische reward shaping schema's (Minmax, Meanstd).
• Kwaliteitsmetrieken: De methoden behielden diversiteit en coherentie die vergelijkbaar is met baselines, wat wijst op geen degradatie in taalkwaliteit.
• GPT-4 Evaluatie: In head-to-head vergelijkingen behaalden SRS-methoden een gemiddelde Win-Tie rate van meer dan 66% tegen alle baselines (Vanilla, Minmax, Meanstd), wat suggereert dat verbeteringen niet worden gedreven door reward hacking.
• Robuustheid:
  • Generalisatie: SRS toonde superieur presteren wanneer getraind op één distributie (SHP) en geëvalueerd op out-of-distribution prompts (HH-RLHF en HarmfulQA).
  • Ablatie: De specifieke drempelstructuur van SRS presteerde beter dan eenvoudige reward capping en mean-based thresholding, wat het belang bevestigt van de zelfconsistente drempelwaarde afgeleid uit de speltheoretische formulering.
• Overhead: Voor SRS-CD is reward shaping offline met nul inference-time overhead. Voor SRS-ARGS is de overhead verwaarloosbaar (~0,03 seconden per 128-token reactie).

5. Betekenis en Claims

Het paper claimt dat reward shaping niet louter een heuristiek is, maar een optimalisatieprobleem dat systematisch kan worden opgelost. Door het alignmentproces te bekijken door een speltheoretische lens, demonstreren de auteurs dat het beloningsmodel dat de LLM stuurt, niet noodzakelijkerwijs (en ook niet zou moeten zijn) een directe proxy voor de bruikersutiliteit. In plaats daarvan moet het "gevormd" worden om rekening te houden met de beperkingen die door het basismodel (KL-divergentie) worden opgelegd.

De betekenis ligt in het bieden van een geprincieerd framework om de biases van het basismodel tegen te gaan zonder reward hacking te induceren. De auteurs benadrukken dat hun aanpak lichtgewicht is, naadloos integreert in bestaande inference-time alignment methoden en een robuuste oplossing biedt voor het aanpassen van vaste modellen aan diverse gebruikersvoorkeuren.

Door de auteurs genoemde beperkingen:

• Het framework gaat uit van een welwillende gebruiker wiens voorkeuren overeenkomen met menselijke waarden. Als het beloningsmodel kwaadaardige voorkeuren reflecteert, kan het resulterende beleid misaligned raken.
• De methode is afhankelijk van de beschikbaarheid van een beloningsmodel en Monte Carlo-samples van het basisbeleid.
• De huidige empirische focus ligt op inference-time alignment, hoewel de auteurs opmerken dat er potentieel is voor uitbreiding naar train-time algoritmen (zoals GRPO) als toekomstig werk.