Inference-time Alignment in Continuous Space

Dit paper introduceert Simple Energy Adaptation (SEA), een efficiënt algoritme voor inferentie-tijd uitlijning van taalmodellen dat door middel van gradiëntgebaseerde bemonstering in een continue latente ruimte direct antwoorden aanpast, waardoor het de beperkingen van discrete zoekmethodes overtreft en aanzienlijk betere prestaties levert.

De kern

Stel je voor dat je een zeer slimme, maar soms wat onvoorspelbare robot hebt die verhalen kan schrijven of vragen kan beantwoorden. Deze robot is getraind op een enorme hoeveelheid tekst, maar hij heeft nog niet geleerd wat we mensen echt belangrijk vinden: wat is veilig, wat is eerlijk en wat is slim?

In de wereld van kunstmatige intelligentie noemen we dit proces "aligneren" (op één lijn brengen). De vraag is: hoe maken we deze robot veilig en behulpzaam zonder hem maandenlang opnieuw te trainen?

Dit paper introduceert een nieuwe, slimme manier om dat te doen, genaamd SEA (Simple Energy Adaptation). Hier is hoe het werkt, vertaald naar alledaagse taal:

1. Het oude probleem: "Zoek in een doos met losse blokken"

Tot nu toe probeerden mensen de robot te verbeteren door een trucje te gebruiken dat lijkt op Best-of-N (de beste van N).

• Hoe het werkte: Je vraagt de robot om 10 of 100 verschillende antwoorden te bedenken. Dan laat je een strenge juf (een "beloningsmodel") al die antwoorden beoordelen. De juf kiest het beste antwoord uit die hoop.
• Het probleem: Stel dat de robot erg dom is of dat je maar een klein aantal antwoorden hebt. Dan is de kans heel groot dat geen enkel van die antwoorden goed genoeg is. Het is alsof je probeert een gouden munt te vinden in een emmer met modder, maar je gooit er maar een paar keer in. Als de emmer klein is of de modder heel diep, vind je de munt nooit.

2. De nieuwe oplossing: "De GPS-geleide wandeling"

De auteurs van dit paper zeggen: "Waarom zoeken we niet gewoon in de richting van het goede antwoord, in plaats van blindelings te gooien?"

Ze introduceren SEA. In plaats van losse antwoorden te verzamelen, laten ze de robot een wandeling maken in een continue ruimte van gedachten.

• De Analogie: Stel je voor dat je in een donker landschap staat (de ruimte van alle mogelijke antwoorden). Je doel is een helder, warm licht (het perfecte, veilige antwoord).
  • De oude methode (Best-of-N): Je gooit 100 kaarten in het donker en hoopt dat één kaart op het licht valt.
  • De SEA-methode: Je hebt een kompas (de beloning) dat je vertelt waar het licht is. Je begint met een willekeurig punt en loopt stap voor stap in de richting van het licht. Je corrigeert je koers continu. Je hoeft niet te gokken; je stapt gewoon naar het goede antwoord toe.

3. Hoe werkt SEA precies? (De "Energie" en de "Stap")

Het paper gebruikt wat wiskundige termen, maar het idee is simpel:

1. Energie: Ze noemen het goede antwoord "lage energie" en het slechte antwoord "hoge energie". Het doel is om de energie te verlagen.
2. De Ladder: De robot begint met een ruw antwoord. Dan kijkt hij naar de "helling" van het landschap (de gradient). Hij ziet: "Ah, als ik dit woord iets anders maak, wordt het antwoord veiliger."
3. Stap voor stap: De robot past het antwoord langzaam aan, alsof hij een trap afloopt naar beneden, totdat hij bij de beste, veiligste oplossing is.

Waarom is dit zo cool?

• Het werkt zelfs als de robot dom is: Zelfs als de robot niet slim genoeg is om direct het perfecte antwoord te bedenken, kan hij er toch naartoe "wandelen" door de aanwijzingen van het kompas te volgen.
• Het is dieper: Oude methoden kijken vaak alleen naar het begin van een zin. Als de robot begint met "Natuurlijk, hier is hoe je een bom maakt...", is het vaak te laat. SEA kijkt naar het hele antwoord en kan het hele verhaal veilig maken, zelfs als het begin wat twijfelachtig was.
• Het is sneller en slimmer: In plaats van duizenden antwoorden te genereren en er één te kiezen, bouwt SEA één antwoord op dat steeds beter wordt.

