Few-Step Diffusion Sampling Through Instance-Aware Discretizations

Resumen Técnico: INDIS (Instance-Specific Discretization)

1. El Problema

Los modelos de difusión y de ajuste de flujo (flow matching) generan datos de alta fidelidad simulando trayectorias definidas por Ecuaciones Diferenciales Ordinarias (ODE) o Estocásticas (SDE). Para acelerar el proceso de muestreo y reducir el número de pasos de inferencia (NFE - Number of Function Evaluations), se utilizan solucionadores numéricos avanzados.

Sin embargo, la estrategia de discretización (la selección de los pasos de tiempo $\tau$) es un componente crítico que a menudo se pasa por alto.

• Limitación actual: La mayoría de los métodos existentes, tanto heurísticos (como Uniform o LogSNR) como los basados en optimización (que buscan un horario de pasos óptimo), imponen un horario de pasos global y compartido para todas las muestras.
• La causa raíz: Este enfoque asume que la complejidad del proceso generativo es uniforme para todos los datos. En realidad, diferentes muestras de ruido inicial ($x_T$) y condiciones ($c$) generan trayectorias de muestreo con dinámicas intrínsecamente distintas. Un horario global es un compromiso óptimo en promedio, pero subóptimo para cada instancia individual, limitando la calidad de la generación en pocos pasos.

2. Metodología: INDIS

Los autores proponen INDIS (Instance-Specific Discretization), un marco que aprende a adaptar la asignación de pasos de tiempo basándose en las características específicas de cada entrada (ruido inicial y condición).

Componentes Clave del Método:

1. Red de Condicionamiento de Prioridad (Prior Conditioning Network):

   • Se introduce una red neuronal ligera, $\phi(\cdot)$, que toma como entrada el ruido inicial $x_T$ y la información condicional $c$ (etiquetas de clase o embeddings de texto).
   • La salida de la red es un conjunto de parámetros de discretización personalizados ($\xi_\phi$) para esa instancia específica.
   • Arquitectura: Utiliza descomposición en valores singulares (SVD) para procesar el ruido $x_T$ y capas lineales/ReLU para integrar las condiciones. Para video, se aplica un pooling en la dimensión temporal para eficiencia.

2. Optimización y Entrenamiento:

   • El objetivo es minimizar la distancia entre la muestra generada por el estudiante (con pocos pasos y discretización aprendida) y la muestra "maestra" (generada con muchos pasos y un solucionador de referencia).
   • Función de Pérdida: Se utiliza una métrica de distancia (como MSE, LPIPS o KL-divergencia) entre el punto final de la trayectoria del estudiante y el del maestro.
   • Enfoque "Soft": A diferencia de métodos anteriores que trataban el ruido como un parámetro fijo, INDIS aprende a mapear directamente el ruido a su mejor horario de pasos.

3. Adaptaciones para Escalabilidad y Sesgo:

   • Guía Condicional: La red $\phi$ acepta explícitamente $c$ para manejar la generación condicional (texto a imagen, clases).
   • Factores de Desplazamiento y Escala: Para mitigar el problema de sesgo de exposición (exposure bias), donde la distribución de estados durante el muestreo difiere de la del entrenamiento, la red predice factores de desplazamiento temporal ($\Delta\tau$) y escala ($\gamma$). Esto permite ajustar dinámicamente la evaluación de la función de ruido $\hat{\epsilon}_\theta$.

4. Eficiencia:

   • El costo de inferencia es mínimo: requiere un solo paso forward de la red ligera $\phi$ antes de iniciar el muestreo principal. En configuraciones de 5 pasos, este sobrecosto es inferior al 3% del tiempo total.

3. Contribuciones Principales

• Identificación de la Suboptimalidad Global: Demostraron mediante experimentos sintéticos que un horario de pasos global no puede ser óptimo para todas las instancias, estableciendo una brecha de rendimiento cuantificable entre estrategias globales y específicas por instancia.
• Marco INDIS: Propusieron un método escalable que generaliza la búsqueda de discretización basada en gradientes al contexto generativo condicional, aprendiendo a asignar pasos de tiempo personalizados.
• Validación Exhaustiva: Validaron el método en diversos dominios:
  • Datos sintéticos (análisis teórico).
  • Espacio de píxeles (CIFAR-10, ImageNet, FFHQ).
  • Espacio latente (Stable Diffusion en LSUN-Bedroom).
  • Modelos de flujo para video (LTX-Video).
  • Modelos de última generación (FLUX.1-dev).

4. Resultados Experimentales

Los resultados muestran mejoras consistentes en la calidad de generación (medida por FID - Fréchet Inception Distance) comparado con heurísticas manuales y métodos de optimización global (como DMN, GITS, LD3).

• Mejoras en Espacio de Píxeles: En CIFAR-10, ImageNet y FFHQ, INDIS supera a los mejores baselines en un 35% (NFE=3), 31% (NFE=5) y 15% (NFE=7) en promedio. La ventaja es más pronunciada en configuraciones de muy pocos pasos.
• Espacio Latente y Video: En LSUN-Bedroom y FLUX.1-dev, se observan mejoras significativas en FID y CLIP scores. En video (LTX-Video), INDIS mejora la consistencia del sujeto y la calidad estética según las métricas de VBench.
• Ablaciones: Se demostró que el condicionamiento a nivel de instancia es el factor más crítico para el rendimiento. Los factores de desplazamiento (shift factors) son cruciales para modelos con esquemas de ruido específicos (como EDM-VE), pero menos impactantes en otros.
• Comparación con Distilación: INDIS logra un rendimiento superior o comparable a métodos de distilación de solucionadores (como EPD o AdaSDE), pero con un costo de entrenamiento mucho menor (no requiere reentrenar el modelo base, solo la pequeña red $\phi$).

5. Significado e Impacto

Este trabajo representa un cambio de paradigma en la aceleración de modelos de difusión:

1. De lo Global a lo Local: Pasa de tratar el muestreo como un problema de optimización global a uno adaptativo por instancia, reconociendo la heterogeneidad de los datos.
2. Eficiencia sin Costo de Entrenamiento Pesado: A diferencia de la distilación de modelos (que requiere reentrenar el modelo base), INDIS es un método "training-free" para el modelo base, requiriendo solo un ajuste ligero de una red pequeña.
3. Generalización: Funciona eficazmente en modelos de difusión clásicos, ajuste de flujo (flow matching), imágenes y video, demostrando ser una solución robusta para la aceleración de inferencia en la próxima generación de modelos generativos.

En conclusión, INDIS demuestra que adaptar dinámicamente la estrategia de integración numérica a la complejidad específica de cada muestra es la clave para lograr una generación de alta calidad en pocos pasos, superando las limitaciones de los horarios de tiempo fijos.