How Neural Reward Models Learn Features for Policy Optimization: A Single-Index Analysis

Probleemstelling

Reward modeling (RM) is een kritisch onderdeel bij het afstemmen van grote taalmodellen (bijvoorbeeld via RLHF), maar verschilt fundamenteel van standaard supervised prediction. In KL-geregulariseerde policy-optimalisatie wordt de geleerde rewardfunctie $r(x)$ geëxponentieerd om de ingezette policy te definiëren: $\pi_\beta(dx) \propto e^{r(x)/\beta}\pi_{ref}(dx)$. Bijgevolg hangt de downstream-waarde van het systeem niet alleen af van de voorspellingsnauwkeurigheid onder de referentieverdeling $\pi_{ref}$, maar van fouten in gebieden waar de reward hoog is en de policy zich concentreert. Dit creëert een door de policy veroorzaakte verdelingverschuiving.

Bestaande theoretische analyses van offline alignment gaan vaak uit van vaste, bekende lineaire features, waardoor de uitdaging wordt omzeild hoe neurale netwerken relevante features ontdekken (feature learning) en niet-lineaire rewardlandschappen benaderen onder dergelijke verdelingverschuivingen. Dit paper adresseert deze kloof door te analyseren hoe een twee-laags neurale netwerk features leert en een niet-lineaire rewardfunctie past wanneer het trainingsdoel exponentiële reward-weighing omvat.

Methodologie

De auteurs formuleren het probleem binnen een Gaussian Single-Index Model (SIM) setting:

• Data Generatie: Inputs $x \sim \mathcal{N}(0, I_d)$, en labels $y = r^*(x) + \zeta$, waarbij de ware reward is $r^*(x) = \sigma^*(\langle \theta^*, x \rangle)$. Hierbij is $\theta^*$ een verborgen doelrichting, $\sigma^*$ een bovenafgebonden polynoom linkfunctie, en $\zeta$ begrensd ruis.
• Model Architectuur: Een twee-laags neurale netwerk met $N$ neuronen: $r_a(x) = \frac{1}{N} \sum_{j=1}^N a_j \sigma(\langle w_j, x \rangle + b_j)$.
• Twee-fase Trainingsprotocol:
  1. Feature Recovery (Fase 1): De eerste-laags richtingen $w_j$ worden getraind via online spherical SGD met een reward-gewogen doel met temperatuur $\beta_1$. De tweede-laags coëfficiënten $a$ zijn vast bij initialisatie. Het doel is om $w_j$ te aligneren met de verborgen richting $\theta^*$.
  2. Readout Fitting (Fase 2): Zodra de features zijn hersteld, worden de richtingen $w_j$ en biases $b_j$ bevroren. De tweede-laags coëfficiënten $a$ worden gefit met behulp van gewogen ridge-regressie op nieuwe samples. Het wegingsschema gebruikt een deployment temperatuur $\beta_2$, die kan verschillen van $\beta_1$. Twee wegingsschema's worden geanalyseerd:
     • Label-gewogen: Gewichten zijn $e^{y/\beta_2}$.
     • Surrogaat-gewogen: Gewichten zijn $e^{r_{proxy}(x)/\beta_2}$, waarbij $r_{proxy}$ een bevroren proxy-reward is (bijvoorbeeld uit de feature recovery-fase).

Belangrijkste Bijdragen en Resultaten

1. Feature Learning onder Reward-geïnduceerde Verdelingverschuiving
Het paper karakteriseert de eerste-laags feature recovery onder exponentiële reward-weighing.

• Signaaltransformatie: Exponentiële weighing fungeert als een niet-lineaire labeltransformatie die het Hermite-signaal dat beschikbaar is voor de eerste laag, verandert.
• Exponentreductie: De analyse toont aan dat voor elke feature-learning temperatuur $\beta_1$ boven een dimensievrije $O(1)$ drempel, de effectieve leerflesnek reduceert van de Information Exponent (IE) naar de Generative Exponent (GE).
• Herstelgaranties: Onder aannames van bovenafgebonden polynomen, herstelt een constante fractie van de neuronen de verborgen richting $\theta^*$ met hoge waarschijnlijkheid. De complexiteit van zwak herstel wordt beheerst door de generative exponent en schaalt als $\tilde{O}(\beta_1^{2(I^*-1)} d^{(p_{gen}-1)\vee 1})$, waarbij $I^*$ de macht is die nodig is om de GE te bereiken.
• Betekenis van $\beta_1$: Cruciaal is dat de vereiste ondergrens voor $\beta_1$ onafhankelijk is van de dimensie $d$. Dit verbetert eerdere analyses met afgeknipte exponentiële functies die dimensie-afhankelijke polylogaritmische temperaturen vereisten, wat anders de voordelen van exponentreductie zou tenietdoen.

