Resumen Técnico: Alineación Contrastiva Condicional-No Condicional para Modelos de Difusión de Cola Larga

Planteamiento del Problema

El entrenamiento de datos para la síntesis de imágenes condicionales por clase exhibe frecuentemente una distribución de cola larga, caracterizada por una escasez de imágenes para las clases de la "cola" en comparación con las clases de la "cabeza". Este desequilibrio provoca que los modelos de difusión sufran de colapso de modo (mode collapse) y una reducción en la diversidad al generar imágenes para estas clases subrepresentadas. Los métodos existentes, incluyendo los modelos de difusión probabilística de eliminación de ruido (DDPM) estándar y los modelos de difusión basados en puntuación (SBDM), a menudo fallan en generar imágenes de alta fidelidad y diversidad para estas clases subrepresentadas sin comprometer la calidad de la generación de las clases de la cabeza. Aunque el aprendizaje contrastivo ha demostrado ser efectivo para el reconocimiento de cola larga y la generación basada en GAN, su aplicación a los modelos de difusión para la generación con desequilibrio de clases requiere una adaptación novedosa para abordar la mecánica específica del proceso de eliminación de ruido.

Metodía: Alineación Contrastiva Condicional–No Condicional (CCUA)

Los autores proponen CCUA, un marco que comprende dos funciones de pérdida sinérgicas diseñadas para mejorar la diversidad y la fidelidad de las imágenes de las clases de la cola manteniendo la calidad de las clases de la cabeza. El método es compatible tanto con arquitecturas basadas en U-Net como con Transformers de Difusión (DiT/SiT).

1. Pérdida Contrastiva No Supervisada (UCL) para Generación No Condicional

Para abordar el colapso de modo, los autores introducen una pérdida contrastiva no supervisada aplicada a los latentes de la generación no condicional.

• Mecanismo: La red de estimación de ruido se divide en un codificador latente $e(\cdot)$ y un decodificador $d(\cdot)$. El codificador mapea las imágenes ruidosas $x_t$ a latentes de baja dimensión $h$.
• Formulación de la pérdida: La pérdida utiliza una formulación de estilo InfoNCE con solo muestras negativas. Dentro de un mini-lote (mini-batch), cada imagen sirve como ancla, y todas las demás imágenes en el lote sirven como muestras negativas. No se crean pares positivos mediante aumentación de datos; el ancla es tratada como su propio vector positivo.
• Objetivo: Esto crea una "fuerza repulsiva" en el espacio latente, maximizando la distancia entre las imágenes sintéticas independientemente de su clase. Esto obliga al modelo a diversificar los latentes no condicionales, evitando que colapsen hacia un vector constante o se sobreajusten a muestras de entrenamiento limitadas.
• Justificación Teórica: Aplicar directamente la pérdida contrastiva a los latentes condicionales puede conducir a una solución trivial donde el modelo maximiza la diversidad inter-clase (separando clases) mientras ignora la diversidad intra-clase. Al aplicar la pérdida a los latentes no condicionales, se obliga al modelo a diversificar todas las muestras, lo cual es luego destilado en la generación condicional a través de la segunda pérdida.

2. Pérdida de Alineación (AL) para la Alineación Condicional-No Condicional

El segundo componente alinea la generación condicional con la generación no condicional diversificada.

• Mecanismo: La pérdida minimiza la divergencia KL entre la distribución condicional $p_\theta(x_{t-1}|x_t, c)$ y la distribución no condicional $p_\theta(x_{t-1}|x_t)$. En la práctica, esto se simplifica a una pérdida MSE entre el ruido estimado $\epsilon_\theta(x_t, t, c)$ y $\epsilon_\theta(x_t, t, \emptyset)$.
• Ponderación Temporal: Crucialmente, esta alineación se pondera más fuertemente en los pasos de tiempo grandes (los pasos iniciales del proceso inverso).
• Motivación: Los autores observan que durante los pasos iniciales de eliminación de ruido, los componentes de baja frecuencia de las imágenes de diferentes clases (y entre la generación condicional y la no condicional) son altamente similares. Al alinear estos pasos iniciales, el modelo facilita el intercambio de conocimiento (knowledge sharing) de las clases de la cabeza (datos abundantes) hacia las clases de la cola (datos escasos). Esto permite que las clases de la cola se beneficien de la rica diversidad aprendida por las clases de la cabeza durante las etapas iniciales de la generación.

