Thompson Sampling via Fine-Tuning of LLMs

Het paper introduceert ToSFiT, een schaalbaar algoritme dat Large Language Models fijnstemt om Thompson Sampling toe te passen voor Bayesiaanse optimalisatie in grote discrete ruimtes, waardoor de noodzaak voor het maximaliseren van acquisitiefuncties wordt geëlimineerd en zowel proef- als computerefficiëntie op diverse taken worden verbeterd.

De kern

Titel: Hoe een AI-leraar slimme keuzes leert maken zonder alles uit te proberen

Stel je voor dat je een enorme bibliotheek binnenloopt. Deze bibliotheek bevat niet alleen boeken, maar elke mogelijke combinatie van woorden die je ooit kunt bedenken. Je zoekt het perfecte antwoord op een vraag, of misschien de beste nieuwe medicijnformule, of het efficiëntste ontwerp voor een quantumcomputer. Het probleem? De bibliotheek is zo groot dat er meer boeken zijn dan er atomen in het heelal zijn. Als je elk boek één voor één zou lezen, zou je eeuwen nodig hebben.

Dit is het probleem waar wetenschappers vaak tegenaan lopen: Hoe vind je het beste antwoord in een onmetelijke, chaotische wereld, zonder tijd en geld te verspillen?

Deze paper introduceert een nieuwe methode genaamd TOSFIT. Laten we uitleggen hoe het werkt met een paar simpele metaforen.

1. Het oude probleem: De "Zoekmachine" die vastloopt

Traditionele methoden (zoals Baysean Optimization) werken als een slimme zoekmachine die probeert het beste antwoord te vinden. Ze doen dit door een "beloningssysteem" op te stellen. Ze zeggen: "Ik denk dat dit antwoord goed is, maar ik weet het niet zeker. Laten we een paar opties proberen en kijken wat er gebeurt."

Het probleem is dat deze zoekmachine vaak vastloopt in de enorme bibliotheek. Om het beste antwoord te vinden, moet de computer eerst alle mogelijke opties in gedachten doorrekenen om te zien welke het beste is. In een wereld met oneindige opties is dit onmogelijk. Het is alsof je probeert de snelste route naar huis te vinden door elke mogelijke straat in de stad één voor één te testen.

2. De oplossing: Een slimme gids met een "gevoel"

De auteurs van dit paper hebben een slimme truc bedacht. In plaats van de computer te laten rekenen en vergelijken, laten ze een Grote Taalmodel (LLM) – denk aan een super-intelligente AI die alles heeft gelezen – de keuzes maken.

Maar hoe leer je deze AI?
Stel je voor dat je een jonge student (de AI) hebt die al veel weet (door te trainen op internetdata). Je geeft hem een taak: "Schrijf het beste antwoord op deze vraag."

• Stap 1: De student schrijft een antwoord.
• Stap 2: Jij kijkt naar het antwoord en zegt: "Nee, dit is niet helemaal goed. Maar dit stukje hier is wel goed."
• Stap 3: De student past zijn kennis direct aan op basis van jouw feedback. Hij "fine-tunes" zichzelf.

Dit is wat TOSFIT doet. Het gebruikt de AI niet alleen als een generator van ideeën, maar als een leraar die zichzelf voortdurend verbetert.

3. De "Gokker" die niet blind gokt (Thompson Sampling)

In de wiskunde heet deze methode Thompson Sampling. Laten we het vergelijken met een gokker in een casino met duizenden gokkasten (machines).

• De meeste gokkers proberen gewoon willekeurig een kastje.
• Slimme gokkers proberen de kastjes die ze al hebben gezien als "goed" te herhalen, maar proberen ook af en toe een nieuw kastje om te kijken of er een betere is.

TOSFIT is als een gokker die gevoel heeft. Hij zegt niet: "Ik ga nu alle kastjes uitrekenen." Hij zegt: "Op basis van wat ik tot nu toe heb gezien, voel ik dat dit kastje de meeste kans heeft om de jackpot te geven." Hij "fine-tunes" zijn gevoel direct na elke gok.

4. Waarom is dit zo cool?

De paper toont aan dat deze methode drie grote voordelen heeft:

