Unsupervised Decomposition and Recombination with Discriminator-Driven Diffusion Models

1. Problema y Contexto

El aprendizaje de representaciones y el modelado generativo son pilares fundamentales de la inteligencia artificial. Un objetivo clave es la desconexión (disentanglement): aprender espacios latentes donde cada factor de variación (ej. iluminación, forma del objeto, fondo, acción) esté codificado en dimensiones independientes.

• El Desafío: La descomposición no supervisada es teóricamente imposible sin suposiciones adicionales o señales de retroalimentación externas. Los modelos existentes, como los VAEs o los modelos de difusión composicionales (ej. Decomp Diffusion), a menudo fallan al reorganizar (recombinar) estos factores latentes de diferentes muestras.
• La Consecuencia: Cuando se mezclan factores latentes de dos imágenes o videos distintos para generar una nueva muestra, el resultado suele ser incoherente (artefactos visuales, violaciones de leyes físicas, escenas ilógicas) porque la distribución de los datos reales no soporta tales combinaciones arbitrarias.
• La Necesidad: Se requiere un mecanismo que guíe al modelo para que las recombinaciones de factores latentes sean físicamente y semánticamente plausibles, sin necesidad de etiquetas manuales (supervisión) para cada factor.

2. Metodología Propuesta

Los autores proponen un marco que combina modelos de difusión factorizados con una señal adversaria (discriminador) para mejorar tanto la descubrimiento de factores como la calidad de la generación composicional.

A. Modelo Generativo Factorizado

Se utiliza un modelo de difusión donde una imagen o video $x$ se codifica en $K$ componentes latentes independientes $z = (z_1, ..., z_K)$.

• El modelo de difusión aprende a predecir el ruido (o la imagen original) como el promedio de las contribuciones de cada factor latente:
  $$\epsilon_\theta(x_t, t) = \frac{1}{K} \sum_{k=1}^K \epsilon_\theta(x_t, t, z_k)$$
• Esto permite la recombinación: durante la inferencia, se pueden tomar factores $z_k$ de una imagen A y factores $z_j$ de una imagen B para crear un código latente híbrido $\tilde{z}$ y generar una nueva imagen.

B. Entrenamiento con Discriminador (La Innovación Clave)

Para evitar que las recombinaciones caigan fuera de la variedad de datos real (off-manifold), se introduce un discriminador $D_\psi$.

• Objetivo del Discriminador: Distinguir entre:
  1. Muestras de fuente única: Generadas a partir de latentes de una sola imagen.
  2. Muestras recombinadas: Generadas a partir de latentes mezclados de múltiples imágenes.
• Mecanismo de Entrenamiento:
  • El discriminador se entrena para maximizar la capacidad de distinguir estos dos casos.
  • El generador (y el codificador) se entrena adversarialmente para engañar al discriminador, minimizando la pérdida adversaria $L_{adv}$.
  • La pérdida total del generador combina la pérdida de reconstrucción estándar de difusión ($L_{MSE}$) y la pérdida adversaria:
    $$L_{gen} = L_{MSE} + \lambda L_{adv}$$
• Efecto: Al forzar al modelo a que las muestras recombinadas sean indistinguibles de las muestras reales, se incentiva que los factores latentes sean verdaderamente independientes y que sus combinaciones respeten la coherencia física y semántica del dominio.

C. Aplicación a Robótica (Video)

El método se extiende a videos de demostraciones robóticas.

• Se codifican secuencias de video en factores latentes que capturan primitivas de acción (ej. "agarrar", "mover", "soltar") y contextos de escena.
• Al recombinar estos factores, se generan nuevas trayectorias de tareas que no existían en los datos de entrenamiento, permitiendo una exploración más amplia del espacio de estados.

3. Contribuciones Clave

1. Señal Adversaria para Difusión: Introducen un nuevo mecanismo de entrenamiento que utiliza un discriminador para guiar la descomposición y recombinación en modelos de difusión, sin necesidad de supervisión a nivel de factores.
2. Mejora en Desconexión y Calidad: Demuestran que esta señal adversaria mejora tanto la calidad perceptual de las imágenes recombinadas como la desconexión de los factores latentes (medida por métricas como MIG y MCC).
3. Aplicación en Robótica: Aplican el método al benchmark LIBERO, mostrando que la recombinación de componentes de acción aprendidos genera trayectorias de video diversas que aumentan significativamente la cobertura del espacio de estados para la exploración de robots.

4. Resultados Experimentales

Los autores evaluaron su método en varios conjuntos de datos de imágenes y video:

• Imágenes (CelebA-HQ, Virtual KITTI, CLEVR, Falcor3D):

  • Calidad Perceptual: El método propuesto superó consistentemente a la línea base Decomp Diffusion, logrando puntuaciones FID (Fréchet Inception Distance) significativamente más bajas (mejores). Por ejemplo, en CelebA-HQ, el FID bajó de 82.70 a 43.98.
  • Desconexión: Se observaron mejoras en las métricas de desconexión MIG (Mutual Information Gap) y MCC (Mean Correlation Coefficient), indicando que los factores aprendidos son más independientes y semánticamente significativos.
  • Cualitativo: Las imágenes recombinadas mostraron coherencia estructural (ej. iluminación y objetos consistentes), mientras que la línea base a menudo producía artefactos o superposiciones inconsistentes.

• Robótica (LIBERO):

  • Se evaluó la cobertura del espacio de estados mediante la ejecución de trayectorias generadas.
  • El método con discriminador logró un aumento masivo en la exploración: +35.8% en Scene 5 y +73.8% en Scene 6 en comparación con la línea base sin discriminador.
  • Esto demuestra que el modelo aprendió a generar planes de acción novedosos pero físicamente realizables, superando las limitaciones de los datos de demostración originales.

• Experimento Sintético:

  • Un experimento controlado mostró que la retroalimentación del discriminador por sí sola es suficiente para "reparar" la geometría de los datos latentes, reduciendo la distancia de Mahalanobis de las muestras recombinadas hacia la distribución real, incluso sin entrenamiento de difusión complejo.

5. Significado e Impacto

• Superación de la Imposibilidad Teórica: El trabajo ofrece una solución práctica al problema de la descomposición no supervisada, demostrando que una señal de retroalimentación basada en la plausibilidad de la recombinación (en lugar de etiquetas manuales) es suficiente para aprender representaciones factorizadas útiles.
• Generalización Composicional: Permite a los modelos de difusión no solo generar datos, sino componer nuevos datos a partir de partes aprendidas, lo cual es crucial para la generalización fuera de distribución.
• Avance en Robótica: Proporciona una vía para la exploración autónoma en robots. Al poder recombinar componentes de acción aprendidos de forma no supervisada, los agentes pueden descubrir comportamientos nuevos y efectivos sin necesidad de reentrenamiento costoso o datos masivos adicionales.
• Eficiencia: A diferencia de métodos anteriores que requerían optimización iterativa costosa durante la generación (como los modelos basados en energía), este enfoque utiliza la dinámica de denoising de los modelos de difusión, manteniendo la escalabilidad y estabilidad.

En resumen, el papel presenta un marco robusto que utiliza la adversarialidad para "enseñar" a los modelos de difusión a entender la estructura composicional de los datos, logrando una desconexión superior y una capacidad de generación de nuevos escenarios (tanto visuales como de acción robótica) que son coherentes y realistas.