3. Sinergia

El marco combina la pérdida estándar de DDPM con la pérdida propuesta CCUA:
$$\mathcal{L}_{total} = \mathcal{L}_{ddpm} + \alpha \cdot \mathcal{L}_{ucl} + \gamma \cdot \mathcal{L}_{al}$$
La sinergia reside en el hecho de que la UCL diversifica el espacio latente no condicional, y la AL alinea el espacio latente condicional con este espacio no condicional diversificado. Esta diversificación indirecta de los latentes condicionales es empíricamente más efectiva que aplicar directamente la pérdida contrastiva a los latentes condicionales.

Principales Contribuciones

1. Marco CCUA: Un método novedoso para modelos de difusión con desequilibrio de clases que combina el aprendizaje contrastivo no supervisado en latentes no condicionales con una pérdida de alineación condicional-no condicional.
2. Estrategia de Alineación: La propuesta de alinear la generación condicional y no condicional específicamente en los pasos de tiempo grandes (pasos iniciales de eliminación de ruido) para aprovechar la similitud de los componentes de baja frecuencia y permitir la transferencia de conocimiento de la cabeza a la cola.
3. Regularización Contrastiva No Supervisada: El uso de una pérdida contrastiva no supervisada con solo muestras negativas en latentes no condicionales para maximizar la diversidad intra-clase e inter-clase sin requerir etiquetas de clase para los pares contrastivos.
4. Agnosticismo de Arquitectura: Se demuestra que el método es efectivo tanto en arquitecturas basadas en U-Net como en Diffusion Transformers (SiT).

Resultados Experimentales

El método fue evaluado en los conjuntos de datos ImageNet-LT (256×256), TinyImageNet-LT (64×64), Places-LT (64×64) y CIFAR-LT.

• Rendimiento Cuantitativo:
  • En ImageNet-LT, CCUA supera significativamente a los modelos de referencia incluyendo SiT, CBDM, REPA y Dispersive Loss.
  • A los 900k pasos de entrenamiento, CCUA logró un FID de 15.1 (frente a 19.9 de SiT) y un FIDtail de 22.5 (frente a 33.9 de SiT), lo que representa una mejora de ~11.4 puntos en la fidelidad de la clase de la cola.
  • El método también mejoró las métricas de Recall (diversidad) y Precision (fidelidad) en todos los ámbitos.
  • En TinyImageNet-LT, CCUA alcanzó el mejor FID (12.0) y FIDtail (20.8), superando a DDPM, CBDM y OCLT.
• Resultados Cualitativos: Las visualizaciones muestran que CCUA genera imágenes con mayor diversidad y fidelidad para las clases de la cola (por ejemplo, 'red wine', 'espresso', 'window shade') en comparación con el SiT estándar, que a menudo produce imágenes repetitivas o de baja calidad para estas clases.
• Estudios de Ablación:
  • La combinación de ambas pérdidas es necesaria; usar solo una produce resultados subóptimos.
  • Aplicar la pérdida de alineación en pasos de tiempo grandes (ponderación lineal) es superior a la ponderación constante o cuadrática.
  • El uso de representaciones latentes del bloque N/4 del transformer generalmente ofrece el mejor equilibrio entre el FID general y el FID espacial/Recall.

Significado y Reivindicaciones

El artículo afirma que CCUA aborda eficazmente el problema de la generación de cola larga adaptando principios de las GAN de cola larga (específicamente el entrenamiento no condicional en resoluciones más bajas/etapas tempranas) al paradigma de la difusión. Al desacoplar la tarea de diversificación (manejada por el aprendizaje contrastivo no condicional) de la tarea de condicionamiento (manejada por la alineación), el método evita los problemas de la regularización contrastiva condicional directa.

Los autores enfatizan que su enfoque:

• Mejora la generación de clases de la cola sin degradar la calidad de las clases de la cabeza.
• Es fácil de implementar dentro de los flujos de trabajo existentes de DDPM y SBDM.
• Supera a trabajos concurrentes como Dispersive Loss y DiffROP, particularmente en escenarios de desequilibrio de clases extremo (por ejemplo, ImageNet-LT con un factor de desequilibrio de 0.001).
• Proporciona una justificación teórica y empírica de por qué alinear la generación condicional y no condicional en los pasos iniciales del proceso inverso es beneficioso para el intercambio de conocimiento en entornos de cola larga.