Wasserstein Gradient Flows for Batch Bayesian Optimal Experimental Design

Deze paper introduceert een nieuwe aanpak voor batch Bayesiaans optimaal experimenteel ontwerp door het optimalisatieprobleem te verheffen naar de ruimte van kansmaten en te benaderen via entropisch geregulariseerde objectives en Wasserstein-gradiëntstromen, wat leidt tot schaalbare deeltjesalgoritmen die effectief multimodale landschappen kunnen verkennen.

De kern

Stel je voor dat je een groot, donker raadsel moet oplossen. Je hebt een set van vragen nodig om het antwoord te vinden, maar je mag maar een beperkt aantal vragen stellen. De kunst is om die vragen zo slim te kiezen dat je met zo min mogelijk moeite het meeste leert. Dit noemen wetenschappers Bayes Optimal Experimental Design (BOED).

Het probleem is echter: het vinden van die perfecte vragen is als het zoeken naar de top van een berglandschap in een dikke mist. Soms zit je vast in een klein heuveltje (een lokaal maximum) en denk je dat je de top hebt gevonden, terwijl er ergens anders een veel hogere bergpiek is. Vooral als je meerdere vragen tegelijk moet kiezen (een "batch"), wordt het landschap zo complex dat traditionele methoden erin vastlopen.

In dit paper introduceert Louis Sharrock een nieuwe manier om dit probleem aan te pakken, gebaseerd op Wasserstein Gradient Flows. Laten we dit uitleggen met een paar creatieve metaforen.

1. Van één pijl naar een hele regen van pijlen

De oude manier (Punt-optimisatie):
Stel je voor dat je een darter bent. Je probeert één pijl te gooien die precies in het midden van het doelwit landt. Je kijkt naar je vorige worp, past je houding een beetje aan en gooit weer. Als je echter in een hoek van het doelwit begint (een slechte startpositie), loop je het risico dat je voor altijd in die hoek blijft hangen, zelfs als het echte doel in het midden ligt. Je zoekt naar één perfecte oplossing.

De nieuwe manier (Distributie-optimisatie):
In plaats van één darter, stuur je nu een heel leger darters de zaal in. Ze gooien allemaal tegelijk, verspreid over de hele zaal. In plaats van te proberen één pijl perfect te gooien, proberen we de verdeling van al die pijlen te optimaliseren. We willen dat de meeste pijlen uiteindelijk in de beste gebieden van het doelwit belanden.

Dit is wat de auteur doet: hij stopt niet met het zoeken naar één perfecte vraag, maar zoekt naar de beste verdeling van vragen.

2. De "Temperatuur" van de exploratie

Om te voorkomen dat al je darters in één klein hoekje blijven hangen, voegt de auteur een beetje "chaos" of "warmte" toe.

• Koude temperatuur: Alle pijlen willen direct naar het punt met de hoogste score. Dit is risicovol; als je startpunt slecht is, mis je de top.
• Warme temperatuur: De pijlen bewegen wat wilder. Ze verkennen het hele doelwit. Ze kunnen over heuvels springen die ze anders niet zouden overwinnen.

De auteur gebruikt een wiskundige techniek (entropische regularisatie) om deze "warmte" te regelen. Het is alsof je eerst een warme, vloeibare substantie hebt die over het hele landschap stroomt en alle dalen en toppen verkent. Naarmate je de temperatuur verlaagt, stolt deze vloeistof langzaam en vormt hij een perfecte vorm die precies de beste plekken op het landschap bedekt.

3. De stroom van de rivier (Wasserstein Gradient Flow)

Hoe laat je die vloeistof nu stromen naar de beste plekken?
Stel je voor dat het landschap van je vragen een rivierbedding is. De "stroom" (de Gradient Flow) zorgt ervoor dat de vloeistof altijd stroomt naar beneden, richting de diepste valleien (of in dit geval, de hoogste pieken, omdat we het probleem omdraaien).