Samenvatting in één zin

In plaats van te hopen dat je toevallig het juiste antwoord vindt door veel te proberen (zoals een visser die 100 haken in het water gooit), gebruikt SEA een visser met een sonar die de vis precies volgt en hem netjes in het net leidt.

Dit maakt AI veiliger, eerlijker en slimmer, zonder dat we de robot opnieuw hoeven te "leren" (trainen), maar gewoon door hem tijdens het gesprek een beetje bij te sturen.

Technische samenvatting

Titel: Inference-time Alignment in Continuous Space

Auteurs: Yige Yuan, Teng Xiao, et al. (ICT, CAS; University of Washington; Allen Institute for AI; Alibaba Group)

---

1. Het Probleem

Het uitlijnen van grote taalmodellen (LLMs) met menselijke voorkeuren is cruciaal voor veiligheid en bruikbaarheid. Bestaande methoden voor inference-time alignment (uitlijning tijdens het genereren, zonder extra training) vertrouwen voornamelijk op het genereren van meerdere antwoorden en het selecteren van het beste antwoord op basis van een beloningsmodel (Reward Model).

De belangrijkste beperkingen van deze bestaande benaderingen zijn:

• Discrete Zoekruimte: Methoden zoals Best-of-N (BoN) opereren in een discrete ruimte. Ze genereren $N$ kandidaat-antwoorden en kiezen de beste.
• Afhankelijkheid van het Basismodel: Als het basismodel zwak is of de kans op een goed antwoord laag is, moet $N$ exponentieel groeien om een goed resultaat te vinden. Dit is computatief zwaar en vaak inefficiënt.
• Beperkte Exploratie: Deze methoden kunnen geen nieuwe, informatieve kandidaten "ontdekken" die niet door het basismodel zijn gegenereerd. Ze zijn beperkt tot de kansverdeling van het basismodel.
• Shallow Alignment: Veiligheidsuitlijning gebeurt vaak alleen in de eerste paar tokens (de "prefix"), waardoor modellen kwetsbaar blijven voor aanvallen die de eerste tokens manipuleren (bijv. Prefilling Attacks).

2. Methodologie: Simple Energy Adaptation (SEA)

De auteurs stellen Simple Energy Adaptation (SEA) voor, een algoritme dat het paradigma verschuift van "zoeken in een discrete ruimte" naar "optimalisatie in een continue ruimte".

Kernconcepten:

• Energy-Based Model (EBM): De optimale RLHF-beleid (policy) wordt geformuleerd als een energie-functie $E(x, y)$ over de logits (soft outputs) van het model:
  $$E(x, y) = \log \pi_{ref}(y|x) + \alpha r(x, y)$$
  Waarbij $\pi_{ref}$ het referentie-model is, $r(x, y)$ de beloning is, en $\alpha$ een regularisatie-coëfficiënt.
• Continue Latente Ruimte: In plaats van te werken met discrete tokens, werkt SEA direct met de continue logits van het model. Dit maakt de ruimte differentieerbaar.
• Langevin Dynamics: SEA gebruikt een iteratief optimalisatieproces gebaseerd op Langevin Dynamics om de initialisatie (het oorspronkelijke antwoord van het basismodel) te verfijnen. De update-regel is:
  $$y^{(n+1)} \leftarrow y^{(n)} - \eta \nabla_y E(x, y^{(n)}) + \epsilon^{(n)}$$
  Hierbij is $\nabla_y E$ de gradiënt van de energie (afgeleid van de beloningsgradiënt en de log-probabiliteit van het referentiemodel) en $\epsilon$ Gaussisch ruis.
