Wasserstein Gradient Flows for Batch Bayesian Optimal Experimental Design

Cet article propose une nouvelle approche pour la conception optimale d'expériences bayésienne par lots en relevant le problème d'optimisation vers l'espace des mesures de probabilité, où l'optimisation d'une régularisation entropique de l'utilité attendue mène à des algorithmes basés sur des particules dérivés des flots de gradient de Wasserstein.

L'essentiel

Imaginez que vous êtes un chef cuisinier très célèbre, mais que vous devez préparer un grand banquet pour des centaines de convives. Votre objectif est de choisir les meilleurs ingrédients et les meilleures techniques pour que le repas soit aussi informatif et délicieux que possible.

Dans le monde de la science, ce "repas", c'est une expérience. Le but est de concevoir l'expérience parfaite pour apprendre le maximum de choses sur un sujet mystérieux (comme la météo, une maladie ou un nouveau matériau). C'est ce qu'on appelle le Design Expérimental Bayésien Optimal (BOED).

Le problème ? Trouver la recette parfaite est un cauchemar.

1. Le paysage est accidenté : Imaginez une montagne avec des milliers de pics. Certains pics sont très hauts (très bons), d'autres sont des petits sommets trompeurs (moyens). Si vous cherchez la meilleure recette en grimpant simplement vers le haut, vous risquez de rester coincé sur un petit pic et de rater le sommet principal.
2. Le groupe est difficile : Souvent, on ne fait pas une seule expérience, mais un lot (un "batch") de plusieurs expériences en même temps. C'est comme essayer de coordonner 50 cuisiniers qui doivent tous cuisiner ensemble sans se marcher sur les pieds. La complexité explose.

Voici comment l'auteur de cet article, Louis Sharrock, propose de résoudre ce problème avec une idée brillante et simple.

1. Au lieu de chercher un point, cherchez une "Carte de Probabilité"

L'approche classique : On essaie de trouver un seul endroit précis sur la carte où l'expérience sera parfaite. C'est comme essayer de trouver la seule aiguille dans une botte de foin en regardant un point précis à la fois. Si vous commencez au mauvais endroit, vous ne trouverez rien.

L'approche de l'article : Au lieu de chercher un point, on cherche une distribution, c'est-à-dire une carte qui nous dit : "Il y a une forte probabilité que la meilleure expérience se trouve ici, une petite probabilité là-bas, et presque aucune chance ailleurs."

L'analogie du nuage :
Imaginez que vous cherchez le meilleur endroit pour planter un champ de fleurs.

• Méthode classique : Vous plantez une graine à un endroit précis. Si le sol est mauvais, vous avez perdu votre temps.
• Méthode de l'article : Vous lancez un nuage de graines (des particules) dans le ciel. Ce nuage va se déplacer, s'agglutiner autour des zones fertiles et s'éloigner des zones rocailleuses. À la fin, vous regardez où le nuage s'est concentré pour savoir où planter.

2. La "Température" et l'Exploration

Pour que ce nuage de graines ne reste pas coincé dans une petite vallée (un mauvais pic), l'auteur introduit un concept clé : la température (ou régularisation entropique).

• Température élevée (Froid) : Le nuage est très agité, il vole partout. Il explore tout le terrain, même les zones moins intéressantes. C'est utile au début pour ne rien rater.
• Température basse (Chaud) : Le nuage se calme et se concentre lentement sur les zones les plus fertiles.

C'est comme si vous aviez un groupe d'explorateurs. Au début, ils courent partout (température haute) pour repérer les ressources. Ensuite, ils se regroupent progressivement autour des meilleures trouvailles (température basse) pour en profiter.

3. La "Danse des Particules" (Flot de Gradient de Wasserstein)

Comment fait-on bouger ce nuage de manière intelligente ? L'article utilise une mathématique appelée Flot de Gradient de Wasserstein.

L'analogie de la rivière :
Imaginez que votre nuage de graines est une rivière. Le but est de faire couler cette rivière vers le point le plus bas (ou le plus haut, selon comment on regarde le problème), qui représente la meilleure expérience possible.