De Wasserstein Gradient Flow is een heel slimme manier om die stroom te berekenen. Het kijkt niet alleen naar de helling onder je voeten, maar naar hoe de hele rivier zich moet verplaatsen om het meest efficiënt naar de top te komen. Het is alsof je een stroompje water hebt dat zichzelf herschikt om de kortste weg naar de top te vinden, terwijl het tegelijkertijd alle obstakels omzeilt.

4. Het probleem van de schaal (De "Zwerm")

Als je 100 vragen tegelijk moet kiezen, wordt het landschap gigantisch groot. Het is onmogelijk om alles in één keer te berekenen.
De auteur lost dit op met twee slimme trucs:

• De Zwerm (Mean-Field): In plaats van 100 individuele darters die allemaal met elkaar praten (wat te veel rekenkracht kost), laten we ze allemaal praten met een "gemiddeld" beeld van de groep. Het is alsof elke darter kijkt naar waar de rest van de zwerm zit, en daarop reageert. Dit maakt het berekenen veel sneller.
• De Identieke Zwerm (i.i.d.): Nog simpeler: we nemen aan dat alle darters exact hetzelfde gedrag vertonen. Ze zijn allemaal kopieën van elkaar. Dit is de snelste manier, maar vereist dat we later een slimme "selectie" doen om te zorgen dat ze niet allemaal op exact dezelfde plek landen (want dat zou saai zijn).

5. Het eindresultaat: De "Best-of-N" strategie

Uiteindelijk krijg je geen één antwoord, maar een ontwerpwet: een kansverdeling die je vertelt waar je waarschijnlijk de beste vragen kunt vinden.
In de praktijk doe je dit zo:

1. Je gebruikt de "stroom" om een verdeling te leren die de beste gebieden verkent.
2. Je trekt duizenden mogelijke sets van vragen uit deze verdeling.
3. Je kiest de allerbeste set uit die duizenden (de "Best-of-N").

Waarom is dit belangrijk?

In de echte wereld (zoals bij het testen van nieuwe medicijnen of het plaatsen van sensoren) zijn experimenten duur en tijdrovend.

• Oude methoden lopen vaak vast in slechte oplossingen omdat ze te snel beslissen.
• Deze nieuwe methode is als een verkenner die eerst het hele terrein afloopt, de beste routes in kaart brengt, en pas daarna de perfecte route kiest. Het is robuuster, vindt betere oplossingen in complexe situaties, en kan zelfs grotere problemen aan.

Kort samengevat:
De auteur heeft een manier bedacht om niet naar één perfecte oplossing te jagen, maar om een slimme verdeling van mogelijke oplossingen te laten "stromen" naar de beste plek. Door een beetje chaos (warmte) toe te voegen en slimme wiskundige stromen te gebruiken, kunnen we complexe raadsels oplossen die voor oude methoden te moeilijk waren. Het is het verschil tussen een darter die blindelings schiet, en een heel leger dat samenwerkt om het doelwit te veroveren.

Technische samenvatting

Hier is een gedetailleerde technische samenvatting van het artikel "Wasserstein Gradient Flows for Batch Bayesian Optimal Experimental Design" van Louis Sharrock, geschreven in het Nederlands.

Probleemstelling

Bayes Optimaal Experimenteel Ontwerp (BOED) biedt een raamwerk om experimenten te selecteren die de verwachte nuttigheid van de verzamelde data maximaliseren, vaak gemeten door de Verwachte Informatiewinst (EIG - Expected Information Gain). In de praktijk is het optimaliseren van de EIG echter computationeel zeer uitdagend vanwege:

1. Neststructuur: De EIG vereist geneste verwachtingen (over parameters $\theta$ en observaties $y$), wat leidt tot complexe schattingsproblemen met bias-variatie trade-offs.
2. Niet-convexiteit: Het landschap van de verwachte nuttigheid is vaak sterk niet-convex en multimodaal, wat lokale optima creëert waar traditionele methoden in vastlopen.
3. Batch-uitdagingen: In de batch-setting (waar $m$ experimenten simultaan worden ontworpen) groeit de dimensie van de zoekruimte naar $m \times d$. De interacties tussen ontwerp-punten maken het landschap nog complexer, en het vinden van een optimale batch is nog moeilijker dan het vinden van een enkel punt.

