CMAD: Cooperative Multi-Agent Diffusion via Stochastic Optimal Control

Resumen Técnico: CMAD – Difusión Cooperativa Multiagente mediante Control Óptimo Estocástico

Enunciado del Problema
Los modelos generativos de tiempo continuo, particularmente los modelos de difusión basados en puntuación, han logrado un rendimiento de vanguardia en la síntesis y restauración de imágenes. Aunque existen técnicas para la generación controlada (por ejemplo, guía mediante clasificador, modelado de puntuación condicional), persiste un desafío significativo en la generación composicional: la agregación efectiva de múltiples modelos de difusión preentrenados para satisfacer objetivos complejos y específicos de la tarea.

Los enfoques actuales tratan mayoritariamente la composición como una combinación algebraica de densidades de probabilidad (por ejemplo, productos de expertos o mezclas). Estos métodos asumen que la distribución objetivo es una función algebraica explícita de las marginales de los componentes (por ejemplo, $p_{prod}(x) \propto \prod q_i(x)$). Sin embargo, los autores argumentan que esta suposición a menudo es poco realista; por ejemplo, no existe garantía teórica de que combinar un modelo para el realismo visual con uno para la alineación de texto resulte en un promedio geométrico simple o un producto de densidades. Además, el muestreo de tales distribuciones de producto explícitas a menudo requiere acceso a la puntuación de la distribución combinada, la cual no se descompone linealmente, lo que conduce a aproximaciones sesgadas y a una calidad de generación deficiente.

Metodología
El artículo propone CMAD (Difusión Cooperativa Multiagente), un cambio de paradigma que formula la generación composicional como un problema de Control Óptico Estocástico (COE) Cooperativo en lugar de un problema de composición de densidades.

1. Formulación Multiagente:

   • Los modelos de difusión preentrenados se tratan como agentes independientes.
   • Cada agente $i$ sigue una Ecuación Diferencial Estocástica (EDE) controlada que representa su dinámica en tiempo inverso:
     $$dX_t^{u,i} = [b_i(X_t^{u,i}, t) + g(t)u_i(X_t^{u,i}, t; \{X_t^{u,j}\}_j)]dt + g(t)dW_t$$
     Aquí, $b_i$ es la deriva del modelo preentrenado, y $u_i$ es una política de control aprendida. Crucialmente, el control $u_i$ depende de los estados de todos los demás agentes, creando un sistema de EDE acoplado.

2. Agregación y Objetivo:

   • El objeto generado final se define como un estado agregado $Y_t = \phi(\{X_t^{u,i}\}_{i=1}^N, t; \vartheta)$, donde $\phi$ es un agregador (por ejemplo, enmascaramiento/concatenación sin superposición) y $\vartheta$ son parámetros aprendibles.
   • El objetivo es minimizar una función objetivo $J$ sobre los controles $\{u_i\}$ y los parámetros del agregador $\vartheta$:
     $$J = \mathbb{E}\left[ \int_0^T \left( \sum_{i=1}^N \lambda_i \|u_i\|^2 + c(Y_t, t) \right) dt + \Psi(Y_T) \right]$$
   • Costo Terminal ($\Psi$): Codifica la recompensa específica de la tarea (por ejemplo, la verosimilitud negativa de un clasificador para un dígito objetivo).
   • Costo de Ejecución ($c$): Actúa como estabilizador y sustituto del costo terminal, evaluado sobre una "estimación de Tweedie" (una visión previa desruidificada) del estado agregado para proporcionar gradientes densos a lo largo de la trayectoria.
   • Costo de Control: El término cuadrático $\|u_i\|^2$ regulariza el control, manteniendo a los agentes cerca de sus dinámicas en tiempo inverso preentrenadas.

3. Esquema de Optimización (Descenso por Control):

   • Los autores proponen un algoritmo de optimización basado en la Optimización de Difusión Iterativa (IDO).
   • En lugar de actualizar todos los controles conjuntamente (lo cual puede ser computacionalmente costoso), emplean un esquema de descenso por coordenadas. En cada iteración externa, se actualiza el control $u_i$ de un agente mientras los demás permanecen congelados.
   • La actualización para un solo agente se realiza mediante simulación Monte Carlo de las trayectorias de las EDE acopladas, calculando los gradientes de la pérdida empírica con respecto a los parámetros de control y realizando un paso de descenso de gradiente.
   • La política de control $u_i$ se parametriza con un sesgo inductivo informado por recompensa, combinando una corrección de red neuronal aprendida con un término de guía de gradiente explícito derivado de la estimación de Tweedie del estado agregado.

Contribuciones Clave

1. Marco Novel: Introducción de un marco multiagente cooperativo que plantea la inferencia con múltiples modelos de difusión preentrenados como un problema de COE, evitando la necesidad de definiciones algebraicas explícitas de la densidad objetivo.
2. Algoritmo de Optimización: Propuesta de un esquema de optimización por control basado en IDO, que permite un control multiagente escalable resolviendo iterativamente subproblemas de COE de agente único.
3. Composición Implícita: Demostración de que objetivos complejos (definidos por funciones de pérdida) pueden optimizarse sin conocer la forma de la densidad compuesta resultante.

Resultados Experimentales
El marco se validó en una tarea de prueba de concepto utilizando la generación de dígitos MNIST, donde múltiples agentes eran responsables de diferentes franjas horizontales de la imagen.

• Configuración: Los agentes controlaban regiones no superpuestas de la imagen. El costo terminal fue la verosimilitud negativa de un clasificador MNIST preentrenado para un dígito objetivo, aumentado con una pérdida de continuidad de costura para garantizar la coherencia espacial.
• Líneas Base: Comparación contra una línea base ingenua en tiempo de inferencia (CDPS) que utiliza guía de gradiente sin control cooperativo aprendido.
• Rendimiento:
  • CMAD (tanto las variantes conjuntas como las de control por control) logró consistentemente una pérdida terminal menor ($\Psi$) en comparación con CDPS.
  • La precisión de clasificación fue alta en todos los métodos, pero CMAD produjo muestras visualmente más realistas.
  • CDPS ocasionalmente produjo artefactos antinaturales en las costuras entre agentes, mientras que CMAD generó dígitos coherentes.
  • El esquema de optimización por control logró un rendimiento competitivo (por ejemplo, 99.90% de precisión para el dígito 3 con 3 agentes) mientras ofrecía una alternativa escalable a la optimización conjunta.

Significado y Trabajo Futuro
El artículo afirma su importancia al proporcionar una alternativa flexible a la composición basada en densidades, permitiendo la optimización de objetivos complejos e implícitos definidos por funciones de pérdida. Los autores señalan que los resultados actuales son experimentos iniciales en MNIST.

El trabajo futuro esbozado en el artículo incluye:

• Escalar el marco a aplicaciones de alta dimensión (por ejemplo, imágenes naturales).
• Investigar las condiciones de convergencia del esquema de actualización por control.
• Extender la conexión a juegos diferenciales y dinámicas de juego ficticio.
• Explorar configuraciones donde el operador de agregación $\phi$ sea en sí mismo parametrizado y aprendido conjuntamente con los controles.

Los autores mantienen un tono modesto, presentando el trabajo como un paso fundamental hacia la difusión cooperativa multiagente en lugar de un problema resuelto para todas las tareas generativas.