Uncertainty Estimation for Molecular Diffusion Models

Resumen Técnico: Estimación de la Incertidumbre en Modelos de Difusión Molecular

Planteamiento del Problema
Los modelos de difusión se han convertido en el paradigma dominante para la generación de moléculas 3D. Sin embargo, existe un cuello de botella práctico crítico: estos modelos carecen de una señal fundamentada para identificar cuándo una molécula generada es químicamente inválida o inestable. Mientras que los modelos predictivos suelen utilizar la estimación de la incertidumbre para señalar salidas poco fiables, los métodos existentes para modelos de difusión se han centrado principalmente en el dominio de la imagen. En el dominio molecular, donde la evaluación posterior (por ejemplo, docking o validación en laboratorio húmedo) es costosa, depender de filtros rudimentarios basados en reglas o de la propia verosimilitud del modelo —que a menudo está mal calibrada— sigue siendo el estándar. Existe una falta de trabajos análogos que proporcionen señales de calidad de la muestra para modelos de difusión molecular.

Metodología
Los autores proponen un método post-hoc para estimar la incertidumbre por muestra para modelos de difusión molecular preentrenados sin reentrenar el generador. El enfoque se basa en la intucción de que las muestras sobre las cuales el modelo tiene incertidumbre interna exhibirán trayectorias de eliminación de ruido más variables.

1. Aproximación de Laplace: Para mantener la compatibilidad con cualquier modelo preentrenado, el método ajusta una aproximación de Laplace a los parámetros de la red de eliminación de ruido. Específicamente, se construye una distribución posterior $q(\theta) \approx \mathcal{N}(\theta | \hat{\theta}, \Sigma)$ alrededor de los parámetros preentrenados $\hat{\theta}$. La covarianza $\Sigma$ se aproxima utilizando la matriz de información de Fisher empírica (aplicada solo a las capas finales para asegurar la escalabilidad) más un término de precisión de la distribución previa.
2. Variabilidad de la Predicción de Ruido: Durante la trayectoria de generación, para un subconjunto seleccionado de pasos temporales $\mathcal{T}$, el método extrae $M$ muestras de pesos $\theta_m$ de la distribución posterior $q(\theta)$. Para cada muestra, el eliminador de ruido predice el ruido $\epsilon^m_t = f_{\theta_m}(\hat{x}_t, t)$.
3. Agregación de la Puntuación de Incertidumbre: La incertidumbre intermedia $u_t$ en cada paso temporal se define como la varianza de la muestra elemento a elemento de estas predicciones de ruido a través de las $M$ muestras. Esta incertidumbre por paso se agrega a través de los pasos temporales seleccionados, todos los átomos ($N_a$) y todas las dimensiones de características ($D$) para producir una única puntuación de incertidumbre escalar $u(\hat{x})$ para la molécula generada:
   $$u(\hat{x}) := \frac{1}{|\mathcal{T}| \cdot N_a \cdot D} \sum_{t \in \mathcal{T}} \sum_{n=1}^{N_a} \sum_{d=1}^{D} [u_t]_{n,d}$$
   Empíricamente, los autores encuentran que medir la incertidumbre solo en el 10% final del proceso de generación (donde $t$ es pequeño) es suficiente para el rendimiento, limitando la carga computacional.

Contribuciones Clave

• Primera Aplicación a la Difusión Molecular: Este trabajo extiende los enfoques de incertidumbre de difusión desde el dominio de la imagen al de la generación molecular, abordando la necesidad específica de señales de calidad de la muestra en dominios científicos con evaluación posterior costosa.
• Marco Post-hoc: El método no requiere el reentrenamiento del modelo de difusión subyacente, lo que lo hace aplicable a arquitecturas de vanguardia existentes como EDM y GeoLDM.
• Escalado en Tiempo de Prueba: El artículo demuestra cómo la puntuación de incertidumbre puede utilizarse como criterio de filtrado en un entorno de escalado de tiempo de prueba (sobremuestreo y filtrado) para mejorar la calidad de las moléculas generadas.

Resultados
Los autores evaluaron su enfoque en los conjuntos de datos QM9 y GEOM-Drugs utilizando modelos EDM y GeoLDM.

• Correlación con la Calidad: En el conjunto de datos QM9, la puntuación de incertidumbre propuesta exhibe una correlación negativa estadísticamente significativa con métricas de calidad establecidas (estabilidad molecular, estabilidad atómica y validez). Esto indica que puntuaciones de incertidumbre más altas corresponden a muestras de menor calidad.
• Comparación con la Verosimilitud: La puntuación de incertidumbre supera consistentemente a la log-verosimilitud negativa (NLL) del modelo de difusión como predictor de la calidad de la muestra. Por ejemplo, en EDM/QM9, la puntuación de incertidumbre logró una correlación de Spearman de $-0.284$ con la estabilidad molecular, comparado con $-0.150$ para NLL.
• Rendimiento del Escalado en Tiempo de Prueba: Al generar 20K moléculas y filtrar para mantener las 10K con la menor incertidumbre, los autores observaron mejoras significativas en la estabilidad molecular ($\sim10\%$), estabilidad atómica ($\sim1\%$) y validez ($\sim5\%$) para el modelo EDM. Estas ganancias fueron comparables a la mejora lograda al cambiar la arquitectura del modelo subyacente de EDM a GeoLDM.
• Estudios de Ablación:
  • Aproximación de Fisher: Eliminar la aproximación de Fisher (reemplazando la covarianza con una Gaussiana isotrópica) produjo un rendimiento casi idéntico, lo que sugiere que la puntuación actúa más como una medida de sensibilidad a pequeñas perturbaciones que como una verdadera incertidumbre epistémica bayesiana.
  • Selección de Pasos Temporales: La señal es más fuerte en el extremo "limpio" de la trayectoria de generación (pequeño $t$), validando la elección de centrar la medición de la incertidumbre en los pasos finales.
• Limitaciones: El método no produjo mejoras significativas en el conjunto de datos más grande y complejo GEOM-Drugs, lo que indica que el enfoque actualmente no se transfiere a conjuntos de datos con mayor complejidad conformacional.

Significancia y Reivindicaciones
El artículo afirma dar un "primer paso" hacia la extensión de los enfoques de incertidumbre de difusión al dominio molecular. La principal significancia radica en proporcionar una herramienta post-hoc práctica para filtrar moléculas generadas, mejorando así la eficiencia de los procesos de evaluación posteriores. Los autores posicionan su puntuación de incertidumbre como un reemplazo natural para los verificadores basados en NLL en métodos de escalado adaptativo en tiempo de prueba.

Los autores mantienen la modestia respecto a sus afirmaciones, señalando explícitamente que el enfoque actualmente falla en conjuntos de datos más grandes y complejos (GEOM-Drugs) y que el mecanismo subyacente parece ser una medida de sensibilidad más que una verdadera incertidumbre epistémica bayesiana. Sugieren que el trabajo futuro debería explorar el cálculo de la incertidumbre en un espacio de características "semánticas" utilizando codificadores moleculares preentrenados, de forma análoga a los éxitos recientes en la generación de imágenes de alta dimensión.