1. Het werkt in de chaos: Of je nu op zoek bent naar een nieuw medicijn (eiwitten), een betere tekst voor een klantenservice, of een quantumcomputer-code. Het maakt niet uit hoe complex of "ongestructureerd" de wereld is; de AI kan erin navigeren.
2. Het is snel en slim: Omdat de AI direct leert van zijn fouten, hoeft hij niet duizenden malen te proberen. Hij vindt het beste antwoord met veel minder pogingen dan andere methoden.
3. Het onthoudt zijn basis: Een groot risico bij het trainen van AI is dat ze vergeten wat ze al wisten (ze "vergeten" hun basis). TOSFIT is zo ontworpen dat de AI zijn oorspronkelijke kennis behoudt, maar die voorzichtig aanpast. Het is alsof je een expert bent die een nieuwe specialiteit leert, zonder zijn oude kennis te verliezen.

Samenvattend

Stel je voor dat je een schat zoekt op een eiland dat zo groot is dat je er nooit rond kunt lopen.

• De oude methode: Je probeert elke hoek van het eiland systematisch af te lopen. Je wordt moe en vindt de schat misschien nooit.
• De TOSFIT-methode: Je hebt een lokale gids (de AI) die het eiland kent. Je zegt: "Hier is een schat die we hebben gevonden." De gids kijkt, leert direct waar de schatten waarschijnlijk liggen, en past zijn kaart direct aan. Hij loopt niet blindelings, maar gebruikt zijn "gevoel" (de kansberekening) om de beste plek te kiezen.

Deze paper laat zien dat we door AI op deze manier te gebruiken, complexe problemen in wetenschap en technologie veel sneller en efficiënter kunnen oplossen. Het is een stap in de richting van machines die niet alleen rekenen, maar leren hoe ze slim moeten zoeken.

Technische samenvatting

Titel: Thompson Sampling via Fine-Tuning of LLMs (TOSFIT)

Auteurs: Nicolas Menet, Aleksandar Terzić, Michael Hersche, Andreas Krause, Abbas Rahimi (IBM Research Zürich & ETH Zürich).

---

1. Het Probleem

Bayesian Optimization (BO) is een krachtige methode voor het optimaliseren van dure, zwarte-doos functies, waarbij een posterior-verdeling over de beloning wordt bijgehouden om de zoektocht naar optimale configuraties te sturen. Traditionele BO-methoden gebruiken een acquisition function (zoals Upper Confidence Bound) die gemaximaliseerd moet worden om de volgende kandidaat te selecteren.

• De Uitdaging: In grote, ongestructureerde discrete ruimtes (zoals aminozuursequenties voor eiwitten, code voor quantum circuits, of tekstsequenties) is het maximaliseren van deze acquisition functie computationeel onhaalbaar. Er zijn geen gradiënten beschikbaar en de zoekruimte is te groot om exhaustief te doorzoeken (bijv. $20^{100}$ mogelijke eiwitsequenties).
• Huidige Beperkingen: Bestaande methoden zoals in-context Thompson sampling of evolutionaire zoekalgoritmen missen vaak de schaalbaarheid, de sample-efficiëntie, of de mogelijkheid om gebruik te maken van de rijke voorafgaande kennis (priors) van grote taalmodellen (LLMs).

2. Methodologie: TOSFIT

Het paper introduceert TOSFIT (Thompson Sampling via Fine-Tuning), een schaalbaar alternatief dat de noodzaak van acquisition function maximalisatie elimineert door de kans op maximaliteit (Probability of Maximality - PoM) direct te parametriseren met een generatief taalmodel.

Kernprincipes:

1. Directe Parametrisatie van PoM: In plaats van een functie te trekken uit de posterior en vervolgens het maximum te zoeken (wat in discrete ruimtes onmogelijk is), gebruikt TOSFIT een vooraf getraind LLM om direct te genereren volgens de verdeling van de PoM. Het model genereert kandidaten die per definitie "Thompson samples" zijn.
2. Variational Bayesian Optimistic Sampling (VBOS): De methode bouwt voort op VBOS, waarbij het beleid $\pi$ wordt geoptimaliseerd om een variational doel te maximaliseren:
   $$V(\pi) = \mathbb{E}_{x \sim \pi} [\mu_x + \sqrt{-2 \ln(\pi_x)} \cdot \sigma_x]$$
   Hierbij is $\mu_x$ de verwachte beloning en $\sigma_x$ de onzekerheid. De term $\sqrt{-2 \ln(\pi_x)}$ fungeert als een adaptieve exploratiebonus (vergelijkbaar met entropy regularisatie).
3. Fine-Tuning Strategie:
   • Initialisatie: Het beleid wordt geïnitieerd met een prompt-gestuurd, vooraf getraind LLM (gebruikmakend van sterke priors).
   • Adaptatie: Het model wordt voorzichtig fine-tuned via gradiëntopwaartse beweging op de VBOS-objective. Dit gebeurt in batches om de variance van de gradiëntenschatters te verlagen.
   • Stabilisatie: Om hoge variance in de gradiënten te voorkomen, wordt de Reinforce Leave-One-Out (RLOO) techniek gebruikt, gecombineerd met normalisatie van de voordeel-functie (vergelijkbaar met GRPO - Group Relative Policy Optimization).
