JADAI: Jointly Amortizing Adaptive Design and Bayesian Inference

Resumen Técnico: JADAI – Diseño Adaptativo y la Inferencia Bayesiana Amortizadas Conjuntamente

Planteamiento del Problema
Muchos dominios científicos e ingenieriles implican la inferencia de parámetros no observados $\theta$ a partir de observaciones $x$ generadas por simuladores complejos. Si bien la Inferencia Basada en Simulación (SBI, por sus siglas en inglés) ha avanzado en la capacidad de aproximar distribuciones posteriores $p(\theta|x)$ sin evaluaciones explícitas de la verosimilitud, estos métodos suelen asumir un diseño experimental fijo. Por el contrario, el Diseño Experimental Bayesiano (BED) y el Diseño Adaptativo (BAD) se centran en seleccionar variables de diseño $\xi$ para maximizar la ganancia de información (por ejemplo, la Ganancia de Información Esperada, EIG), pero a menudo tratan la inferencia como un paso separado y computacionalmente costoso que se realiza después de la recolección de datos. Los enfoques existentes suelen desacoplar estas tareas: los métodos de BAD suelen depender de una inferencia no amortizada y lenta o de familias paramétricas simples, mientras que los métodos de SBI no optimizan la adquisición de datos. Esta separación conduce a una ganancia de información y una eficiencia computacional subóptimas, particularmente en entornos de alta dimensionalidad o multimodales.

Metodología
Los autores presentan JADAI (Diseño Adaptativo e Inferencia Bayesiana Amortizados Conjuntamente), un marco que unifica el diseño adaptativo y la inferencia posterior en un único sistema entrenable de extremo a extremo. La innovación central es la amortización conjunta de tres componentes:

1. Una Red de Política ($\pi_\phi$): Propone el siguiente diseño $\xi_t$ basado en el historial experimental actual.
2. Una Red de Historial/Resumen ($\eta_\omega$): Codifica la secuencia de pares diseño-observación pasados $\{(\xi_k, x_k)\}_{k=1}^t$ en un estadístico de resumen de dimensión fija $h_t$.
3. Una Red de Inferencia ($q_\psi$): Aproxima la distribución posterior $p(\theta|h_t)$ en cada paso de la secuencia experimental.

Objetivo de Optimización
JADAI minimiza una función de pérdida genérica que agrega reducciones incrementales del error posterior a lo largo de secuencias experimentales. En lugar de maximizar la Ganancia de Información Esperada (EIG) directamente —lo cual requiere evaluaciones de densidad—, el marco utiliza una suma telescópica de pérdidas posteriores.

• El objetivo de entrenamiento se deriva de un límite inferior variacional de la Ganancia de Información Esperada Total (TEIG).
• La pérdida en el paso $t$, $\ell_t$, está definida por el estimador de la posterior (por ejemplo, pérdida de difusión o de ajuste de flujo/flow matching).
• La función de utilidad $u_T$ se construye como una suma de diferencias entre pérdidas consecutivas: $u_T = \sum (\ell^*_{t-1} - \ell_t)$, donde $\ell^*$ denota una versión desacoplada (sin gradiente) de la pérdida anterior.
• Esta formulación asegura que la red de política y la red de resumen sean entrenadas para maximizar la reducción de la pérdida posterior en cada paso, maximizando efectivamente la ganancia de información sin requerir evaluaciones explícitas de la densidad de la posterior.

Opciones Arquitectónicas Clave

• Modelos de Difusión: La red de inferencia $q_\psi$ se instancia utilizando modelos de difusión (o ajuste de flujo), lo que permite al marco aproximar distribuciones posteriores de alta dimensión y multimodales en cada paso experimental, una capacidad que carecen los métodos de BAD anteriores restringidos a flujos de normalización.
• Estabilidad del Entrenamiento: Para evitar la retropropagación a través del tiempo (BPTT) anidada y asegurar un entrenamiento estable, el marco emplea:
  • Estados de Historial Desacoplados: Los diseños se generan a partir de un estado de resumen desacoplado para evitar rutas de gradiente cíclicas entre la red de política y la red de historial.
  • Aprendizaje por Currículo: La longitud del despliegue (rollout) se incrementa gradualmente durante el entrenamiento, y se utiliza una estrategia de "mezcla de prior de diseño" al inicio del entrenamiento para exponer las redes a diseños diversos antes de depender únicamente de la política aprendida.

Resultados
JADAI fue evaluado en tres pruebas de referencia que representan una complejidad creciente en la estructura de la política y la posterior:

1. Localización de Ubicación (LF): Una tarea que requiere la inferencia de ubicaciones de fuentes a partir de mediciones de intensidad con ruido. JADAI logró puntuaciones de sPCE (Estimación Contrastiva de Prior Secuencial) superiores o competitivas respecto a las líneas base del estado del arte (por ejemplo, DAD, iDAD, RL-BOED), particularmente para horizontes largos ($T > 10$). La política aprendida equilibró eficazmente la exploración y la explotación, y la posterior basada en difusión capturó correctamente la naturaleza multimodal e intercambiable de las fuentes.
2. Sustitución Elástica Constante (CES): Un problema de diseño desafiante que involucra la inferencia de preferencias donde los diseños informativos se encuentran en un estrecho "punto ideal". JADAI superó a todas las líneas base existentes en el horizonte estándar $T=10$, demostrando robustez en escenarios donde los diseños aleatorios no son informativos.
3. Descubrimiento de Imágenes MNIST: Una tarea de alta dimensión donde el objetivo es reconstruir una imagen a partir de mediciones parciales. JADAI alcanzó puntuaciones de Índice de Similitud Estructural (SSIM) y Error Cuadrático Medio Normalizado (NRMSE) de vanguardia, superando a la línea base CoDiff. El método identificó con éxito las clases de dígitos tempranamente y refinó los detalles locales, manteniendo el rendimiento incluso bajo ruido de medición aditivo.

Significancia y Reivindicaciones
El artículo afirma que JADAI representa un paso significativo al unificar el diseño y la inferencia, eliminando así la necesidad de pasos de inferencia separados y costosos durante el proceso adaptativo. Las contribuciones clave incluyen:

• Amortización Conjunta: Es el primer marco que amortiza conjuntamente el diseño adaptativo y la inferencia posterior utilizando una utilidad general que agrega reducciones incrementales del error posterior, evitando la necesidad de evaluaciones de densidad.
• Escalabilidad a Posteriors Complejas: Al aprovechar los modelos de difusión, JADAI puede manejar posteriores de alta dimensión y multimodales en cada paso experimental, superando las limitaciones de expresividad de los métodos anteriores que dependían de flujos de normalización o mezclas gaussianas.
• Rendimiento: El marco logra un rendimiento superior o competitivo a través de diversas pruebas de referencia, incluyendo la reconstrucción de imágenes de alta dimensión, manteniendo la eficiencia computacional durante el despliegue (los despliegues se completan en milisegundos).

Los autores señalan que, aunque JADAI depende actualmente de simuladores diferenciables, el enfoque es muy adecuado para el creciente panorama de los marcos de simulación diferenciable. Reconocen que el muestreo de la posterior sigue siendo un cuello de botella para el despliegue de sub-segundo en casos de alta dimensión, pero sugieren la destilación post-hoc como una posible dirección futura. El trabajo no pretende resolver la optimización de simuladores de caja negra sin gradientes, sino que se posiciona como una solución robusta para entornos diferenciables donde los procesos clásicos son computacionalmente impracticables.