Interpretable epistemic uncertainty decomposition in sequential generative models via polynomial chaos surrogates

Resumen Técnico: Descomposición Interpretable de la Incertidumbre Epistémica en Modelos Generativos Secuenciales mediante Surrogados de Caos Polinomial

Enunciado del Problema

Los modelos generativos secuenciales, como las Redes de Flujo Generativo (GFlowNets), son cada vez más centrales en el descubrimiento científico impulsado por la IA, la optimización combinatoria y el diseño biológico. Estos modelos a menudo operan bajo condiciones de incertidumbre epistémica, específicamente cuando las funciones de recompensa que guían su generación son en sí mismas inciertas (por ejemplo, surrogados de redes neuronales entrenados con datos experimentales limitados). Si bien los métodos existentes, como los conjuntos profundos (deep ensembles) y las redes neuronales bayesianas, pueden cuantificar la magnitud de la incertidumbre de predicción, no logran proporcionar una descomposición interpretable de dicha incertidumbre. En consecuencia, los practicantes no pueden determinar qué componentes específicos de la incertidumbre de la recompensa de entrada impulsan qué decisiones generativas, ni pueden distinguir entre señales químicas robustas y artefactos estadísticos derivados de sesgos en el conjunto de entrenamiento.

Metodología

Los autores proponen un marco de modelado de surrogados que propaga la incertidumbre epistémica a través de las GFlowNets utilizando Expansiones de Caos Polinomial (PCE). El enfoque procede en tres etapas:

1. Entrenamiento y Parametrización del Conjunto: Se entrena un modesto conjunto de GFlowNets bajo condiciones de recompensa variadas. El espacio de la función de recompensa se parametriza en un espacio latente de baja dimensión $\mu \in \mathbb{R}^d$ utilizando Análisis de Componentes Principales (PCA) o expansiones analíticas (por ejemplo, Karhunen–Loève). Se prefiere el PCA porque sus componentes ortogonales aproximan entradas gaussianas independientes, un prerrequisito para una descomposición analítica de Sobol válida.
2. Ajuste del Surrogado PCE: Para cada acción $k$ en el paso de trayectoria $t$, se ajusta un PCE para mapear los parámetros de recompensa de baja dimensión $\mu$ a las probabilidades de acción transformadas mediante log-ratio (utilizando una transformación Log-Ratio Aditiva relativa a una acción de referencia). Los coeficientes se estiman mediante regresión de crestas regularizada.
3. Descomposición Analítica: Los coeficientes ajustados del PCE permiten el cálculo analítico de los índices de sensibilidad de Sobol. A diferencia de los métodos de Monte Carlo, que requieren una integración costosa, estos índices se derivan directamente de los coeficientes polinomiales, proporcionando una descomposición exacta basada en la varianza de cuánto contribuye cada componente de la incertidumbre de la recompensa a la varianza en la distribución de la política.

El marco incluye garantías formales de convergencia para el surrogado y sus índices de Sobol. Cuatro de las cinco propiedades teóricas clave (convergencia del PCE, convergencia del índice de Sobol, preservación del simplex y límites de propagación de incertidumbre) han sido verificadas formalmente en el asistente de pruebas Lean 4.

Contribuciones Clave

1. Marco General: Un método computacionalmente eficiente para propagar y descomponer la incertidumbre epistémica en modelos generativos secuenciales condicionados utilizando surrogados PCE ajustados a pequeños conjuntos.
2. Verificación Formal: Garantías teóricas de convergencia para el surrogado y sus índices de sensibilidad, con cuatro de cinco teoremas verificados formalmente en Lean 4.
3. Descomposición de Sensibilidad Interpretable: El primer método analítico para descomponer la incertidumbre de la política en modelos generativos, revelando exactamente qué componentes de la recompensa impulsan decisiones generativas específicas. Esta capacidad no está disponible en conjuntos profundos estándar, redes neuronales bayesianas o dropout de Monte Carlo.
4. Validación Empírica: Demostraciones en siete tareas de complejidad creciente, incluidas tres aplicaciones del mundo real y cuatro tareas de validación sintéticas.

Resultados y Aplicaciones

El marco fue evaluado en tres dominios principales del mundo real:

• Optimización de Reacciones Químicas (Buchwald–Hartwig): En el conjunto de datos Doyle–Dreher, el marco analizó la selección de catalizador, base, haluro de arilo y aditivo. Reveló que la selección de catalizador es robusta (índice de Sobol bajo, $D \approx 71$), indicando una señal química genuina, mientras que la selección de aditivo es frágil (índice de Sobol alto, $D \approx 179$, 2.5$\times$ más alto), indicando susceptibilidad al sesgo del conjunto de entrenamiento. Esto aborda directamente el "debate Chuang–Keiser" sobre los artefactos estadísticos en los predictores de rendimiento.
• Descubrimiento Causal Bayesiano (Señalización de Proteínas Sachs): Aplicado a la red Sachs, el análisis mostró que, aunque todas las aristas exhibieron alta varianza debido a datos limitados (200 observaciones), la incertidumbre se descompuso en regímenes distintos. Las aristas de la vía MAPK fueron impulsadas por un componente principal de la variación de covarianza, mientras que las aristas del centro PKA/PKC fueron impulsadas por otro. Esto proporciona un mapa dirigido para experimentos de perturbación.
• Diseño Molecular Basado en Fragmentos: En una tarea de diseño de similitud con fármacos, el marco revirtió la suposición convencional de que los andamios son robustos y las decoraciones son frágiles. Los resultados mostraron que la posición del conector fue el componente más frágil ($D \approx 28.2$), mientras que las posiciones de andamio y decoración fueron más robustas. Esto sugiere que los esfuerzos experimentales deberían priorizar la diversidad de conectores.

En tareas sintéticas (mundos de cuadrícula, regresión simbólica y proxies de razonamiento de LLM), el surrogado PCE logró una cobertura de calibración del 95% entre 0.97 y 1.00 para los pasos dominantes. Computacionalmente, el surrogado evaluó 10,000 muestras de política en milisegundos, ofreciendo una aceleración de $10^3$ a $10^4\times$ en comparación con el reentrenamiento exhaustivo del conjunto.

Significado y Afirmaciones

El artículo afirma que este marco ofrece una capacidad única: descomposición interpretable de la incertidumbre. Mientras que otros métodos pueden afirmar que una predicción es incierta, este enfoque explica por qué atribuyendo la varianza a componentes específicos de la recompensa. Esto permite a los científicos:

• Distinguir entre señales robustas y artefactos estadísticos en los datos de entrenamiento.
• Identificar qué partes de un espacio de diseño (por ejemplo, conectores frente a andamios) requieren validación experimental adicional.
• Asignar recursos experimentales de manera eficiente al apuntar a las fuentes específicas de incertidumbre epistémica.

Los autores enfatizan que el método es agnóstico a la arquitectura de la GFlowNet y al objetivo de entrenamiento. Señalan limitaciones, como la dependencia de una estructura de bajo rango en el paisaje de recompensas para que el PCA sea efectivo y la necesidad de volver a ajustar el surrogado para diferentes trayectorias (aunque no de reentrenar el conjunto). El trabajo se posiciona como un puente entre la cuantificación de incertidumbre en ingeniería computacional y la IA generativa moderna, proporcionando un camino escalable para comprender la "caja negra" de la toma de decisiones secuenciales bajo incertidumbre.