• Proces:
  1. Het basismodel genereert een initiële respons (logits).
  2. De logits worden iteratief aangepast via gradiëntafstijging om de energie te minimaliseren (d.w.z. de beloning te maximaliseren).
  3. Na $N$ stappen wordt de finale continue logits-reeks gedecodeerd naar discrete tekst.

Voordelen ten opzichte van discrete zoekmethoden:

• Diepe Uitlijning: Omdat alle tokens gelijktijdig worden geoptimaliseerd binnen een globaal receptief veld, kan SEA "herstellen" van schadelijke startcondities, in plaats van alleen de eerste tokens te beïnvloeden.
• Efficiëntie: Het vermijdt de noodzaak om duizenden kandidaten te genereren (zoals bij BoN) en zoekt direct in de richting van de optimale oplossing via gradiënten.

3. Belangrijkste Bijdragen

1. Nieuw Paradigma: Introductie van een continue optimalisatie-approach voor inference-time alignment, in tegenstelling tot de dominante discrete zoekmethoden.
2. Algoritme SEA: Een eenvoudig maar krachtig algoritme dat gradiëntgebaseerde sampling (Langevin Dynamics) toepast op de logits van een LLM, geleid door een beloningsmodel.
3. Oplossing voor "Shallow Alignment": Bewijs dat SEA diepere veiligheidsuitlijning bereikt door de hele respons te optimaliseren, wat kwetsbaarheden voor Prefilling Attacks effectief elimineert.
4. Efficiëntie: SEA bereikt betere resultaten met minder rekenkracht dan methoden die afhankelijk zijn van grote kandidaatsets (zoals BoN-64 of BoN-128).

4. Resultaten

De auteurs evalueren SEA op drie taken: Veiligheid (Safety), Waarachtigheid (Truthfulness) en Redeneren (Reasoning), met verschillende LLaMA-3 modellen.

• Veiligheid (AdvBench):
  • SEA overtreft de beste bestaande methoden (zoals BoN-64) aanzienlijk.
  • Op het LLaMA-3.2-1B-Base model verbeterde SEA de schadelijkheidsratio (Harmful Rate) met 91,54% ten opzichte van SFT, terwijl BoN-64 slechts een verbetering van ~33% boekte.
  • SEA is extreem robuust tegen Prefilling Attacks, waarbij de aanvalssuccesratio (ASR) van BoN stijgt met meer vooraf ingevoegde tokens, terwijl SEA stabiel blijft.
• Waarachtigheid (TruthfulQA):
  • SEA verbetert zowel de waarheidsgetrouwheid als de informativiteit en diversiteit, terwijl BoN-methoden vaak een afname in diversiteit vertonen bij grotere $N$.
• Redeneren (MATH & GSM8K):
  • Op de MATH-dataset boekte SEA een relatieve verbetering van 16,36% in nauwkeurigheid en een enorme stijging in de beloningswaarde (74,96%), terwijl zoekmethoden vaak faalden om de beloningsruimte effectief te verkennen.
• Efficiëntie:
  • SEA is sneller en memory-efficiënter dan token-level zoekmethoden (zoals ARGS) en presteert beter dan sentence-level BoN met grote $N$, ondanks het gebruik van iteratieve stappen.

5. Betekenis en Conclusie

Dit paper toont aan dat continue optimalisatie een onderbenutte maar krachtige richting is voor het uitlijnen van LLM's tijdens inferentie.

• Flexibiliteit: SEA werkt als een "plug-and-play" oplossing voor elk onuitgelijnd LLM zonder extra training.
• Robuustheid: Het lost fundamentele beperkingen op van zoekmethoden, vooral wanneer het basismodel zwak is of de zoekruimte groot is.
• Veiligheid: Het biedt een oplossing voor het probleem van "shallow alignment", wat essentieel is voor de veiligheid van AI-systemen in de praktijk.

De auteurs concluderen dat SEA een eenvoudige, maar effectieve manier biedt om de prestaties van LLM's te verbeteren door gebruik te maken van gradiëntinformatie in een continue ruimte, wat leidt tot superieure resultaten in veiligheid, waarheid en redenering vergeleken met de huidige state-of-the-art. De code is openbaar beschikbaar.