Improving Diffusion Generalization with Weak-to-Strong Segmented Guidance

Resumen Técnico: Mejora de la Generalización en Modelos de Difusión con Guía Segmentada de Débil a Fuerte (SGG)

1. El Problema

Los modelos de difusión y de emparejamiento de flujo (flow matching) son el estándar actual para la síntesis de imágenes, generando resultados realistas mediante un proceso iterativo de refinamiento. Sin embargo, enfrentan un desafío fundamental conocido como sesgo de exposición (exposure bias): existe una desalineación entre el objetivo de entrenamiento (que es libre de simulación y local) y el proceso de muestreo iterativo global.

Esta discrepancia provoca la acumulación de errores de gradiente a lo largo de la trayectoria de muestreo, lo que lleva a:

• Resultados insatisfactorios y falta de generalización.
• Muestras que se desvían de la distribución real (out-of-distribution).
• Fallos en tareas condicionales complejas (como texto-a-imagen).

Para mitigar esto, se utilizan técnicas de guía (guidance) en la inferencia, como la Guía Libre de Clasificador (CFG) y la AutoGuía (AG). Aunque son efectivas, sus regímenes operativos efectivos no están bien definidos:

• CFG (Guía Dependiente de la Condición): Tiende a buscar modos (mode-seeking), lo que mejora la adherencia al prompt pero reduce la diversidad y puede colapsar la estructura intra-clase.
• AG (Guía Agnóstica a la Condición): Utiliza un modelo "inferior" alineado con la condición. Mejora la diversidad y los detalles intra-clase, pero puede fallar en corregir errores graves o desviaciones de la variedad de datos en modelos no bien ajustados.

La ambigüedad actual radica en seleccionar la técnica adecuada sin un marco teórico claro sobre cuándo y por qué falla cada una.

2. Metodología Propuesta

Los autores proponen un marco unificado basado en el principio "Débil a Fuerte" (Weak-to-Strong, W2S) y una nueva técnica híbrida llamada Guía Segmentada (SGG).

A. Principio Débil a Fuerte y Análisis de Regímenes
El trabajo categoriza los métodos de guía en dos clases basadas en cómo se construye la señal "débil" ($\tilde{v}$):

1. Guía Dependiente de la Condición (CDG): Ejemplificada por CFG. La señal débil se obtiene eliminando la condición ($\tilde{c} = \emptyset$). Es efectiva para separar clases y encontrar la variedad correcta (manifold) en niveles de ruido alto.
2. Guía Agnóstica a la Condición (CAG): Ejemplificada por AG. La señal débil se obtiene usando un modelo inferior pero manteniendo la condición ($\tilde{c} = c$). Es efectiva para refinar detalles intra-clase en niveles de ruido bajo.

Mediante experimentos sintéticos y análisis en ImageNet, los autores demuestran que:

• CDG es superior en las etapas tempranas (ruido alto) para corregir la dirección semántica y evitar desviaciones de clase.
• CAG es superior en las etapas tardías (ruido bajo) para refinar texturas y detalles sin colapsar la diversidad.

B. Guía Segmentada (SGG)
SGG es una implementación híbrida que combina lo mejor de ambos mundos dividiendo el proceso de muestreo en el tiempo ($t$):

• Fase de Alto Ruido ($t > \tau$): Se utiliza CDG (CFG) para asegurar que la generación se dirija hacia la variedad condicional correcta.
• Fase de Bajo Ruido ($t \le \tau$): Se cambia a CAG (usando un modelo inferior o perturbado) para refinar los detalles dentro de esa variedad sin perder la adherencia al prompt.

La fórmula de la velocidad guiada $v_w$ es:
$$v_w(x_t, t, c) = v(x_t, t, c) + (w - 1) \cdot g(x_t, t, c)$$
Donde el vector de guía $g$ cambia dinámicamente según el umbral de tiempo $\tau$.