• La rivière suit le relief naturel (les mathématiques de l'information).
• Mais comme il y a des obstacles (les pics locaux), la rivière a besoin d'un peu de turbulence pour sauter par-dessus les petits rochers et continuer vers le grand océan.

L'auteur propose de simuler cette rivière avec des particules (des points individuels) qui interagissent entre elles.

• Si deux particules sont trop proches, elles se repoussent légèrement (pour éviter que tout le monde ne choisisse la même chose, ce qui serait redondant).
• Elles sont attirées par les zones où l'information est riche.

4. Pourquoi c'est génial pour les grands groupes (Batch) ?

Quand on doit faire 100 expériences en même temps, calculer la meilleure combinaison est impossible directement (trop de calculs).

L'auteur propose deux astuces de simplification :

1. Le modèle "Indépendant" (i.i.d.) : Au lieu de calculer comment les 100 expériences interagissent toutes ensemble (ce qui est un casse-tête), on apprend une seule règle de base qui s'applique à chaque expérience individuellement. C'est comme apprendre une seule recette de base que chaque cuisinier adapte légèrement.
2. L'ajout de "Repulsion" : Pour s'assurer que les 100 expériences ne sont pas toutes identiques (ce qui serait inutile), on ajoute une règle qui dit : "Éloignez-vous un peu les uns des autres !". Cela force le groupe à couvrir différentes zones du terrain, maximisant ainsi l'information totale.

En résumé

Cet article propose de passer d'une chasse à l'aiguille (trouver un seul point parfait) à une chasse au trésor collective (faire évoluer un nuage de possibilités).

• Le problème : Les expériences sont complexes, coûteuses et pleines de pièges (pics locaux).
• La solution : Utiliser un nuage de particules qui "flotte" sur une carte de probabilités, guidé par une température qui permet d'explorer puis de se concentrer.
• Le résultat : On obtient des groupes d'expériences (batches) beaucoup plus intelligents, qui évitent les pièges et couvrent mieux le terrain, même quand le nombre d'expériences est très grand.

C'est comme si, au lieu d'envoyer un seul détective chercher un criminel, on envoyait une équipe de détectives qui se parlent, se repoussent pour ne pas se marcher dessus, et finissent par encercler le coupable beaucoup plus efficacement.

Résumé technique

1. Problématique : Conception Expérimentale Bayésienne Optimale (BOED) par Lots

La Conception Expérimentale Bayésienne Optimale (BOED) vise à sélectionner des expériences afin de maximiser l'utilité attendue des données collectées, généralement mesurée par le Gain d'Information Attendu (EIG). L'EIG correspond à l'information mutuelle entre les paramètres inconnus $\theta$ et les observations futures $y$.

L'article se concentre sur le cas par lots (batch), où l'on doit concevoir simultanément $m$ expériences $\xi_{1:m} = (\xi_1, \dots, \xi_m)$. Ce problème présente plusieurs défis majeurs :

• Non-convexité et Multimodalité : Le paysage de l'utilité (EIG) est souvent fortement non-convexe avec de nombreux optima locaux, rendant les méthodes d'optimisation par gradient classiques (comme la descente de gradient stochastique) sensibles à l'initialisation et sujettes à la convergence vers des solutions sous-optimales.
• Complexité computationnelle : L'évaluation de l'EIG et de son gradient implique des espérances imbriquées (nested expectations) sur les paramètres et les données, nécessitant des approximations par Monte Carlo coûteuses et biaisées.
• Dimensionnalité : Dans le cadre par lots, la dimension de l'espace de conception passe de $d$ à $m \times d$, exacerbant la difficulté d'exploration.

2. Méthodologie Proposée

L'auteur propose une reformulation distributionnelle du problème d'optimisation, passant d'une optimisation sur un point de conception $\xi$ à une optimisation sur une mesure de conception $\nu$ (une loi de probabilité sur l'espace des designs).

A. Levée Probabiliste et Régularisation Entropique

Au lieu de chercher un vecteur $\xi^*$, on cherche une mesure $\nu_m \in \mathcal{P}(\Xi^m)$ minimisant une fonctionnelle d'énergie libre régularisée :
$$F_{\lambda, m}(\nu_m) = -\mathbb{E}_{\nu_m}[\text{EIG}_m(\xi_{1:m})] + \lambda_m \text{KL}(\nu_m \| \rho_m)$$
où $\rho_m$ est une mesure de référence et $\lambda_m > 0$ est un paramètre de régularisation (température).

• Avantage : Cette régularisation rend le problème strictement convexe sur l'espace des mesures, garantissant l'existence d'un minimiseur unique de forme Gibbs explicite :
  $$\frac{d\nu^*_{\lambda, m}}{d\rho_m} \propto \exp\left(\frac{\text{EIG}_m(\xi_{1:m})}{\lambda_m}\right)$$
• Interprétation : Pour $\lambda_m \to 0$, la mesure se concentre sur les designs optimaux. Pour $\lambda_m$ plus grand, la mesure explore l'espace, évitant les pièges locaux.

B. Approximations Évolutives pour la Scalabilité

Pour les grands lots ($m$ élevé), l'optimisation directe sur $\mathcal{P}(\Xi^m)$ est impossible. L'article introduit deux restrictions tractables :

1. Famille Mean-Field (Produit de mesures) : $\nu_m = \mu_1 \otimes \dots \otimes \mu_m$. Permet des marginales différentes pour chaque position du lot.
2. Famille i.i.d. (Produit identique) : $\nu_m = \mu^{\otimes m}$. On optimise une seule loi $\mu$ sur $\Xi$, et le lot est généré par tirages i.i.d.

