Resumen Técnico: Sobre el Espacio de Diseño de la Adaptación en Línea de la Difusión Discreta para la Optimización Molecular

Planteamiento del Problema

La optimización molecular comienza típicamente con un modelo generativo preentrenado (por ejemplo, un modelo de difusión discreta) que captura un prior amplio sobre estructuras químicas válidas. Sin embargo, el objetivo práctico en el tiempo de prueba no es muestrear de este prior, sino identificar moléculas que maximicen una función de recompensa específica de la tarea $r(x)$ (por ejemplo, afinidad de unión o aptitud proteica) utilizando un presupuesto limitado de oráculo. El oráculo, que evalúa la recompensa, suele ser costoso (por ejemplo, ensayos de laboratorio húmedo o procesos de acoplamiento computacional costosos).

El desafío central es la adaptación del modelo en línea: dirigir un generador preentrenado $p_{\theta_0}$ hacia regiones de alta recompensa del espacio químico manteniendo la validez y la diversidad. Esto implica un ciclo estrechamente acoplado donde el sistema debe decidir:

1. Adquisición: Qué candidatos evaluar dado el presupuesto del oráculo.
2. Adaptación: Cómo convertir las recompensas observadas en actualizaciones del modelo.
3. Reutilización de Retroalimentación: Qué datos históricos retener para el entrenamiento.
4. Eliminación de Sesgo del Prior (Prior Debiasing): Qué tanto alejar el generador de sus modos preentrenados sin colapsar la validez.

Si bien estos componentes han sido estudiados de forma aislada, no está claro si se complementan, interfieren o se vuelven redundantes cuando se integran en un ciclo completo de adaptación en línea.

Metodología

Los autores proponen un estudio controlado de un "arnés" modular para la adaptación en línea de modelos de difusión discreta. El ciclo (Figura 1) alterna entre el muestreo de candidatos, la selección de un lote para la evaluación del oráculo y el ajuste fino (fine-tuning) del generador. El estudio evalúa cinco elecciones de diseño de "plug-in" específicas (Tabla 1) a través de nueve tareas: seis tareas de afinidad de unión de moléculas pequeñas y tres tareas de aptitud proteica.

Componentes Clave de Diseño

1. Direccionamiento (Adquisición):

   • Muestreo de Thompson: En lugar de una selección codiciosa basada en la predicción de la media de un modelo subrogante, los autores utilizan el muestreo de Thompson para seleccionar candidatos. Esto equilibra la exploración y la explotación mediante el muestreo de estimaciones de recompensa $\hat{r}_i = \mu_i + \epsilon_i \sigma_i$, donde $\sigma_i$ representa el desacuerdo del ensamble (incertidumbre). Esto incentiva la evaluación de candidatos con alta incertidumbre incluso si su predicción media es menor.

2. Adaptación (Modelado de Recompensa y Eliminación de Sesgo):

   • Modelado de Recompensa por CVaR: Para centrar las actualizaciones en la cola de alta recompensa en lugar del promedio, la recompensa de entrenamiento se moldea utilizando el Valor en Riesgo Condicional (CVaR). La recompensa se transforma a $\tilde{r}(x) = (r(x) - Q_\tau)_+ / (1-\tau)$, donde $Q_\tau$ es el $\tau$-cuantil (establecido en 0.8). Esto concentra las actualizaciones de gradiente en las muestras de mejor desempeño.
   • Regularización de Entropía de Densidad (Eliminación de Sesgo del Modelo): Para evitar que el modelo colapse en los modos de alta densidad ya existentes del prior, se añade un regularizador: $\log \tilde{r}_{debias}(x) = \log r(x) - \gamma \log p_{\theta_t}(x)$. Dado que calcular la verosimilitud marginal exacta $\log p_{\theta_t}(x)$ para la difusión discreta es intratable, los autores utilizan un estimador de Monte Carlo computacionalmente eficiente basado en el ELBO de la difusión (una aproximación de un solo paso) para aproximar el logaritmo negativo de la verosimilitud.

3. Correcciones de Desajuste de Objetivo:

   • Buffer de Replay (Replay Buffer): Para evitar que el modelo olvide los descubrimientos de alta recompensa realizados en rondas tempranas a medida que la distribución de muestreo cambia, se utiliza un buffer FIFO de capacidad fija con expulsión por prioridad (eliminando primero las muestras de menor recompensa). Los mini-lotes se extraen mediante muestreo estratificado para asegurar que los valores atípicos de alta recompensa no sean diluidos.
   • Penalización de Salida Inválida: Para mantener la validez (por ejemplo, cadenas SMILES válidas), los candidatos que fallan en el parseo (por ejemplo, por RDKit) reciben una penalización negativa fija $r_{inv}$. Esto atrae al generador de vuelta hacia el manifold molecular válido cuando la exploración produce estructuras inválidas.

Configuración Experimental

