In-Context Learning for Latent Space Bayesian Optimization

Resumen Técnico: Aprendizaje en Contexto para la Optimización Bayesiana en el Espacio Latente

Planteamiento del Problema
La Optimización Bayesiana (BO) es un enfoque estándar para optimizar objetivos de caja negra costosos sobre espacios estructurados, como el diseño molecular y la ingeniería de proteínas. La Optimización Bayesiana en el Espacio Latente (LSBO) aborda la naturaleza combinatoria de estos espacios mediante el aprendizaje de una representación latente continua (a través de un codificador-decodificador, típicamente un VAE) y la realización de la BO dentro de dicho espacio latente.

Si bien los modelos fundacionales tabulares (TFMs) y las redes de ajuste previo de datos (PFNs) han logrado un rendimiento de vanguardia en tareas de regresión al aprender la predicción bayesiana aproximada mediante el aprendizaje en contexto, su aplicación directa a LSBO enfrenta un desajuste crítico. Los TFMs actuales han sido preentrenados en grandes colecciones de tareas de regresión sintéticas que no reflejan la geometría específica de los espacios latentes ni la naturaleza de los objetivos moleculares (p. ej., logP, QED). En consecuencia, el "prior implícito" aprendido por estos modelos durante el preentrenamiento no se alinea con el entorno de despliegue de LSBO, donde el mapa de la clave latente al valor del objetivo está moldeado por el decodificador y las propiedades específicas del dominio. Tratar a los TFMs como regresores genéricos de tipo plug-in impide aprovechar todo su potencial en este entorno especializado.

Metodología
Los autores proponen LILBO (Latent In-context Learning for Bayesian Optimization), un marco que adapta un modelo fundacional tabular preentrenado (específicamente TabPFN-3) al dominio de LSBO mediante un preentrenamiento continuo. La metodología consta de dos componentes principales:

1. Generación de Tareas Sintéticas Conscientes de la BO:
   En lugar de utilizar priors de regresión genéricos, los autores construyen una distribución de preentrenamiento sintética fundamentada en objetivos moleculares.

   • Tareas Base: Los objetivos sintéticos se derivan de una biblioteca de tareas moleculares base (p. ej., logP, QED, similitud, redescubrimiento) extraídas de la base de datos GuacaMol.
   • Construcción del Objetivo: Para cada episodio de entrenamiento, se selecciona un subconjunto de tareas base y se crea un objetivo escalar sintético mediante combinaciones lineales (ponderadas por distribuciones de Dirichlet) o combinaciones no lineales (usando MLPs aleatorios).
   • Muestreo de Contexto: Para imitar el comportamiento de la BO, que se enfoca en regiones de alto valor, el contexto de los datos no se muestrea de forma uniforme. En su lugar, se induce una distribución sesgada hacia el valor (distribución de Boltzmann) sobre el conjunto molecular, enfatizando las regiones prometedoras del espacio latente del VAE mientras se mantiene la diversidad.

2. Preentrenamiento Continuo con Regularización:
   Los autores inicializan el modelo de sustitución (surrogate) desde un checkpoint de TabPFN-3 preentrenado en lugar de entrenar desde cero. Esto preserva el amplio prior de regresión aprendido durante el preentrenamiento a gran escala.

   • Objetivo: El modelo se somete a un preentrenamiento continuo sobre los episodios sintéticos de LSBO molecular.
   • Regularización: Para evitar el olvido catastrófico y la especialización excesiva, se aplica un regularizador L2-SP (L2-Soft Parameter). Esto ancla los pesos del modelo adaptado ($\theta$) a los pesos del checkpoint original ($\theta_0$), asegurando que el modelo conserve sus capacidades de regresión de propósito general mientras se adapta a la geometría específica de los espacios latentes moleculares.
   • Bucle de Entrenamiento: El modelo se entrena durante 8,000 pasos en aproximadamente 500,000 episodios de LSBO molecular.

Contribuciones Clave

• Adaptación Específica para LSBO: El artículo introduce un método para adaptar los sustitutos de aprendizaje en contexto específicamente para la optimización en el espacio latente, abordando el desajuste distributivo entre el preentrenamiento de TFMs genéricos y el despliegue de LSBO.
• Diseño de Distribución Sintética: Los autores diseñan una canalización de generación de tareas sintéticas que incorpora tareas base moleculares y un muestreo de contexto sesgado por valor para alinearse mejor con los requisitos de la BO (centrarse en regiones de alto valor).
• Aprendizaje Continuo Regularizado: El uso de un objetivo de preentrenamiento continuo regularizado permite que el modelo se especialice en LSBO molecular sin descartar los amplios priors del modelo fundacional.
• Diagnósticos: El artículo introduce diagnósticos en el espacio latente y en el espacio de objetivos (vía UMAP) para verificar que la distribución de preentrenamiento sintético cubra las regiones relevantes del manifold latente y del espacio de objetivos.

Resultados
El modelo propuesto, LILBO, fue evaluado en pruebas de optimización molecular de la suite GuacaMol (p. ej., Osimertinib, Fexofenadine, Ranolazine), que no fueron vistos durante el preentrenamiento.

• Rendimiento: LILBO logró un sólido desempeño, superando al baseline de conexión directa de TabPFN-3 y a varios baselines especializados de LSBO (incluyendo GP-LSBO, LOLBO, CoBO e InvBO) en términos de ranking promedio a través de las tareas.
• Eficiencia: La mejora respecto al baseline estándar de TabPFN-3 fue modesta pero notable, lograda con una adaptación relativamente ligera (500k episodios) en comparación con la escala masiva del preentrenamiento del modelo fundacional original (billones de tokens).
• Viabilidad: Los resultados demuestran que los sustitutos probabilísticos en contexto son una alternativa viable a los Procesos Gaussianos en los flujos de trabajo modernos de LSBO, siempre que el prior de preentrenamiento esté apropiadamente alineado con el dominio.

Significancia y Reivindicaciones
El artículo sostiene que, para que los modelos fundacionales tabulares sean efectivos en LSBO, su distribución de adaptación debe reducir el desajuste entre las tareas de preentrenamiento genéricas y la estructura específica de los objetivos del espacio latente y los historiales de BO encontrados en el despliegue. Los autores argumentan que la adaptación específica de LSBO es crucial para desbloquear todo el potencial de los sustitutos de aprendizaje en contexto en el diseño científico.

El trabajo reconoce limitaciones, señalando que el enfoque actual asume un VAE congelado durante el preentrenamiento, mientras que el despliegue a menudo implica el reentrenamiento periódico del VAE (deriva latente). Las direcciones futuras sugeridas incluyen simular trayectorias de optimización cortas para los contextos de entrenamiento, considerar la deriva latente en el prior, extender el método a la optimización multiobjetivo y aplicar el marco a otros dominios estructurados como el diseño de proteínas.