C. Integración en el Entrenamiento
Un avance clave es migrar el principio W2S y SGG directamente al objetivo de entrenamiento, no solo a la inferencia.

• Se modifica la función de pérdida para incluir un término de guía derivado de la diferencia entre la señal fuerte y la débil.
• Se proponen métodos eficientes para construir la señal débil durante el entrenamiento, como Branching (BR), donde una rama auxiliar de una capa intermedia del transformador actúa como el modelo débil. Esto evita la necesidad de entrenar un modelo separado (como en AG) y reduce el costo computacional.
• El objetivo de entrenamiento se vuelve: $L = ||v_\theta - (u + w \cdot \text{sg}[g])||^2$, donde $u$ es el objetivo óptimo y $g$ es la señal de guía.

3. Contribuciones Clave

1. Análisis Teórico y Empírico: Se categorizan y analizan los regímenes operativos de la CDG y CAG bajo la perspectiva W2S, identificando que la granularidad de la condición y la capacidad de ajuste del modelo determinan qué método falla o tiene éxito.
2. SGG (Guía Segmentada): Se introduce una técnica simple pero efectiva que alterna entre CDG y CAG durante el muestreo, logrando un equilibrio óptimo entre adherencia al prompt y calidad estética/diversidad.
3. Migración al Entrenamiento: Se demuestra que el principio W2S puede integrarse en el entrenamiento para mejorar la capacidad de generalización de modelos de difusión sin guía, reduciendo la dependencia de costos adicionales de guía durante la inferencia.
4. Validación Exhaustiva: Se prueban los métodos en modelos de última generación (SD3, SD3.5) y arquitecturas Transformer (SiT), tanto en escenarios condicionales como incondicionales.

4. Resultados

• Inferencia (SD3 y SD3.5):
  • SGG supera a las variantes de guía existentes (CFG, SLG, AG, etc.) en métricas de adherencia al prompt (HPSv2.1) y calidad estética (Aesthetic Score).
  • En SD3.5, SGG logra un HPSv2.1 de 29.736 y una puntuación estética de 5.717, superando a CFG (que tiene mejor adherencia pero peor estética) y a SLG (que tiene mejor estética pero peor adherencia).
• Entrenamiento (SiT en ImageNet):
  • La integración de W2S/SGG en el entrenamiento acelera la convergencia y mejora la generalización del modelo base sin guía.
  • En generación condicional, el modelo entrenado con SGG alcanza un FID de 4.58 (sin guía en inferencia), superando al modelo base con guía CFG (FID 6.02).
  • La variante BR (Branching) ofrece mejoras significativas con un sobrecosto computacional de entrenamiento mínimo (~2%).
• Análisis de Ablación: Se confirma que el punto de segmentación $\tau$ es crítico; un punto intermedio (ni muy temprano ni muy tarde) ofrece el mejor equilibrio (Pareto frontier).

5. Significado e Impacto

Este trabajo es significativo porque:

• Resuelve la dicotomía Guía vs. Diversidad: Proporciona una solución práctica al compromiso clásico entre la fidelidad al prompt (CFG) y la diversidad/calidad visual (AG), demostrando que no son mutuamente excluyentes si se aplican en los momentos correctos del proceso de difusión.
• Unifica Entrenamiento e Inferencia: Al migrar la lógica de guía al entrenamiento, se reduce la necesidad de cálculos extra durante la inferencia, haciendo que los modelos sean más eficientes y robustos "out-of-the-box".
• Marco Generalizable: El principio de "Débil a Fuerte" y la segmentación temporal ofrecen una nueva dirección para el diseño de arquitecturas y objetivos de entrenamiento en modelos generativos, aplicable más allá de la generación de imágenes (ej. video, audio).

En resumen, SGG representa un avance hacia modelos de difusión más robustos y generalizables, eliminando la ambigüedad en la selección de técnicas de guía y mejorando tanto la calidad de salida como la eficiencia computacional.