Denoising, Fast and Slow: Difficulty-Aware Adaptive Sampling for Image Generation

1. El Problema: Homogeneidad vs. Heterogeneidad en la Generación

Los modelos modernos de generación de imágenes basados en difusión y flujo (como DiT y SiT) asignan recursos computacionales de manera uniforme en todo el espacio de la imagen. En cada paso de denoising, todas las "parches" (regiones espaciales) se actualizan con el mismo nivel de ruido global y el mismo número de evaluaciones de funciones.

Limitaciones de este enfoque:

• Ignora la heterogeneidad: Las imágenes naturales no son uniformes. Algunas regiones (fondos, áreas de baja frecuencia) son fáciles de denoising, mientras que otras (bordes de objetos, texto pequeño, oclusiones) requieren mucha más refinación y contexto.
• Desperdicio de recursos: Tratar todas las regiones por igual fuerza al modelo a gastar cómputo innecesario en áreas fáciles, mientras que las áreas difíciles reciben recursos insuficientes.
• Falta de contexto dinámico: Las regiones difíciles se benefician de tener información de regiones ya resueltas (contexto), pero los métodos estándar no permiten que las partes "fáciles" de la imagen avancen más rápido para guiar a las "difíciles".

2. Metodología: Patch Forcing (PF)

Los autores proponen Patch Forcing (PF), un marco flexible que permite que diferentes regiones de una imagen sigan trayectorias de denoising distintas (diferentes escalas de ruido y timesteps) para intercambiar información contextual.

Componentes Clave:

A. Muestreo de Timesteps a Nivel de Parche (Training)
El desafío principal es cerrar la brecha entre entrenamiento e inferencia.

• El problema del muestreo ingenuo: Si se muestrean timesteps independientes para cada parche, la distribución tiende a concentrarse alrededor de $t \approx 0.5$, lo que significa que durante el entrenamiento el modelo ve casi siempre señales limpias (contexto parcial), algo que no ocurre en la inferencia (que empieza desde ruido puro).
• La solución (LTG Sampler): Los autores introducen un muestreador Gaussiano Truncado Log-Normal (LTG).
  1. Se muestrea un $t_{max}$ (el tiempo máximo de información disponible en la muestra) desde una distribución Log-Normal.
  2. Luego, se muestrean los timesteps individuales de cada parche ($t_i$) condicionalmente a $t_{max}$, restringiéndolos a la mitad inferior de una distribución gaussiana centrada en $t_{max}$.
  3. Objetivo: Controlar la información máxima disponible durante el entrenamiento para que coincida mejor con la distribución de inferencia (ruido puro), evitando la fuga de contexto.

B. Cabeza de Dificultad (Uncertainty Head)

• Se añade una cabeza ligera al modelo que predice la incertidumbre (o dificultad) por parche.
• Esta incertidumbre se modela como la desviación estándar predicha de la velocidad del modelo.
• Hallazgo clave: Existe una fuerte correlación entre la incertidumbre predicha y el error de validación. Las regiones con alta incertidumbre son las que realmente necesitan más refinamiento y contexto.

C. Estrategias de Muestreo Adaptativo (Inferencia)
Utilizando el mapa de dificultad, se proponen dos estrategias para asignar cómputo dinámicamente:

1. Dual-Loop (Bucle Dual):
   • Bucle exterior: Avanza los parches "confiables" (baja incertidumbre) con pasos de tiempo más grandes.
   • Bucle interior: Condiciona los parches "inciertos" (alta incertidumbre) en el estado actualizado de los parches confiables, realizando pasos más pequeños para refinarlos.
   • Esto permite que las regiones fáciles proporcionen contexto a las difíciles en tiempo real.
2. Look-Ahead (Mirada al Futuro):
   • Propaga las representaciones de los parches confiables hacia un timestep futuro ($t + \delta$) para crear un "contexto futuro" limpio.
   • Este contexto avanzado se utiliza para guiar el denoising de las regiones inciertas en el timestep actual.

3. Contribuciones Clave

1. Patch Forcing (PF): Un nuevo marco que introduce timesteps heterogéneos a nivel de parche para la síntesis de imágenes, superando las limitaciones de los enfoques globales.
2. Muestreador LTG: Una estrategia de entrenamiento que controla la información máxima ($t_{max}$) para eliminar la brecha entre entrenamiento e inferencia, superando a métodos anteriores como Spatial Reasoning Models (SRM).
3. Muestreo Consciente de la Dificultad: Demostración de que predecir la incertidumbre por parche permite asignar recursos computacionales donde más se necesitan, mejorando la calidad sin aumentar el presupuesto de cómputo total.
4. Ortogonalidad: El método es compatible y se combina eficazmente con técnicas existentes de alineación de representaciones (como REPA) y guía libre de clasificadores (CFG).

4. Resultados Experimentales

Los autores evaluaron su método en ImageNet (256x256) y en síntesis de texto a imagen (T2I).

• ImageNet (Clase-condicional):
  • El modelo PFT (Patch Forcing Transformer) supera consistentemente a las bases SiT (Scalable Interpolant Transformers) de tamaños B, L y XL.
  • Con el muestreador Look-Ahead, el modelo PFT-XL/2 logra un FID de 9.8 (comparado con 17.2 de SiT-XL/2 base y 7.9 de SiT-XL/2 + REPA).
  • La combinación de PFT + Look-Ahead + REPA alcanza un FID de 6.7, estableciendo un nuevo estado del arte competitivo.
• Síntesis Texto-a-Imagen (T2I):
  • Escalaron el método a un modelo de 1.2B parámetros entrenado en COYO.
  • En benchmarks como T2I-CompBench++ y GenEval, el modelo PFT con muestreo adaptativo supera o iguala a modelos de vanguardia como SDXL y DALL-E 3 en métricas de alineación y composición.
  • Calidad de Texto: Un hallazgo notable es que PF mejora significativamente la renderización de texto en comparación con modelos entrenados con Flow Matching estándar, logrando mayor precisión en la ortografía y legibilidad.

5. Significado e Impacto

• Cambio de Paradigma: El trabajo desafía la suposición de que el denoising debe ser uniforme. Propone que la adaptatividad espacial es crucial para la eficiencia y la calidad.
• Eficiencia Computacional: Permite obtener imágenes de mayor calidad con el mismo presupuesto de pasos (NFE) o lograr la misma calidad con menos pasos, al enfocarse en las regiones problemáticas.
• Fundamento para Futuros Modelos: Sugiere que los horarios de denoising a nivel de parche son una base prometedora para una nueva clase de muestreadores adaptativos.
• Aplicabilidad General: El método no solo mejora la generación de imágenes estáticas, sino que también tiene beneficios tangibles en la generación de texto y la comprensión de la estructura semántica de la imagen.

En resumen, Patch Forcing demuestra que tratar las partes "fáciles" y "difíciles" de una imagen de manera diferente, permitiendo que las primeras guíen a las segundas, es una estrategia superior para la generación de imágenes de alta fidelidad.