Technische Samenvatting: Wanneer afstand afleidt: Representatieafstand-bias in BT-loss voor beloningsmodellen

1. Probleemstelling

Beloningsmodellen (Reward Models, RM's) zijn centraal in Reinforcement Learning from Human Feedback (RLHF) en dienen als de proxy voor menselijke voorkeuren die de afstemming van Large Language Models (LLM's) sturen. De standaard doelfunctie voor het trainen van deze modellen is de Bradley-Terry (BT) loss, die leert van gepaarde data bestaande uit "gekozen" ($y_w$) en "afgewezen" ($y_l$) responsen.

De auteurs identificeren een kritiek gebrek in de update-dynamiek van de standaard BT-loss: Representatieafstand-bias. Hoewel de BT-gradiënt bedoeld is om gedreven te worden door de voorspellingsfout (de omvang van de fout in de rangschikking), onthult de analyse dat de gradiëntnorm wordt bepaald door twee onderscheidende componenten:

1. Voorspellingsfout: Het verschil in voorspelde beloningen tussen de gekozen en de afgewezen responsen.
2. Representatieafstand: De afstand tussen de twee responsen in de outputruimte van de laatste laag ($\|h_w - h_l\|$).

Deze koppeling creëert een vals leersignaal waarbij de omvang van de update wordt bepaald door hoe ver de representaties van het paar uit elkaar liggen, in plaats van uitsluitend door hoe foutief de voorspelling is.

• Paren met kleine afstand: Paren die semantisch of structureel zeer vergelijkbaar zijn (bijv. een correct codefragment versus één met een enkele logische fout) hebben vaak kleine representatieafstanden. Onder standaard BT-loss ontvangen deze paren verdwijnend zwakke updates, zelfs wanneer het model ze verkeerd rangschikt. Dit is problematisch omdat fijne onderscheidingen vaak het meest cruciaal zijn voor taken zoals redeneren en coderen.
• Paren met grote afstand: Paren die zeer verschillend zijn (bijv. een weigering versus een antwoord op een schadelijke query) hebben van nature grote representatieafstanden. Deze paren ontvangen disproportioneel sterke updates, wat potentieel de leersignalen van de subtielere, fijnmazige paren kan overschaduwen.

Gevolgen daarvan traint het model onvoldoende op de moeilijke, fijnmazige onderscheidingen waar het het meest moeite mee heeft, terwijl het te veel de nadruk legt op oppervlakkige verschillen.

2. Methodologie: NormBT

Om deze bias aan te pakken, stellen de auteurs NormBT voor, een adaptief paar-normalisatieschema dat de gradiëntbijdrage van elk voorkeurspaar herschaalt.

Theoretische basis

De auteurs deconstrueren de BT-gradiëntnorm $\|\nabla_\theta L_{BT}\|$ en tonen aan dat deze proportioneel is aan:
$$\|\nabla_\theta L_{BT}\| \approx |\sigma(d) - 1| \cdot \|h_w - h_l\|$$
waarbij $d$ het verschil in beloning is en $\|h_w - h_l\|$ de representatieafstand is. Om de omvang van de update primair afhankelijk te maken van de voorspellingsfout-term $|\sigma(d) - 1|$, introduceren de auteurs een gewicht $w_i$ dat invers schaalt met de representatieafstand.

Implementatiedetails

Het direct berekenen van de exacte gradiëntnorm voor weging is computationeel duur. In plaats daarvan gebruikt NormBT de representatieafstand van de laatste laag ($\|h_w - h_l\|$) als een proxy, die is aangetoond een exacte proxy te zijn voor de lineaire scorekop en een sterke benadering voor het volledige model.

De voorgestelde verliesfunctie is:
$$\mathcal{L}_{NormBT}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \tilde{w}(y_w, y_l) \cdot \log \sigma(r_w - r_l) \right]$$

Waar het gewicht $\tilde{w}$ wordt gedefinieerd als:
$$\tilde{w}(y_w, y_l) = \frac{\mu_t}{\|h_w - h_l\| + \epsilon}$$

Belangrijke ontwerpkeuzes zijn onder meer:

