Test-Time Compositional Generalization in Diffusion Models via Concept Discovery

Resumen Técnico: Generalización Composicional en Tiempo de Prueba en Modelos de Difusión mediante Descubrimiento de Conceptos

Enunciado del Problema

La generalización composicional —la capacidad de recombinar partes familiares en configuraciones novedosas— es una característica distintiva de la inteligencia humana, pero sigue siendo un desafío significativo para los modelos de aprendizaje profundo. Si bien las redes neuronales suelen interpolar bien dentro de las distribuciones de entrenamiento, tienen dificultades cuando los primitivos familiares aparecen en combinaciones desconocidas. Los enfoques existentes para la generación composicional en modelos de difusión dependen típicamente de una biblioteca predefinida de conceptos, señales de condicionamiento o etiquetas externas para guiar la generación de composiciones novedosas. Esto crea una brecha entre el descubrimiento de conceptos y la generalización composicional: los métodos anteriores asumen que los primitivos relevantes ya son conocidos y están disponibles, sin abordar escenarios donde un modelo debe descubrir conceptos específicos de la consulta desde cero en tiempo de prueba, sin un vocabulario fijo ni supervisión externa.

Metodología

El artículo propone un marco que reutiliza un modelo de difusión preentrenado como una jerarquía implícita de modos de densidad indexados por tiempo para descubrir y componer conceptos en tiempo de prueba. El método opera en tres etapas principales:

1. Descubrimiento de Conceptos mediante Ascenso de Modo:
   En lugar de utilizar el modelo de difusión únicamente para el muestreo inverso, los autores aprovechan la función de puntuación aprendida indexada por tiempo $s_\theta(x_t, t) \approx \nabla_{x_t} \log p_t(x_t)$ para localizar modos locales de densidad de las marginales ruidosas $p_t(x_t)$. Partiendo de una única consulta fuera de distribución (OOD) $x_q$, el método realiza un ascenso de gradiente sobre la puntuación incondicional en múltiples pasos de tiempo de ruido $t$ para recuperar modos locales.

   • Base Teórica: Aprovechando la fórmula de Tweedie, los autores demuestran que estos modos corresponden a "prototipos de concepto" libres de ruido en el espacio de datos limpio ($x_0$), actuando efectivamente como centroides discretos en diversos niveles de abstracción dependiendo del nivel de ruido.
   • Representación: Cada modo descubierto se mapea de vuelta al espacio $x_0$ y se aproxima como un experto gaussiano local $q_j(x_0) = \mathcal{N}(x_0; m_j, \Sigma_j)$, donde la media se deriva del modo y la covarianza se estima mediante sondas de Hutchinson.

2. Composición Producto de Expertos (PoE):
   Los expertos gaussianos descubiertos se combinan en un único modelo maestro analítico utilizando una formulación Producto de Expertos (PoE).

   • Selección: Un algoritmo de optimización submodular codiciosa selecciona el subconjunto de prototipos más relevante para la consulta, maximizando un objetivo de cobertura de verosimilitud por dimensión.
   • Composición: Los expertos seleccionados se combinan en una distribución PoE ponderada $q_T$. Dado que los expertos son gaussianos, la composición resultante también es gaussiana con una función de puntuación de forma cerrada. Esto permite derivar una predicción de ruido analítica $\epsilon_{q_T}(x_t, t)$ sin requerir un token condicional entrenado para la composición novedosa.

3. Muestreo Composicional y Destilación:

   • Muestreo Analítico: La puntuación analítica del maestro PoE se integra en un marco de guía libre de clasificador (CFG). El modelo contrasta la predicción PoE frente a la predicción incondicional para guiar el muestreo hacia la composición descubierta.
   • Destilación LoRA: Para mejorar la generación dentro de la variedad y crear una condición reutilizable, el método destila opcionalmente al maestro PoE en el modelo base. Esto implica entrenar una nueva incrustación de clase y un adaptador de Adaptación de Bajo Rango (LoRA) sobre un conjunto de muestras generadas por el maestro PoE. Este paso absorbe la composición descubierta de nuevo en la variedad de puntuación del modelo.