Bestaande methoden, zoals Monte Carlo-gebaseerde benaderingen of variatiele inferentie, worstelen vaak met deze schaalbaarheid en de neiging om in lokale optima vast te komen, vooral bij grote batchgroottes.

Methodologie

De auteur introduceert een fundamenteel nieuwe aanpak: het vervangen van het optimaliseren van een enkel ontwerppunt $\xi$ door het optimaliseren over de ruimte van kansmaat (probability measures) $\nu$ op het ontwerp-domein.

1. Probabilistische Lift en Entropische Regularisatie

In plaats van te zoeken naar een deterministisch $\xi^*$, wordt het probleem gelift naar het vinden van een optimale verdeling $\nu_m$ over de batchruimte $\Xi^m$. Om het probleem goed gesteld te maken en uniekheid te garanderen, wordt een entropische regularisatie toegevoegd. Het doel is het minimaliseren van de vrije-energie functional:
$$F_{\lambda, m}(\nu_m) = -\mathbb{E}_{\nu_m}[EIG_m(\xi_{1:m})] + \lambda_m KL(\nu_m \| \rho_m)$$
Hierbij is $\lambda_m$ een temperatuurparameter en $\rho_m$ een referentiemaat.

• Resultaat: Onder milde voorwaarden heeft dit probleem een unieke minimizer met een expliciete Gibbs-vorm:
  $$\frac{d\nu^*_{\lambda, m}}{d\rho_m} \propto \exp\left(\frac{EIG_m(\xi_{1:m})}{\lambda_m}\right)$$
  Dit zorgt voor een strikt convex optimalisatieprobleem dat de kans op het vastlopen in lokale optima verkleint door exploratie via de entropie-term.

2. Schaalbare Benaderingen (Mean-Field en i.i.d.)

Het direct simuleren van de gezamenlijke Gibbs-verdeling voor grote $m$ is onhaalbaar. Daarom worden twee tractabele restricties voorgesteld:

• Mean-Field Familie: $\nu_m = \mu_1 \otimes \dots \otimes \mu_m$. Hierbij zijn de componenten onafhankelijk maar niet noodzakelijk identiek (verschillende randverdelingen).
• i.i.d. Familie: $\nu_m = \mu^{\otimes m}$. Hierbij wordt gezocht naar één enkele ontwerp-wet $\mu$ die voor alle $m$ componenten geldt. Dit reduceert de complexiteit aanzienlijk.

3. Wasserstein Gradient Flows (WGF)

Voor de i.i.d. (en formeel voor de mean-field) doelstelling wordt een Wasserstein Gradient Flow afgeleid. Dit is een dynamisch systeem dat de kansverdeling $\mu_t$ in de tijd evolueert in de richting van de steilste daling van de vrije energie in de ruimte van kansmaten (uitgerust met de $W_2$-metriek).

• De flow wordt beschreven door een McKean-Vlasov Fokker-Planck PDE.
• De stationaire oplossing van deze flow komt overeen met de Gibbs-minimizer.

4. Interacting Particle Systems (IPS) en Doubly Stochastic Algoritmen

Om de WGF numeriek op te lossen, wordt deze benaderd door een Interacting Particle System (IPS).

• Doubly Stochastic: Omdat de gradiënt van de EIG ($\nabla EIG$) vaak niet analytisch beschikbaar is en via geneste Monte Carlo moet worden geschat, introduceert de auteur een doubly stochastic algoritme.
  • Outer stochasticity: Subsampling van tuples van deeltjes om de interactie-term (gemiddelde over de batch) te schatten.
  • Inner stochasticity: Gebruik van een stochastische gradiëntschatter (bijv. geneste Monte Carlo) voor de EIG-gradiënt zelf.