• Tareas: Optimización de moléculas pequeñas para la unión a seis objetivos proteicos (usando GenMol como generador y FlashAffinity/Boltz-2 como oráculos) y optimización de la aptitud proteica para tres familias (GB1, CreiLOV, TrpB).
• Líneas Base (Baselines): El ciclo completo en línea se compara contra:
  • Ajuste fino fuera de línea (entrenamiento sobre un conjunto de datos estático).
  • Búsqueda en tiempo de inferencia (Beam Search, MCTS, Best-of-N, etc.) con un generador fijo.
  • Versiones con ablación del ciclo en línea (eliminando un componente a la vez).
• Métricas: El desempeño se mide mediante la mejor recompensa encontrada, las medias de top-k y la recompensa normalizada, evaluadas tanto contra llamadas acumuladas al oráculo como contra horas de GPU. Las métricas secundarias incluyen validez, diversidad (Tanimoto), QED y accesibilidad sintética (SA).

Resultados Clave

1. Complementariedad de los Componentes

El estudio encuentra que los componentes propuestos proporcionan beneficios complementarios, particularmente para las tareas de moléculas pequeñas donde los candidatos de alta recompensa se encuentran lejos del prior preentrenado.

• Moléculas Pequeñas: La adquisición (muestreo de Thompson), el modelado de CVaR y la Regularización de Entropía de Densidad ofrecen las ganancias de recompensa más claras. El muestreo de Thompson mejora la asignación de llamadas al oráculo; el CVaR se enfoca en la cola de la recompensa; y la eliminación de sesgo empuja la exploración lejos de los modos de alta ocupación del prior.
• Aptitud Proteica: Los efectos de la ablación son generalmente más débiles. Los priors específicos de la familia para proteínas ya cubren regiones de alta aptitud, lo que hace que los grandes cambios de distribución sean menos críticos. Sin embargo, el muestro de Thompson aún proporciona ganancias consistentes al mejorar la asignación del oráculo.

2. Mecanismos de Estabilidad

• Los Buffers de Replay y las Penalizaciones de Validez actúan como estabilizadores en lugar de impulsores primarios de recompensa. El buffer de replay alinea los gradientes con la distribución promediada en el tiempo de los datos evaluados, evitando la pérdida de descubrimientos de alta recompensa tempranos. La penalización de validez evita el "hackeo de recompensa" (reward hacking), donde el modelo genera cadenas químicamente inválidas pero con puntuaciones altas, un modo de falla que ocurre rápidamente sin la penalización.

3. Adaptación en Línea vs. Búsqueda en Tiempo de Inferencia

El ciclo completo en línea supera significativamente a las líneas base de búsqueda en tiempo de inferencia (por ejemplo, Beam Search, MCTS) bajo presupuestos iguales de llamadas al oráculo y de horas de GPU.

• Eficiencia: En las tareas de moléculas pequeñas, el ciclo en línea alcanza la puntuación final de la mejor línea base de búsqueda en aproximadamente 58 veces menos llamadas al oráculo. En las tareas de proteínas, la aceleración es de 16 veces.
• Razonamiento: Los métodos de búsqueda gastan el presupuesto generando y puntuando candidatos de un generador fijo. En contraste, la adaptación en línea reutiliza la retroalimentación del oráculo para actualizar el propio generador, multiplicando el valor de cada llamada al oráculo. Los métodos de búsqueda también consumen un cómputo significativo para la generación, dejando poco tiempo para los pasos de gradiente de ajuste fino dentro de un presupuesto de tiempo de ejecución fijo.

4. Análisis de Cambio de Distribución

Las ganancias de desempeño son más pronunciadas cuando la solución óptima requiere un gran cambio respecto al prior preentrenado.

• La Regularización de Entropía de Densidad es el principal motor del cambio de distribución, empujando la masa del generador lejos de las regiones de alta densidad del prior hacia áreas de baja densidad y alta recompensa.
• CVaR y el Muestreo de Thompson facilitan este cambio condicionalmente, dependiendo de dónde se encuentran los candidatos de alta recompensa respecto al prior.
• Las tareas de proteínas muestran cambios de distribución más débiles porque los priors específicos de la familia ya están concentrados cerca de las regiones de alta aptitud.

Significancia y Reivindicaciones

El artículo afirma proporcionar una receta práctica para la optimización de moléculas con eficiencia de retroalimentación: ajuste fino en línea con adquisición, modelado de recompensa, eliminación de sesgo, replay y control de validez.

• Novedad: Este trabajo integra y evalúa sistemáticamente estos componentes dentro de un único marco de difusión discreta, demostrando que no son redundantes sino que se refuerzan mutuamente.
• Eficiencia: El método propuesto supera tanto al ajuste fino fuera de línea como a las líneas base de búsqueda en tiempo de inferencia, particularmente en escenarios donde el oráculo es costoso y las moléculas objetivo están fuera de la distribución del modelo preentrenado.
• Robustez: La receta se muestra robusta a través de diferentes objetivos de ajuste fino (DDPP-LB y VIDD) y diferentes dominios (moléculas pequeñas frente a proteínas), aunque la magnitud de las ganancias varía según la distancia entre el prior y la distribución objetivo.

Los autores concluyen que, si bien los componentes son más críticos para el diseño de moléculas pequeñas (donde el prior es amplio y el objetivo es distante), el marco proporciona un enfoque generalizable para dirigir modelos generativos discretos en tareas de descubrimiento científico con retroalimentación limitada.