• Stop-Gradient: De afstandsterm wordt berekend met een stop-gradient operator om te voorkomen dat de normalisatie-weging zelf wordt geoptimaliseerd.
• EMA-stabilisatie: Om schaaldrift tijdens de training te voorkomen (omdat de magnitudes van embeddings kunnen verschuiven), gebruiken de auteurs een Exponential Moving Average (EMA) van de gemiddelde representatieafstand per batch ($\mu_t$) om de gewichten te normaliseren. Dit zorgt ervoor dat de gewichten stabiel blijven en adaptief zijn aan distributieverschuivingen.
• Lichtgewicht Integratie: NormBT vereist geen architecturale wijzigingen, behoudt de probabilistische basis van BT en brengt verwaarloosbare computationele overhead met zich mee.

3. Belangrijkste bijdragen

1. Gradiëntanalyse: Het artikel biedt een gedetailleerde theoretische en empirische analyse die aantoont dat BT-gradiëntnormen gezamenlijk worden bepaald door de voorspellingsfout en de representatieafstand, wat leidt tot bevoordeelde leersignalen.
2. Identificatie van Bias: De auteurs demonstreren dat deze koppeling ervoor zorgt dat paren met een kleine afstand (vaak de meest kritieke voor fijnmazige taken) zwakke updates ontvangen, terwijl paren met een grote afstand de leersignalen domineren.
3. NormBT Voorstel: Ze introduceren NormBT, een "drop-in" modificatie die updates normaliseert naar de representatieafstand, waardoor leersignalen zich richten op de voorspellingsfouten.
4. Empirische Validatie: Uitgebreide experimenten over verschillende LLM-backbones (Gemma-2b-it, Llama-3.2-3B) en datasets (Unified-Feedback, Skywork-Reward) tonen consistente prestatieverbeteringen aan.

4. Experimentele Resultaten

De auteurs evalueerden NormBT tegenover de standaard BT-baseline en verschillende varianten (BT + Margin, BT + Label Smoothing) op de RewardBench benchmark en downstream Best-of-N (BoN) selectietaken.

• RewardBench Prestaties: NormBT presteerde consequent beter dan de BT-baseline in alle settings. De meest significante winst werd waargenomen in de categorie Redeneren (Reasoning), waar de nauwkeurigheid gemiddeld met meer dan 5% verbeterde. Dit komt overeen met de hypothese dat redeneertaken steunen op fijnmazige onderscheidingen (kleine representatieafstanden) die door standaard BT worden onderdrukt.
• Downstream RLHF (BoN): In Best-of-N selectie-experimenten behaalden modellen getraind met NormBT hogere "gold scores" (kwaliteit van geselecteerde responsen) vergeleken met BT-baselines.
• Ablatie-studies:
  • NormBT presteerde beter dan andere trainingstechnieken zoals LayerNorm en Gradient Clipping, die er niet in slaagden het relatieve wegingprobleem tussen paren aan te pakken.
  • De methode bleek complementair: het verbeterde de prestaties zelfs wanneer het werd toegepast bovenop andere BT-varianten (bijv. NormBT + Margin).
  • Het verwijderen van de EMA-stabilisatie of het gebruik van alternatieve gelijkenis-metrieken (cosine similarity, average-pooled embeddings) resulteerde in prestatievermindering, wat de specifieke ontwerpkeuzes valideert.

5. Betekenis en Claims

Het artikel claimt dat NormBT een geprincieerde en lichtgewicht oplossing biedt voor een structurele beperking in de standaard BT-loss. Door de update-omvang te ontkoppelen van de representatieafstand, zorgt NormBT ervoor dat het beloningsmodel getrouwer leert van fijnmazige onderscheidingen, die vaak de meest uitdagende maar cruciale zijn voor hoogwaardige afstemming.

De auteurs benadrukken dat hun werk niet direct gericht is op de kwaliteit van de data of de bias van annotatoren, maar de efficiëntie en getrouwheid van de leersignalen verbetert die uit bestaande gepaarde data worden geëxtraheerd. Zij positioneren NormBT als een fundamentele verbetering voor het paradigma van discriminatieve beloningsmodellering, wat potentieel leidt tot betrouwbaardere beloningssignalen in downstream RLHF-toepassingen, met name in domeinen die precieze redenering en subtiele foutdetectie vereisen.