C. Flots de Gradient de Wasserstein (WGF)

Pour optimiser ces lois $\mu$, l'auteur dérive les Flots de Gradient de Wasserstein associés à la fonctionnelle d'énergie libre.

• Ces flots sont décrits par des équations de Fokker-Planck non linéaires de type McKean-Vlasov.
• Ils sont approximés par des Systèmes de Particules Interagissantes (IPS).
• Algorithme Doublement Stochastique : Comme le gradient de l'EIG est intrinsèquement intractable (nécessite du Monte Carlo imbriqué), l'algorithme combine :
  1. Un échantillonnage de tuples de particules pour approximer les interactions (réduction de complexité $O(N^m)$ à $O(N)$).
  2. Des estimateurs Monte Carlo internes pour approximer le gradient de l'EIG.

D. Extraction de Lots Déterministes

La méthode produit une loi de conception $\mu$. Pour obtenir un lot déterministe $\hat{\xi}_{1:m}$, on utilise une stratégie "Best-of-n" (BoN) : on génère $n$ candidats selon $\mu^{\otimes m}$ et on sélectionne celui qui maximise l'EIG estimé.

3. Contributions Clés

1. Formulation Variationnelle : Reformulation du BOED par lots comme un problème d'optimisation variationnelle régularisé par l'entropie sur l'espace des mesures de probabilité, garantissant un minimiseur unique de forme Gibbs.
2. Développement Théorique : Dérivation des flots de gradient de Wasserstein pour les objectifs i.i.d. et mean-field, caractérisation de leur comportement à long terme (convergence vers l'équilibre) et analyse d'erreur finie (décomposition de l'erreur en effets de nombre fini de particules, discrétisation temporelle et approximation stochastique).
3. Algorithmes Scalables : Proposition d'algorithmes basés sur des particules, y compris des variantes doublement stochastiques capables de gérer les gradients de Monte Carlo imbriqué, rendant la méthode applicable à des problèmes réalistes.
4. Validation Empirique : Démonstration de l'efficacité sur des benchmarks synthétiques et réels, montrant une robustesse supérieure aux méthodes par points classiques dans des paysages non convexes.

4. Résultats Expérimentaux

Les expériences comparent les méthodes WGF (Joint, Mean-Field, i.i.d., i.i.d. avec répulsion) à des baselines (Gradient Ascent, SGA, SMC, méthodes d'échange de coordonnées).

• Optimisation 1D et 2D (Paysages multimodaux) :
  • Les méthodes WGF évitent la collapse de mode (convergence vers un optimum local) qui affecte gravement l'ascension de gradient (GA), même avec de multiples redémarrages.
  • Elles découvrent systématiquement les optima globaux et produisent des designs avec une entropie postérieure plus faible (plus d'information).
• Design de Capteurs et Placement sur le Tore :
  • Pour les grands lots, les approximations Mean-Field et i.i.d. surpassent souvent la méthode "Joint" (optimisation directe sur l'espace complet) en raison de la difficulté d'exploration en haute dimension.
  • L'ajout d'un terme de répulsion (pour encourager la diversité intra-lot) améliore les performances dans les scénarios où la redondance est pénalisante.
• Benchmarks Réels (Pharmacocinétique et FitzHugh-Nagumo) :
  • Sur des problèmes de conception de temps d'échantillonnage pour des modèles ODE non linéaires, les méthodes WGF (notamment les variantes Mean-Field et i.i.d.) atteignent des EIG supérieurs ou comparables aux meilleures méthodes existantes (comme l'échange de coordonnées ou SMC), tout en étant plus robustes aux initialisations.
  • Les designs obtenus capturent correctement les structures dynamiques critiques (phases de montée rapide et de queue d'élimination).

5. Signification et Impact

Cet article apporte une avancée significative dans le domaine de la BOED en :

• Changeant de paradigme : Passer d'une optimisation de point à une optimisation de loi de probabilité permet de traiter naturellement l'incertitude et la multimodalité.
• Résolvant le problème de la non-convexité : L'utilisation de flots de gradient stochastiques régularisés offre une voie théorique et pratique pour explorer des paysages d'utilité complexes sans être piégé par les optima locaux.
• Offrant une scalabilité : La combinaison de restrictions structurelles (i.i.d./Mean-Field) et d'estimateurs doublement stochastiques rend la conception de grands lots réalisable pour des modèles complexes et coûteux.
• Fournissant un cadre théorique rigoureux : Les garanties de convergence et les bornes d'erreur établies pour les algorithmes de particules dans ce contexte spécifique renforcent la crédibilité de la méthode pour des applications critiques (essais cliniques, ingénierie).

En résumé, cette approche transforme la BOED par lots en un problème d'inférence variationnelle dynamique, offrant une alternative robuste et efficace aux méthodes d'optimisation déterministes traditionnelles.