4. Gaussische Processen (GP): Voor de reward-modellering wordt een lineaire Kernel GP gebruikt met diepe features (embeddings) van het LLM. Dit maakt inferentie en marginal likelihood maximalisatie mogelijk in gesloten vorm met constante complexiteit ten opzichte van het aantal observaties.

3. Belangrijkste Bijdragen

1. Theoretische Regret Bound:

   • De auteurs leiden een nieuwe cumulatieve regret-bound af voor een variational formulering van Thompson Sampling.
   • Voor exacte VBOS verbeteren ze de bound van $\tilde{O}(\sqrt{T|X|})$ naar $\tilde{O}(\sqrt{T\gamma_T})$, waarbij $\gamma_T$ de maximale informatiewinst is. Dit maakt de methode schaalbaar voor combinatorisch grote ruimtes waar $|X|$ exponentieel groot is.
   • Ze generaliseren de bound voor benaderde (gradiënt-gebaseerde) VBOS, waarbij de fout afhankelijk is van de Bregman-divergentie tussen het gesampelde beleid en de exacte optimizer.

2. Het TOSFIT Algorithm:

   • Een praktisch algoritme dat LLMs gebruikt om Thompson sampling te implementeren zonder acquisition function maximalisatie.
   • Het benadrukt het belang van voorzichtige adaptatie: het beleid moet dicht bij de initiele prior blijven (door pre-training) om te voorkomen dat het model "vergeet" hoe het geldige oplossingen genereert, terwijl het toch leert om naar de posterior toe te bewegen.

3. Empirische Validatie:

   • Toepassing op drie diverse, complexe taken:
     • FAQ Response Refinement: Het optimaliseren van tekstantwoorden op basis van semantische alignement.
     • Thermally Stable Protein Search: Het vinden van aminozuursequenties met hoge thermische stabiliteit (een ruimte groter dan het aantal atomen in het waarneembare universum).
     • Quantum Circuit Design: Het ontwerpen van Qiskit circuits die lage-energie toestanden voorbereiden.

4. Resultaten

• Sample Efficiency: TOSFIT presteert state-of-the-art op alle drie de taken. Het overtreft methoden zoals Unguided Generation, Post-Generation Thompson Sampling (die beperkt is tot een vaste pool), Actor-Critic methoden, en evolutionaire zoekalgoritmen.
• Computational Efficiency: Ondanks de extra kosten van fine-tuning per ronde, is TOSFIT computatie-efficiënter dan concurrenten omdat het veel minder evaluaties van de dure reward-functie nodig heeft om een goede oplossing te vinden.
• Batched Optimization: De methode werkt uitstekend in batched settings (parallelle evaluatie), wat de wandtijd (wall-clock time) verder verlaagt.
• Ablatie Studies:
  • Priors: Het gebruik van sterke priors (via prompts) is cruciaal; zwakkere context leidt tot slechtere prestaties.
  • Learning Rate: Een te hoge learning rate tijdens fine-tuning leidt tot het vergeten van de prior en stagnatie, wat de noodzaak van "careful adaptation" bevestigt.

5. Significantie en Conclusie

Dit paper biedt een fundamentele doorbraak in het toepassen van Bayesian Optimization op grote, ongestructureerde discrete ruimtes, een domein waar traditionele methoden vaak vastlopen.

• Paradigmaverschuiving: Het verschuift de focus van het maximaliseren van een acquisition functie naar het genereren van samples die direct de posterior van maximaliteit volgen, gebruikmakend van de generatieve capaciteit van LLMs.
• Synergie: Het combineert de principes van Bayesian Optimization (onzekerheidsgebaseerde exploratie) met de kracht van Foundation Models (voorafgaande kennis en generatie).
• Toepassingsbereik: De methode is breed toepasbaar in wetenschappelijke ontdekkingen (drug discovery, material science) en engineering (quantum computing), waar de zoekruimtes te groot zijn voor exhaustieve zoektochten en te complex voor traditionele gradiënt-methoden.

Kortom, TOSFIT demonstreert dat het fine-tunen van LLMs als een Thompson Sampling-strategie een krachtig, schaalbaar en theoretisch onderbouwd raamwerk biedt voor complexe discrete optimalisatieproblemen.