• Dit resulteert in een schaalbaar algoritme dat geschikt is voor grote batchgroottes en complexe modellen.

Belangrijkste Bijdragen

1. Formulering: BOED wordt geformuleerd als een entropisch geregulariseerd variatieel optimalisatieprobleem over de ruimte van ontwerp-maten, wat leidt tot een unieke Gibbs-oplossing.
2. Schaalbaarheid: Introductie van mean-field en i.i.d. restricties die de dimensie van het probleem reduceren, vergezeld van afgeleide vaste-puntvergelijkingen.
3. WGF Afleiding: Afleiding van de bijbehorende Wasserstein Gradient Flows voor deze doelstellingen, gekarakteriseerd als niet-lineaire McKean-Vlasov SDE's.
4. Algoritmen: Ontwikkeling van schaalbare, deeltjesgebaseerde algoritmen (Euler-Maruyama discretisatie) die "doubly stochastic" zijn, waardoor ze om kunnen gaan met de inherente onnauwkeurigheid van EIG-schattingen.
5. Theoretische Garanties: Bewijzen van bestaan, uniciteit, en convergentie (exponentiële contractiviteit onder sterke convexiteitsvoorwaarden) van de flows, evenals foutdecomposities voor de eind-tot-eind benadering (partikel-aantal, tijdsdiscretisatie, en stochastische gradiëntfout).

Resultaten

De methode is getest op diverse numerieke experimenten, waaronder synthetische 1D/2D problemen en gevestigde benchmarks uit de literatuur (Farmacokinetiek en FitzHugh-Nagumo modellen).

• Multimodale Exploratie: In tegenstelling tot traditionele methoden zoals Stochastic Gradient Ascent (SGA) of Gradient Ascent (GA), die vaak in lokale optima vastlopen afhankelijk van de initialisatie, slaagt de WGF erin om het globale maximum te vinden door de toegevoegde ruis (temperatuur) en de verdelingsbenadering.
• Batch-Optimalisatie: Voor grote batchgroottes ($m > 1$) presteren de i.i.d. en mean-field WGF-methoden consistent beter dan puntsgewijze optimalisatie en herhaalde single-design benaderingen. Ze vermijden redundantie binnen de batch en vinden diverse, informatieve ontwerpen.
• Vergelijking met State-of-the-Art: In de farmacokinetische en FHN benchmarks presteerden de WGF-methoden (vooral de mean-field varianten) concurrerend of superieur ten opzichte van geavanceerde baselines zoals Annealed SMC, Coordinate Exchange (CE) algoritmen en variatiele benaderingen, vaak met een hogere EIG en minder variabiliteit tussen verschillende seeds.
• Robuustheid: De methoden tonen minder gevoeligheid voor slechte initialisatie en zijn robuuster in niet-convexe landschappen.

Significantie

Dit artikel biedt een fundamentele verschuiving in hoe Batch BOED wordt benaderd: van het optimaliseren van een vector naar het optimaliseren van een verdeling.

• Theoretische Diepgang: Het koppelt BOED aan de theorie van optimal transport en gradient flows in de ruimte van kansmaten, wat nieuwe inzichten biedt in convergentie en stabiliteit.
• Praktische Toepasbaarheid: De voorgestelde "doubly stochastic" deeltjesalgoritmen bieden een praktische, schaalbare oplossing voor een van de meest computationeel intensieve problemen in het veld: het ontwerpen van grote batches experimenten in niet-lineaire, niet-convexe settings.
• Flexibiliteit: Het raamwerk is modulair en kan worden gecombineerd met verschillende methoden voor het schatten van de EIG-gradiënt (variational, likelihood-free, etc.), waardoor het breed toepasbaar is in moderne, complexe wetenschappelijke toepassingen.

Kortom, de paper levert een krachtig, wiskundig onderbouwd en computationeel efficiënt raamwerk dat de barrières voor schaalbaar en robuust Batch BOED verlaagt.