Technische Samenvatting: Alignment-Aware Decoding (AAD)

Probleemstelling

Het afstemmen (alignen) van Large Language Models (LLM's) op menselijke voorkeuren blijft een centrale uitdaging in de natuurlijke taalverwerking. Hoewel methoden tijdens de training, zoals Reinforcement Learning from Human Feedback (RLHF) en Direct Preference Optimization (DPO), effectief zijn gebleken, kennen zij aanzienlijke beperkingen. Deze methoden zijn gevoelig voor imperfecte voorkeurssignalen, gevoelig voor overoptimalisatie (reward hacking), en resulteren vaak in beleid dat de biases van het referentiemodel dat tijdens de training werd gebruikt, overneemt. Bovendien is de afstemming van een model zodra het getraind is, vastgelegd; het wijzigen ervan vereist computationeel dure hertraining of complexe interventies tijdens de inferentie die vaak vertrouwen op hulp-beloningsmodellen (auxiliary reward models), uitgebreide zoekprocedures of zorgvuldig afgestelde hyperparameters.

Methodologie: Alignment-Aware Decoding (AAD)

Het artikel introduceert Alignment-Aware Decoding (AAD), een methode die ontworpen is om de modelafstemming direct tijdens de inferentie te verbeteren zonder de modelparameters aan te passen of extra training te vereisen. AAD werkt op basis van het theoretische inzicht dat een met DPO getraind model ($\pi_{DPO}$) impliciet een beloningssignaal codeert ten opzichte van zijn supervised fine-tuning (SFT) referentiemodel ($\pi_{SFT}$).

Kernmechanisme

AAD maakt gebruik van twee onderscheidende kenmerken van het met DPO afgestemde model:

1. Token-waarschijnlijkheden: Het vermogen om veilige kandidaat-tokens te identificeren voor de volgende decoderingsstap.
2. Token-niveau Credit Assignment: Het vermogen om het voordeel (advantage) van een token te schatten met behulp van de log-waarschijnlijkheidsratio tussen het afgestemde model en het referentiemodel.

In plaats van de waarschijnlijkheidsverdeling van $\pi_{DPO}$ direct te volgen, selecteert AAD tokens door een token-niveau voordeelfunctie te maximaliseren:
$$A(v | x \circ y_{1:t}) = \log \frac{\pi_{DPO}(v | x \circ y_{1:t})}{\pi_{SFT}(v | x \circ y_{1:t})}$$
Deze ratio behandelt het afgestemde model effectief als een token-niveau beloningsfunctie, waarbij wordt geoptimaliseerd voor de impliciete beloning die tijdens de voorkeursoptimalisatie is vastgelegd.

Voorkomen van Overoptimalisatie

Om het probleem van overoptimalisatie aan te pakken, waarbij het maximaliseren van de voordeelfunctie mogelijk lage-waarschijnlijkheidstokens selecteert die leiden tot degeneratieve outputs, bevat AAD een min-$\alpha$ filtering mechanisme. Voordat de token met het hoogste voordeel wordt geselecteerd, beperkt de methode de zoekruimte tot tokens die voldoen aan een minimale waarschijnlijkheidsdrempel onder het afgestemde model:
$$V_\alpha = \{ v \in \mathcal{V} \mid \pi_{DPO}(v | x \circ y_{1:t-1}) \geq \alpha \cdot \max_{v' \in \mathcal{V}} \pi_{DPO}(v' | x \circ y_{1:t-1}) \}$$
De uiteindelijke token $y_t$ wordt gekozen als:
$$y_t = \arg \max_{v \in V_\alpha} A(v | x \circ y_{1:t-1})$$
Dit zorgt ervoor dat het model alleen optimaliseert voor afstemming over tokens die plausibel en vloeiend zijn volgens het afgestemde beleid.

Belangrijkste Bijdragen

1. Afstemming tijdens de Inferentie: Het artikel stelt AAD voor, een eenvoudige, training-vrije methode die de afstemming verbetert tijdens de inferentie door het DPO-model te gebruiken als een token-beloningsfunctie en het SFT-model als referentie.
2. Empirische Superieuriteit: Uitgebreide experimenten over diverse benchmarks (Ultrafeedback, Argilla, Skywork, Nectar, etc.) en model-schalen (0.6B tot 13B parameters) tonen aan dat AAD consequent beter presteert dan sterke baselines, inclus_ief Greedy DPO, Best-of-2 (Bo2) en Emulated Fine-Tuning (EFT), onder rekenkundig equivalente condities.
3. Data-efficiënte Synthetische Datageneratie: De auteurs demonstreren dat AAD hoogwaardige synthetische voltooiingen kan genereren. Deze voltooiingen kunnen worden gebruikt in iteratieve DPO-loops om de modelafstemming verder te verbeteren, met name in scenario's waar gelabelde voorkeursdata schaars is.

Resultaten

• Benchmark Prestaties: AAD bereikt de hoogste gemiddelde oracle-beloning en winrates over alle geteste datasets en modelfamilies. Zo behaalde AAD op de Skywork-dataset met een 3B LLaMA-model een beloningsscore van 19.27 vergeleken met 13.64 voor Greedy DPO.
• Robuustheid tegen Dataschaarste: In experimenten waarbij de omvang van de trainingsdataset varieerde, presteerde AAD consistent beter dan Best-of-2 sampling, zelfs wanneer het getraind was op slechts 10% van de beschikbare voorkeursdata.
• Menselijke Evaluatie: Een menselijke evaluatiestudie op de Skywork-dataset toonde aan dat door AAD gegenereerde reacties werden verkozen boven Greedy DPO (58.7% winrate) en Best-of-2 (73.5% winrate), met een Elo-rating van 1610.1, wat de andere baselines aanzienlijk overtreft.
• Iteratieve Verbetering: Het gebruik van door AAD gegenereerde data voor iteratieve DPO stelde modellen die getraind waren op slechts 10% van de data in staat om bijna het prestatieverschil te dichten met modellen die getraind waren op de volledige dataset.
• Beperkingen: De methode vereist twee forward passes per token (één voor $\pi_{DPO}$ en één voor $\pi_{SFT}$), wat resulteert in ongeveer de helft van de throughput van standaard greedy decoding op een enkele GPU. Daarnaast merken de auteurs op dat AAD specifiek is ontwor de SFT+DPO-pipeline; het combineren ervan met PPO-getrainde modellen leverde gemengde resultaten op, omdat PPO de specifieke waarschijnlijkheidsafstand niet behoudt die AAD exploiteert.

Betekenis

Het artikel positioneert AAD als een praktische en theoretisch gefundeerde oplossing voor het verbeteren van LLM-afstemming zonder de computationele overhead van hertraining of de complexiteit van externe beloningsmodellen. Door het DPO-proces te interpreteren als het leren van een Q-functie en het referentiemodel te gebruiken om geërfde biases te mitigeren, biedt AAD een "plug-and-play" decoderingsstrategie die veiligheid en behulpzaamheid verhoogt. De auteurs stellen dat deze aanpak bijzonder waardevol is in scenario's waarin modelgewichten bevroren, eigendom van derden, of wanneer hoogwaardige voorkeursdata beperkt is, wat een pad biedt naar sterkere afstemming via efficiënte optimalisatie tijdens de inferentie en synthetische datageneratie.