Contribuciones Clave

1. Puente entre Descubrimiento y Composición: El trabajo une el descubrimiento de conceptos y la generalización composicional al tratar un modelo de difusión preentrenado como una jerarquía implícita de modos de densidad. Permite la recuperación de prototipos de conceptos discretos en tiempo de prueba sin un conjunto predefinido de primitivos.
2. Mecanismo de Composición en Tiempo de Prueba: Los autores introducen un método que selecciona codiciosamente modos relevantes utilizando un objetivo de verosimilitud submodular y los combina en un modelo maestro PoE con puntuación de forma cerrada, permitiendo el muestreo directo sin reentrenar el modelo base en la composición específica.
3. Validación Empírica: El marco se evalúa en dos puntos de referencia, ColorMNIST y CelebA, demostrando que los prototipos descubiertos y su composición PoE codifican información más rica sobre composiciones no vistas que el muestreo solo con la consulta o la recuperación de la clase entrenada más cercana.

Resultados

Los experimentos en ColorMNIST (identidad del dígito, color, fondo) y CelebA (atributos faciales) muestran que el método propuesto supera a varias líneas base:

• Líneas Base: El método se compara con "Solo consulta" (muestreo alrededor de la consulta con varianza fija), "Top-1/Top-3" (recuperación de las clases entrenadas más cercanas) y CFG estándar.
• Rendimiento: Tanto el muestreador PoE analítico como la variante destilada con LoRA logran métricas superiores de FID, similitud CLIP y Precisión/Recall en comparación con las líneas base.
  • En ColorMNIST, la variante LoRA logra una puntuación F1 de 87.6% (Fidelidad) y 76.3% (Generalización), superando significativamente a la línea base Top-1 (~6.5% F1).
  • En CelebA, el método PoE logra el mejor FID de Generalización (140.3) y puntuaciones CLIP, mientras que LoRA mejora aún más el equilibrio entre cobertura y fidelidad, logrando el mayor Recall y F1 de Generalización.
• Análisis Cualitativo: Los resultados visuales indican que, mientras las líneas base a menudo fallan en preservar la combinación específica de atributos (por ejemplo, cambiando la identidad del dígito o derivando hacia diferentes combinaciones de atributos), el método propuesto reconstruye con éxito la composición OOD y se generaliza a otras instancias de la misma clase.

Significado y Afirmaciones

El artículo afirma que la geometría de puntuación indexada por tiempo de los modelos de difusión contiene conceptos de modo de densidad reutilizables que apoyan la generación composicional en tiempo de prueba sin una biblioteca de conceptos predefinida.

• Insight Teórico: Los resultados sugieren que los modelos de difusión codifican implícitamente una jerarquía de conceptos donde los modos de las marginales ruidosas sirven como prototipos en diferentes niveles de abstracción.
• Implicación Práctica: El método permite la generación de configuraciones novedosas a partir de una única consulta OOD, ofreciendo una vía para la generalización composicional que no depende de supervisión externa ni vocabularios preespecificados.
• Limitaciones: Los autores notan modestamente que el protocolo asume que los factores primitivos son conocidos (configuración composicional estándar) y que el método depende de los sesgos inductivos del backbone preentrenado. Si bien puede interpolar entre primitivos conocidos, no puede fabricar primitivos genuinamente nuevos que extrapolen más allá del soporte de entrenamiento (por ejemplo, un fondo blanco en un modelo entrenado solo con fondos oscuros).

El trabajo proporciona tanto evidencia teórica como empírica de que los modelos de difusión pueden reutilizarse para el descubrimiento y la composición de conceptos, extendiéndose potencialmente a aplicaciones más amplias en modelado de lenguaje y control robótico.