In-Context Learning for Latent Space Bayesian Optimization

Sintesi Tecnica: In-Context Learning per la Ottimizzazione Bayesiana nello Spazio Latente

Definizione del Problema
L'Ottimizzazione Bayesiana (BO) è un approccio standard per ottimizzare obiettivi black-box costosi su spazi strutturati, come il design molecolare e l'ingegneria delle proteine. L'Ottimizzazione Bayesiana nello Spazio Latente (LSBO) affronta la natura combinatoria di questi spazi apprendendo una rappresentazione latente continua (tramite un encoder-decoder, tipicamente un VAE) ed eseguendo la BO all'interno di tale spazio latente.

Sebbene i modelli fondativi tabulari (TFM) e le reti basate su dati pregressi (PFN) abbiano raggiunto prestazioni allo stato dell'arte nei compiti di regressione attraverso l'in-context learning, la loro applicazione diretta alla LSBO affronta un disallineamento critico. Gli attuali TFM sono pre-addestrati su grandi collezioni di compiti di regressione sintetici che non riflettono la specifica geometria degli spazi latenti o la natura degli obiettivi molecolari (ad esempio, logP, QED). Di conseguenza, il "prior implicito" appreso da questi modelli durante il pre-addestramento non si allinea con l'ambiente di distribuzione della LSBO, dove la mappa dal codice latente al valore dell'obiettivo è modellata dal decoder e dalle proprietà specifiche del dominio. Trattare i TFM come generici regressori plug-in impedisce di sfruttarne appieno il potenziale in questo contesto specializzato.

Metodologia
Gli autori propongono LILBO (Latent In-context Learning for Bayesian Optimization), un framework che adatta un modello fondativo tabulare pre-addestrato (specificamente TabPFN-3) al dominio LSBO attraverso il pre-addestramento continuo. La metodologia consiste in due componenti primarie:

1. Generazione di Compiti Sintetici Consapevoli della BO:
   Invece di utilizzare prior di regressione generici, gli autori costruiscono una distribuzione di pre-addestramento sintetico radicata negli obiettivi molecolari.

• Compiti Base: Gli obiettivi sintetici sono derivati da una libreria di compiti molecolari base (ad esempio, logP, QED, similarità, rediscovery) estratti dal database GuacaMol.
• Costruzione dell'Obiettivo: Per ogni episodio di addestramento, viene selezionato un sottoinsieme di compiti base e viene creato un obiettivo scalare sintetico tramite combinazioni lineari (pesate da distribuzioni di Dirichlet) o combinazioni non lineari (utilizzando MLP casuali).
• Campionamento del Contesto: Per simulare il comportamento della BO, che si concentra sulle regioni ad alto valore, il contesto non è campionato uniformemente. Viene invece indotta una distribuzione con bias di valore (distribuzione di Boltzmann) sul pool molecolare, enfatizzando le regioni promettenti dello spazio latente del VAE pur mantenendo la diversità.

2. Pre-addestramento Continuo con Regolarizzazione:
   Gli autori inizializzano il modello di sostituzione (surrogate) da un checkpoint pre-addestrato di TabPFN-3 anziché addestrarlo da zero. Ciò preserva il ampio prior di regressione appreso durante il pre-addestramento su larga scala.

• Obiettivo: Il modello subisce un pre-addestramento continuo sugli episodi di LSBO molecolare sintetica.
• Regolarizzazione: Per prevenire l'oblio catastrofico e la sovra-specializzazione, viene applicato un regolarizzatore L2-SP (L2-Soft Parameter). Questo ancora i pesi del modello adattato ($\theta$) ai pesi del checkpoint originale ($\theta_0$), garantendo che il modello mantenga le sue capacità di regressione general-purpose pur adattandosi alla specifica geometria degli spazi latenti molecolari.
• Ciclo di Addestramento: Il modello viene addestrato per 8.000 step su circa 500.000 episodi di LSBO molecolare.

Contributi Chiave

• Adattamento Specifico per LSBO: Il documento introduce un metodo per adattare i sostituti in-context learning specificamente per l'ottimizzazione dello spazio latente, affrontando il disallineamento tra il pre-addestramento dei TFM generici e la distribuzione della LSBO in fase di deployment.
• Design della Distribuzione Sintetica: Gli autori progettano una pipeline di generazione di compiti sintetici che incorpora compiti molecolari base e un campionamento del contesto con bias di valore per allinearsi meglio ai requisiti della BO (focalizzazione sulle regioni ad alto valore).
• Apprendimento Continuo Regolarizzato: L'uso di un obiettivo di pre-addestramento continuo regolarizzato consente al modello di specializzarsi nella LSBO molecolare senza scartare i ampi prior del modello fondativo.
• Diagnostica: Il documento introduce diagnostiche dello spazio latente e dello spazio degli obiettivi (tramite UMAP) per verificare che la distribuzione di pre-addestramento sintetico copra le regioni rilevanti del manifold latente e dello spazio degli obiettivi.

Risultati
Il modello proposto, LILBO, è stato valutato su benchmark di ottimizzazione molecolare tenuti fuori dal set di addestramento (held-out) dalla suite GuacaMol (ad esempio, Osimertinib, Fexofenadine, Ranolazine), che non erano stati visti durante il pre-addestramento.

• Prestazioni: LILBO ha ottenuto prestazioni elevate, superando il baseline plug-in off-the-shelf di TabPFN-3 e diversi baseline specializzati di LSBO (inclusi GP-LSBO, LOLBO, CoBO e InvBO) in termini di ranking medio tra i compiti.
• Efficienza: Il miglioramento rispetto al baseline standard di TabPFN-3 è stato modesto ma degno di nota, ottenuto con un adattamento relativamente leggero (500k episodi) rispetto alla scala massiccia del pre-addestramento del modello fondativo originale (trilioni di token).
• Viabilità: I risultati dimostrano che i sostituti probabilistici in-context sono un'alternativa valida ai Processi Gaussiani nelle moderne pipeline di LSBO, a condizione che la distribuzione di pre-addestramento sia appropriatamente corrispondente al dominio.

Significato e Rivendicazioni
Il documento sostiene che, affinché i modelli fondativi tabulari siano efficaci nella LSBO, la loro distribuzione di adattamento debba ridurre il disallineamento tra i compiti di pre-addestramento generici e la specifica struttura degli obiettivi dello spazio latente e delle storie di BO incontrate durante il deployment. Gli autori sostengono che l'adattamento specifico per la LSBO sia cruciale per sbloccare il pieno potenziale dei sostituti in-context nel design scientifico.

Il lavoro riconosce i limiti, notando che l'approccio attuale assume un VAE congelato durante il pre-addestramento, mentre il deployment comporta spesso un ri-addestramento periodico del VAE (drift latente). Le direzioni future suggerite includono la simulazione di brevi traiettorie di ottimizzazione per i contesti di addestramento, la gestione del drift latente nel prior, l'estensione all'ottimizzazione multi-obiettivo e l'applicazione del framework ad altri domini strutturati come il design delle proteine.