Technische Samenvatting: Veilige Inferentie-tijd Alignment via Lagrangian Reward Augmentation (LARA)

Probleemstelling

Reinforcement Learning from Human Feedback (RLHF) is het standaardkader voor het afstemmen (alignment) van Large Language Models (LLM's) op menselijke voorkeuren. Traditionele alignment vertrouwt echter op policy fine-tuning (bijv. via RLHF of Direct Preference Optimization zoals DPO), wat computationeel duur is en hertraining vereist wanneer veiligheidseisen of gebruikerspopulaties veranderen.

Hoewel inferentie-tijd alignment-methoden (zoals Best-of-N reranking of token-niveau reward-gestuurde decoding) zijn opgekomen om bevroren modellen te sturen met behulp van aanvullende reward-signalen zonder gewicht-updates, missen zij momenteel een principieel mechanisme voor veilige alignment. Bestaande methoden optimaliseren doorgaans een enkele scalaire score. Wanneer doelstellingen voor behulpzaamheid (helpfulness) en veiligheid (harmlessness) conflicteren, moeten praktijkgebruikers de afweging tussen beide handmatig afstemmen. Bovendien handhaven deze methoden inherent geen expliciete veiligheidsrestricties (bijv. een budget op de verwachte kosten) tijdens de decoding, wat leidt tot ofwel onveilige outputs, ofwel onnodige reducties in behulpzaamheid.

Methodologie: Lagrangian Reward Augmentation (LARA)

De auteurs stellen LARA voor, een decoder-agnostisch kader dat de geconstelde optimalisatiestructuur van Safe RLHF overbrengt naar de inferentie-fase. In plaats van de modelgewichten bij te werken, kalibreert LARA een scalaire duale variabele om een uitgebreid (augmented) reward-signaal te creëren.

1. Theoretische Formulering

De methode begint bij de Safe RLHF-doelstelling, die de verwachte behulpzaamheid ($r$) maximaliseert onder een verwachte kostenrestrictie ($c \le \tau$), geregulariseerd door een KL-divergentie-term tegen een referentiepolicy ($\pi_{ref}$):
$$\max_{\pi} \mathbb{E}[r(x,y)] - \beta D_{KL}(\pi \parallel \pi_{ref}) \quad \text{onder de voorwaarde dat} \quad \mathbb{E}[c(x,y)] \le \tau$$

Gebruikmakend van Lagrangiaanse dualiteit, dualiseert de auteurs de veiligheidsrestrictie. Dit transformeert het oneindig-dimensionale policy-optimalisatieprobleem naar een eendimensionaal convex optimalisatieprobleem over een niet-negatieve duale variabele $\lambda$ (de Lagrange-multiplier).

• Augmented Reward: De duale variabele definieert een uitgebreide reward: $r_\lambda(x, y) = r(x, y) - \lambda c(x, y)$.
• Duale Doelstelling: Er wordt aangetoond dat de duale functie $g(\lambda)$ convex is. De afgeleide komt overeen met het verschil tussen het veiligheidsbudget $\tau$ en de verwachte kosten onder de Gibbs-tilted policy geïnduceerd door $r_\lambda$.
• Optimaliteit: De optimale $\lambda^*$ voldoet aan de complementariteitsvoorwaarde (complementary slackness), wat ervoor zorgt dat de verwachte kosten gelijk zijn aan het budget $\tau$ (indien de restrictie actief is) terwijl de behulpzaamheid wordt gemaximaliseerd.

2. Empirische Kalibratie

LARA vereist geen analytische oplossing van de optimalisatie. In plaats daarvan schat het $\lambda^*$ in via empirische kalibratie op een kleine dataset:

1. Data: Een kleine set prompts wordt gesampled, en meerdere responsen worden gegenereerd vanuit de bevroren referentiepolicy.
2. Schatting: De duale variabele wordt geschat door de empirische duale doelstelling te minimaliseren over een compact interval $[0, \Lambda]$.
3. Algoritme: Vanwege de convexiteit en monotoniciteit van de duale gradiënt, kan de optimale $\lambda$ efficiënt worden gevonden met behulp van binair zoeken.
4. Robuustheid: Om discrepanties tussen kalibratie en deployment aan te pakken, suggereren de auteurs het bootstrappen van de kalibratieset en het selecteren van een bovenste kwantiel van de geschatte $\lambda$-distributie als een conservatieve keuze.

3. Integratie met Decoding

Zodra gekalibreerd, wordt de scalaire $\hat{\lambda}$ gebruikt om de uitgebreide reward $r_{\hat{\lambda}}$ te berekenen. Deze score wordt als een "drop-in" signaal doorgegeven aan bestaande inferentie-tijd alignment-methoden:

• Sequentie-niveau (bijv. Best-of-N): De gekalibreerde $\lambda$ levert een exacte oplossing voor het verwachte-kosten-geconstelde probleem.
• Token-niveau (bijv. Reward-Guided Search): De gekalibreerde $\lambda$ biedt een principieel heuristiek afgeleid van de geconstelde doelstelling, hoewel dit geen exacte populatie-niveau restrictie-naleving garandeert vanwege het hebzuchtige (greedy) karakter van token-niveau decoding.

Belangrijkste Bijdragen

1. Formulering: Het artikel formuleert veilige inferentie-tijd alignment als het duale van de Safe RLHF-doelstelling, waardoor het probleem wordt teruggebracht tot een eendimensionale convexe kalibratietaak.
2. Geometrie & Efficiëntie: De auteurs karakteriseren de geometrie van de duale doelstelling, waarbij ze de convexiteit en de monotoniciteit van de gradiënt bewijzen. Dit maakt een efficiënt binair zoekproces voor kalibratie mogelijk.
3. Theoretische Garanties: Er wordt een eindige-steekproef garantie verstrekt voor de estimator $\hat{\lambda}_{N,K}$, die aantoont dat de fout convergeert rond de populatie-minimizer $\lambda^*$, gecontroleerd door het aantal prompts ($N$) en samples per prompt ($K$).
4. Decoder-Agnostisch Kader: LARA is ontworpen als een kalibratielaag die compatibel is met elke bestaande inferentie-tijd alignment-methode, waarbij handmatige scalaire afstemming wordt vervangen door een principieel geconstelde constructie.

Experimentele Resultaten

De auteurs evalueerden LARA op Llama-3.2-3B en Qwen-2.5-3B modellen met behulp van de BeaverTails dataset.

• Restrictie-naleving: Experimenten bevestigden dat de gekalibreerde $\lambda^*$ zich op de grens van de veiligheidsrestrictie bevindt. Waarden van $\lambda < \lambda^*$ overtraden het veiligheidsbudget, terwijl $\lambda > \lambda^*$ de restrictie respecteerde maar onnodig de behulpzaamheid verminderde.
• Prestaties vs. Baselines:
  • Inferentie-tijd: Onder de inferentie-tijd methoden bereikte Best-of-N (BoN) in combinatie met LARA de beste afweging tussen behulpzaamheid en onschadelijkheid. Het presteerde beter dan SafeDPO op behulpzaamheid en bleef competitief op onschadelijkheid, alles zonder gewicht-updates.
  • Token-niveau: Token-niveau methoden (ARGS, PARGS, CD-Q) presteerden over het algemeen slechter op behulpzaamheid vergeleken met BoN, hoewel CD-Q sterke verbeteringen in onschadelijkheid liet zien met toenemende $\lambda$.
  • vs. Fine-tuning: Hoewel LARA-gestuurde BoN de prestaties van fine-tuning gebaseerde baselines (Safe-RLHF, SafeDPO) benaderde, bereikten de fine-tuned methoden nog steeds de overall beste afweging. LARA biedt echter een aanzienlijk voordeel in flexibiliteit bij deployment en kosten.
• Evaluatie: De resultaten waren consistent over model-gebaseerde evaluaties (met behulp van BeaverTails reward/cost modellen) en GPT-4o-mini gebaseerde evaluaties, hoewel GPT minder gevoelig was voor $\lambda$-variaties dan de reward-modellen.

Betekenis en Claims

Het artikel claimt dat LARA een principieel alternatief biedt voor handmatige penalty-afstemming voor veilige inferentie-tijd alignment. Door de veiligheidsrestrictie te dualiseren, vertaalt LARA een veiligheidsbudget naar een decoding-doelstelling zonder de modelgewichten te wijzigen.

• Voor Sequentie-niveau Samplers: LARA biedt een exacte verwachting van de kosten, wat theoretisch de optimale afweging garandeert onder de geconstelde doelstelling.
• Voor Token-niveau Decoders: Het biedt een principieel heuristiek die de geometrie van de geconstelde doelstelling respecteert, zelfs als het de populatie-restrictie niet strikt kan afdwingen.

De auteurs concluderen dat LARA er succesvol in is geslaagd om veiligheidsrestricties over te dragen van training-tijd alignment naar praktische controle tijdens de decoding, wat de afweging tussen behulpzaamheid en onschadelijkheid verbetert en inferentie-tijd methoden dichter bij de prestaties van fine-tuning gebaseerde benaderingen brengt. Het werk benadrukt dat veiligheidsrestricties effectief beheerd kunnen worden tijdens de inferentie-tijd, waardoor de computationele overhead van herhaalde hertraining wordt vermeden.