Technische Samenvatting: Value-Free Policy Optimization via Reward Partitioning

Probleemstelling

Het afstemmen van Large Language Models (LLM's) op menselijke voorkeuren steunt doorgaans op Reinforcement Learning from Human Feedback (RLHF) of Direct Preference Optimization (DPO). Deze paradigma's kampen echter met aanzienlijke uitdagingen op het gebied van schaalbaarheid en stabiliteit:

• RLHF vereist kostbare menselijke annotatie van paarwijze voorkeuren, het trainen van een apart Reward Model (RM) en onstabiele reinforcement learning-loops (bijv. PPO) die campagnie-inefficiënt zijn en gevoelig zijn voor hyperparameters.
• Direct Preference Optimization (DPO-stijl) methoden elimineren de expliciete RM en de RL-loop, maar vereisen gecureerde paarwijze voorkeursdata. Naarmale modellen verbeteren, wordt het onderscheid tussen hoogwaardige responsen steeds moeilijker, wat de annotatiekosten opdrijft.
• Single-Trajectory Learning biedt een praktisch alternatief met gebruikmakend van $(prompt, response, scalar\_reward)$ tupels, die gemakkelijker te annoteren zijn (bijv. duimpje omhoog/omlaag). Bestaande methoden voor deze setting, zoals Direct Reward Optimization (DRO), vertrouwen op het gezamenlijk leren van een policy en een value function. Deze gezamenlijke optimalisatie introduceert variantie, instabiliteit en gevoeligheid voor off-policy data, omdat fouten in de waarde-inschatting de training van de policy kunnen beïnvloeden.

Het kernprobleem dat wordt aangepakt, is de behoefte aan een stabiele, schaalbare en eenvoudige methode om policies direct te optimaliseren vanuit scalaire beloningsdata zonder de complexiteit en instabiliteit van het leren van een hulp-value function.

Methodologie: Reward Partitioning Optimization (RPO)

De auteurs stellen Reward Partitioning Optimization (RPO) voor, een value-free objectief dat gebruikmaakt van de analytische vorm van de optimale policy onder KL-regularisatie om de noodzaak van een geleerde value function te elimineren.

Theoretische Fundering

Vertrekkend van de optimale policy-formulering in KL-geregulariseerde beloningsmaximalisatie:
$$\pi^*(y|x) = \pi_{ref}(y|x) \exp\left(\frac{1}{\tau} (r(x, y) - V^*(x))\right)$$
waarbij $V^*(x)$ de soft value function is die fungeert als een partitiefunctie om beloningen te normaliseren. Traditionele methoden zoals DRO benaderen $V^*(x)$ met een apart neuraal netwerk, wat leidt tot gekoppelde optimalisatie van de policy ($\pi_\theta$) en de value ($V_\phi$).

RPO omzeilt dit door de partitiefunctie empirisch te schatten direct uit de data. Voor een gegeven prompt $x$, aggregeert de methode alle beschikbare responsen $y_j$ in de dataset om een empirische partitiefunctie $\hat{Z}(x)$ te berekenen:
$$\hat{Z}(x) = \sum_{j \in I_x} \pi_{ref}(y_j|x) \exp\left(\frac{1}{\tau} r(x, y_j)\right)$$
waarbij $I_x$ de verzameling indices is voor samples met prompt $x$. De empirische waarde is vervolgens $\hat{V}(x) = \tau \log \hat{Z}(x)$.

Optimalisatie-objectief

RPO formuleert de leeropdracht als een directe regressie op de log-ratio tussen de geleerde policy en de referentie-policy. De verliesfunctie minimaliseert de kwadratische fout tussen de voorspelde log-ratio en de gecentreerde, geschaalde beloning:
$$\mathcal{L}_{RPO}(\theta) = \frac{1}{2} \mathbb{E}_{(x,y,r) \sim D} \left[ \left( \log \frac{\pi_\theta(y|x)}{\pi_{ref}(y|x)} - \frac{1}{\tau}(r(x, y) - \hat{V}(x)) \right)^2 \right]$$

Belangrijke kenmerken van de methode zijn:

• Value-Free: Geen hulp-value netwerk wordt getraind; $\hat{V}(x)$ wordt on-the-fly berekend vanuit de referentie-policy en geobserveerde beloningen.
• Stabiele Gradiënten: De gradiënt-update hangt alleen af van de policy-parameters, waardoor backpropagation door een geleerde value estimator wordt vermeden.
• Offline Compatibiliteit: De methode is volledig offline en schaalt naar grote datasets zonder dat er online RL-loops nodig zijn.

Belangrijkste Bijdragen

1. Voorstel van RPO: Een nieuwe single-trajectory policy optimalisatiemethode die waarde-estimatie vervangt door empirische beloningsnormalisatie.
2. Analyse van DRO: Een theoretische en praktische kritiek op Direct Reward Optimization, waarbij de instabiliteit wordt belicht door gezamenlijke policy-value optimalisatie.
3. Formele Afleiding: Een bewijs dat aantoont dat RPO de optimale policy-structuur benadert door gebruik te maken van de partitiefunctie-formulering.
4. Uitgebreide Evaluatie: Omvattende experimenten over encoder-decoder (Flan-T5) en decoder-only (Mistral, LLaMA, Qwen) architecturen, die een consistente superioriteit tonen ten opzichte van SFT, KTO en DRO.

Experimentele Resultaten

De auteurs hebben RPO geëvalueerd op de UltraFeedback dataset en over meerdere benchmarks (AlpacaEval, MT-Bench, IFEval, GSM8K).

In-Domain Prestaties (UltraFeedback)

• Metrieken: RPO presteerde consequent beter dan SFT, KTO en DRO op automatische metrieken (BERTScore, ROUGE-L, Distinct-2) en veiligheid (Toxicity).
• LLM-as-a-Judge: RPO behaalde hogere winstpercentages tegenover baselines wanneer geëvalueerd door GPT-4o en Claude-3.5.
• Stabiliteit: KL-divergentie analyse toonde aan dat RPO gladde, gecontroleerde trajecten behoudt, terwijl DRO sterke divergentie-groei vertoonde (vooral in decoder-only modellen) en KTO oscillaties vertoonde.
• Efficiëntie: RPO demonstreerde snellere convergentie en lagere trainingstijd vergeleken met DRO en KTO over alle model-schalen heen.

Out-of-Distribution (OOD) Generalisatie

• RPO toonde robuuste maar taakafhankelijke verbeteringen. De winsten waren het meest uitgesproken in instruction-following benchmarks (AlpacaEval, IFEval).
• Prestaties op redeneertaken (GSM8K) en multi-turn conversaties (MT-Bench) lieten consistente maar meer gematigde verbeteringen zien, variërend per modelfamilie.

Ablatie-studies

• Temperatuur ($\tau$): De methode is robuust tegen matige smoothing ($\tau=5.0$), maar lijdt onder significante prestatievermindering bij zeer scherpe schaling ($\tau=0.1$), wat leidt tot onstabiele optimalisatie.
• Empirische Normalisatie ($\hat{V}(x)$): Het verwijderen van de prompt-niveau beloningsnormalisatie veroorzaakte een consistente degradatie in semantische gelijkenis, diversiteit en veiligheid, wat de cruciale rol ervan bevestigt bij het stabiliseren van beloningsschalen.
• Dataset Dichtheid: Prestaties nemen af naarmate het aantal completions per prompt afneemt, waarbij een 25% dichtheid-instelling zorgde voor substantiële kwaliteitsverliezen. Dit benadrukt de afhankelijkheid van de methode van diverse beloningsverdelingen per prompt voor accurate partitie-estimatie.
• Ruis-robuustheid: RPO blijft stabiel onder matige beloningsruis ($\sigma=0.1, 0.5$), maar prestaties dalen naarmate de ruisintensiteit toeneemt, hoewel het niet instort zelfs onder sterke corruptie ($\sigma=1.0$).

Betekenis en Beperkingen

Het artikel beweert dat RPO een eenvoudige, schaalbare en stabiele alternatief biedt voor bestaande alignment-methoden. Door de noodzaak van het leren van een value function te elimineren, vermindert het de optimalisatiecomplexiteit en variantie, terwijl het absolute beloningssupervisie biedt. De auteurs positioneren RPO als een veelbelovende richting voor schaalbare preferentie-optimalisatie, met name in scenario's waar paarwijze data schaars is maar scalaire feedback beschikbaar is.

Door de auteurs erkende beperkingen:

• Datavereisten: RPO vertrouwt op gegroepeerde data waarbij meerdere completions per prompt bestaan om de partitiefunctie effectief te schatten. Prestaties nemen af als de diversiteit per prompt onvoldoende is.
• Kwaliteit van Beloningen: Hoewel robuust tegen matige ruis, blijft de methode afhankelijk van de kwaliteit van de beloningsannotaties; ernstige corruptie heeft een negatieve impact op de alignment.
• Evaluatie Omvang: De huidige evaluatie richt zich primair op instruction-following en alignment benchmarks. De auteurs merken op dat bredere domeinen, meertalige settings en long-horizon redeneertaken verdere studie vereisen om de generalisatie volledig te beoordelen.