In-Context Learning for Latent Space Bayesian Optimization

Résumé Technique : Apprentissage en Contexte pour l'Optimisation Bayésienne dans l'Espace Latent

Énoncé du Problème
L'optimisation bayésienne (BO) est une approche standard pour optimiser des objectifs boîte noire coûteux sur des espaces structurés, tels que la conception moléculaire et l'ingénierie des protéines. L'optimisation bayésienne dans l'espace latent (LSBO) traite la nature combinatoire de ces espaces en apprenant une représentation latente continue (via un encodeur-décodeur, typiquement un VAE) et en effectuant la BO au sein de cet espace latent.

Bien que les modèles de fondation tabulaires (TFM) et les réseaux ajustés sur des données préal कितनी (PFN) aient atteint des performances de pointe en régression grâce à l'apprentissage en contexte, leur application directe à la LSBO se heurte à un décalage critique. Les TFM actuels sont préentraînés sur de vastes collections de tâches de régression synthétiques qui ne reflètent pas la géométrie spécifique des espaces latents ni la nature des objectifs moléculaires (ex: logP, QED). Par conséquent, le « prior implicite » appris par ces modèles lors du préentraînement ne s'aligne pas sur l'environnement de déploiement de la LSBO, où la carte de l'espace latent vers la valeur de l'objectif est façonnée par le décodeur et des propriétés spécifiques au domaine. Traiter les TFM comme de simples régresseurs génériques ne permet pas de tirer pleinement parti de leur potentiel dans ce cadre spécialisé.

Méthodologie
Les auteurs proposent LILBO (Latent In-context Learning for Bayesian Optimization), un cadre qui adapte un modèle de fondation tabulaire préentraîné (spécifiquement TabPFN-3) au domaine de la LSBO par un préentraînement continu. La méthodologie se compose de deux composantes principales :

1. Génération de Tâches Synthétiques Sensibles à la BO :
   Au lieu d'utiliser des priors de régression génériques, les auteurs construisent une distribution de préentraînement synthétique ancrée dans les objectifs moléculaires.

• Tâches de Base : Les objectifs synthétiques sont dérivés d'une bibliothèque de tâches moléculaires de base (ex: logP, QED, similitude, redécouverte) issues de la base de données GuacaMol.
• Construction de l'Objectif : Pour chaque épisode d'entraînement, un sous-ensemble de tâches de base est échantillonné, et un objectif scalaire synthétique est créé via des combinaisons linéaires (pondérées par des distributions de Dirichlet) ou des combinaisons non linéaires (utilisant des MLP aléatoires).
• Échantillonnage du Contexte : Pour imiter le comportement de la BO, qui se concentre sur les régions de haute valeur, les données de contexte ne sont pas échantillonnées de manière uniforme. Au lieu de cela, une distribution biaisée par la valeur (distribution de Boltzmann) est induite sur le pool moléculaire, mettant l'accent sur les régions prometteuses de l'espace latent du VAE tout en conservant la diversité.

2. Préentraînement Continu avec Régularisation :
   Les auteurs initialisent le modèle de substitution à partir d'un point de contrôle (checkpoint) de TabPFN-3 préentraîné plutôt que de l'entraîner à partir de zéro. Cela préserve le large prior de régression appris lors du préentraînement à grande échelle.

• Objectif : Le modèle subit un préentraînement continu sur les épisodes de LSBO moléculaire synthétique.
• Régularisation : Pour éviter l'oubli catastrophique et la spécialisation excessive, un régularisateur L2-SP (L2-Soft Parameter) est appliqué. Cela ancre les poids du modèle adapté ($\theta$) aux poids du point de contrôle original ($\theta_0$), garantissant que le modèle conserve ses capacités de régression polyvalentes tout en s'adaptant à la géométrie spécifique des espaces latents moléculaires.
• Boucle d'Entraînement : Le modèle est entraîné pendant 8 000 étapes sur environ 500 000 épisodes de LSBO moléculaire.

Contributions Clés

• Adaptation Spécifique à la LSBO : Le papier introduit une méthode pour adapter les substituts d'apprentissage en contexte spécifiquement pour l'optimisation dans l'espace latent, traitant le décalage de distribution entre le préentraînement des TFM génériques et le déploiement de la LSBO.
• Conception de la Distribution Synthétique : Les auteurs conçoivent un pipeline de génération de tâches synthétiques qui incorpore des tâches de base moléculaires et un échantillonnage de contexte biaisé par la valeur pour mieux s'aligner sur les exigences de la BO (se concentrer sur les régions de haute valeur).
• Apprentissage Continu Régularisé : L'utilisation d'un objectif de préentraînement continu régularisé permet au modèle de se spécialiser dans la LSBO moléculaire sans abandonner les larges priors du modèle de fondation.
• Diagnostics : Le papier introduit des diagnostics dans l'espace latent et l'espace des objectifs (via UMAP) pour vérifier que la distribution de préentraînement synthétique couvre les régions pertinentes du manifold latent et de l'espace des objectifs.

Résultats
Le modèle proposé, LILBO, a été évalué sur des benchmarks d'optimisation moléculaire mis de côté de la suite GuacaMol (ex: Osimertinib, Fexofenadine, Ranolazine), qui n'avaient pas été vus lors du préentraînement.

• Performance : LILBO a obtenu des performances solides, surpassant la ligne de base de remplacement TabPFN-3 et plusieurs lignes de base spécialisées de LSBO (incluant GP-LSBO, LOLBO, CoBO et InvBO) en termes de classement moyen à travers les tâches.
• Efficacité : L'amélioration par rapport à la ligne de base standard TabPFN-3 est modeste mais notable, obtenue avec une adaptation relativement légère (500k épisodes) comparée à l'échelle massive du préentraînement du modèle de fondation original (billions de tokens).
• Viabilité : Les résultats démontrent que les substituts probabilistes en contexte sont une alternative viable aux processus gaussiens dans les pipelines de LSBO modernes, à condition que la distribution de préentraînement soit appropriée au domaine.

Signification et Revendications
Le papier affirme que pour que les modèles de fondation tabulaires soient efficaces dans la LSBO, leur distribution d'adaptation doit réduire le décalage entre les tâches de préentraînement génériques et la structure spécifique des objectifs de l'espace latent et des historiques de BO rencontrés lors du déploiement. Les auteurs soutiennent que l'adaptation spécifique à la LSBO est cruciale pour libérer tout le potentiel des substituts en contexte dans la conception scientifique.

Le travail reconnaît des limites, notant que l'approche actuelle suppose un VAE figé pendant le préentraînement, alors que le déploiement implique souvent un réentraînement périodique du VAE (dérive latente). Les directions futures suggérées incluent la simulation de trajectoires d'optimisation courtes pour les contextes d'entraînement, la prise en compte de la dérive latente dans le prior, l'extension à l'optimisation multi-objectif, et l'application de ce cadre à d'autres domaines structurés comme la conception de protéines.