2. Van Gewogen Voorspelling naar Tilted-Policy Regret
De auteurs gaan verder dan standaard foutgrenzen voor voorspelling om de tilted-policy value gap te evalueren:
$$\mathcal{R} = \mathbb{E}_{x \sim \pi_{\beta^*}}[r^*(x)] - \mathbb{E}_{x \sim \hat{\pi}_{\beta_2}}[r^*(x)]$$
waarbij $\pi_{\beta^*}$ de doelpolicy is en $\hat{\pi}_{\beta_2}$ de policy die wordt geïnduceerd door de geleerde reward.

• Decompositie: De regret wordt ontbonden in drie termen:
  1. Temperatuur Mismatch ($T_{temp}$): Hangt af van het verschil tussen de doeltemperatuur $\beta^*$ en de deployment temperatuur $\beta_2$.
  2. Geprojecteerde Truncatie ($T_{cut}$): Ontstaat door het beperken van de analyse tot een geprojecteerde deelruimte om integreerbaarheid te waarborgen.
  3. Leerterm ($T_{learn}$): Hangt af van de gewogen voorspellingsfout van de gefitte reward.
• Gedrag bij Lage Temperatuur: Het paper leidt expliciete grenzen af die aantonen dat de variantie van de reward onder de tilted policy schaalt als $\beta_2^2 / (2p_{max})$, waarbij $p_{max}$ de lokale orde is van het reward-maximum.

3. Label- en Surrogaat-gewogen Regretgrenzen
Het paper biedt expliciete value-gap grenzen voor beide fitting-schema's:

• Label-gewogen Grens: Dient als theoretisch referentiepunt. De grens scheidt temperatuur mismatch, truncatie en leerkosten, en koppelt policy value rates aan de lokale vorm van het niet-lineaire reward-maximum.
• Surrogaat-gewogen Grens: Praktischer voor plug-in pipelines. Deze grens introduceert een extra factor $C_{0,R}(\beta_2)$, die de dichtheidsverhouding vertegenwoordigt tussen de door het surrogaat geïnduceerde maat en de referentieverdeling.
  • Proxy Mismatch Mechanisme: Wanneer de proxy-reward imperfect is, kan $C_{0,R}(\beta_2)$ groot zijn (schalend polynomiëel met $1/\beta_2$). Dit versterkt zowel benaderingsfouten als fouten door eindige steekproeven, en identificeert een statistisch mechanisme achter "proxy mismatch" en reward overoptimization.
• Toelaatbare Temperaturen: De analyse levert een toelaatbare set van deployment temperaturen $\beta_2$ op. Het kiezen van een kleinere $\beta_2$ verbetert de overeenkomst met de doelpolicy, maar versterkt de leerkosten via exponentiële weighing. De optimale keuze balanceert deze trade-offs.

Betekenis en Claims

Het paper claimt een fundamentele karakterisering te bieden van de interactie tussen reward modeling en policy value binnen een gecontroleerde theoretische setting.

• Feature Learning: Het toont aan dat reward-geïnduceerde exponentiële weighing feature learning kan faciliteren door de effectieve leern exponent te reduceren van IE naar GE, mits de temperatuur $\beta_1$ niet te klein is. Het stelt vast dat een dimensievrije $O(1)$ temperatuur voldoende is voor herstel, en corrigeert hiermee eerdere aannames die dimensie-afhankelijke temperaturen vereisten.
• Evaluatiemetric: Het betoogt dat de relevante statistische doelstelling voor reward modeling de tilted-policy value is, en niet de voorspellingsfout onder de referentieverdeling. De afgeleide grenzen koppelen de value gap expliciet aan de lokale geometrie van het rewardlandschap en de keuze van de deployment temperatuur.
• Surrogaat Risico's: Het identificeert een specifiek statistisch mechanisme waarbij het gebruik van een proxy-reward voor weging (surrogaat-weighing) fouten kan versterken als de proxy imperfect is, en biedt een theoretische verklaring voor reward overoptimization fenomenen die in de praktijk worden waargenomen.

De auteurs handhaven dat hoewel het Gaussian SIM een theoretische abstractie is, het de kernmechanismen van feature learning en door policy veroorzaakte verdelingverschuivingen isoleert, en een basis biedt voor toekomstig werk aan rijker inputdistributies en multi-index structuren.