Resumen Técnico: DiffRGD

Declaración del Problema

Los modelos de difusión se han convertido en fundamentales para la generación de imágenes, sin embargo, controlarlos para tareas específicas (por ejemplo, restauración de imágenes, generación condicional) sin un reentrenamiento o ajuste fino costoso sigue siendo un desafío. Los métodos de guía existentes en tiempo de inferencia (por ejemplo, DPS, FreeDoM) dirigen el proceso de muestreo inyectando gradientes de un objetivo diferenciable directamente en el espacio latente. Sin embargo, estos métodos operan en el espacio ambiente sin respetar la estructura gaussiana subyacente del proceso de difusión.

Esta falta de conciencia distributiva conduce al desplazamiento distributivo (distributional drift), donde las variables latentes se desvían del manifold gaussiano original inducido por el modelo. A medida que los pasos de muestreo se acumulan, este fenómeno de "fuera del manifold" degrada la calidad de la muestra limpia final. Intentos previos para restringir las actualizaciones, como MPGD (que asume linealidad local) o DSG (que proyecta a una hiperesfera de radio fijo), o bien dependen de codificadores perfectos inviables, o bien no logran adherirse completamente a las propiedades gaussianas específicas de la distribución latente, degradándose a veces hacia una distribución uniforme.

Metodología: DiffRGD

Los autores proponen DiffRGD, un marco de guía en tiempo de inferencia de tipo plug-and-play que preserva explícitamente la estructura gaussiana latente. La idea central es que en cada paso de muestreo DDIM $t$, la distribución latente es una gaussiana isotrópica centrada en una media predicha $\mu_t$.

1. Descomposición Polar y Construcción del Manifold

El método aprovecha una propiedad de descomposición polar de las distribuciones gaussianas isotrópicas. En cualquier paso de tiempo $t$, una muestra latente $x_t$ puede descomponerse en un componente radial estadísticamente independiente $r$ y un componente direccional $u$:
$$x_t = \mu_t + \sigma_t r u$$
donde $r$ sigue una distribución chi escalada y $u$ está distribuido uniformemente en una esfera unitaria.

En lugar de optimizar en el espacio ambiente o en una esfera de radio fijo, DiffRGD construye un manifold esférico $\mathcal{S}_{t,r}$ inducido por la distribución latente específica:
$$\mathcal{S}_{t,r} := \{ x \in \mathbb{R}^n \mid \|x - \mu_t\|_2 = \sigma_t r \}$$
Aquí, el radio $\sigma_t r$ se muestrea de la escala natural de la distribución, asegurando que el conjunto factible respete las propiedades gaussianas del proceso de difusión.

2. Descenso de Gradiente de Riemann (RGD)

El problema de guía en cada paso de tiempo se formula como un problema de optimización restringida:
$$x_{t-1} = \arg\min_{x \in \mathcal{S}_{t,r}} \mathcal{L}(\hat{x}_0(x, t-1), y)$$
donde $\mathcal{L}$ es la pérdida de guía (por ejemplo, pérdida de reconstrucción para restauración o alineación de características para generación condicional).

Para resolver esto, DiffRGD emplea Descenso de Gradiente de Riemann (RGD) en el manifold esférico $\mathcal{S}_{t,r}$:

1. Proyección al Espacio Tangente: El gradiente euclidiano de la pérdida se proyecta en el espacio tangente de la esfera en el punto actual, eliminando el componente radial para asegurar que la dirección de actualización sea tangente al manifold.
2. Retracción: Después de dar un paso a lo largo del espacio tangente, una operación de retracción mapea el punto actualizado de nuevo al manifold esférico $\mathcal{S}_{t,r}$, preservando el radio $\sigma_t r$.

Este proceso se repite durante $K$ iteraciones internas (típicamente $K=3$) por cada paso de tiempo de difusión. Los autores demuestran que, bajo supuestos de suavidad $L$, esta optimización converge a un punto estacionario.

Contribuciones Clave

1. Geometría Consciente de la Distribución: Los autores identifican una propiedad de descomposición polar que revela una geometría inducida por la distribución para la guía de difusión. Proponen un marco de Riemann que respeta la geometría del manifold de la distribución, mitigando la degradación de la muestra causada por el desplazamiento fuera del manifold.
2. Análisis Teórico: El artículo proporciona un análisis de la preservación de la distribución y demuestra la convergencia del algoritmo de optimización de Riemann hacia un punto estacionario. En la práctica, la optimización consciente de la geometría propuesta converge más rápido que los métodos basados en ADMM.
3. Rendimiento Plug-and-Play: Se demuestra que DiffRGD es un método plug-and-play que no requiere reentrenamiento ni ajuste fino, logrando resultados de vanguardia (state-of-the-art) en diversas tareas.

Resultados Experimentales

Los autores evaluaron DiffRGD frente a los modelos de referencia (baselines) (DPS, FreeDoM, MPGD, DSG, ADMMDiff) utilizando una GPU NVIDIA GeForce RTX 4090 y el framework HuggingFace Diffusers.

Restauración de Imágenes

Evaluado en FFHQ 256×256 e ImageNet 256×256 para cuatro tareas: Inpainting, Super-Resolución (4×), Desenfoque Gaussiano y Desenfoque de Movimiento.

• Cuantitativo: DiffRGD logró mejoras estadísticamente significativas en PSNR, SSIM, LPIPS y FID sobre todos los modelos de referencia en la mayoría de las tareas. Por ejemplo, en Inpainting, DiffRGD alcanzó un PSNR de 34.04 comparado con 32.38 del segundo mejor método (ADMMDiff).
• Robustez: DiffRGD mantuvo un fuerte rendimiento incluso con menos pasos de muestreo (100 pasos), mientras que métodos como DPS y ADMMDiff mostraron una degradación notable.
• Cualitativo: Los resultados visuales mostraron menos artefactos y una mejor consistencia del fondo en comparación con DPS y MPGD, con una calidad comparable o superior a ADMMDiff pero con menor tiempo computacional.

Generación Condicional

Evaluado en CelebA-HQ 256×256 para guía de Mapa de Segmentación, Boceto (Sketch) y FaceID.

• DiffRGD superó consistentemente a los métodos de inferencia existentes (FreeDoM, DSG, ADMMDiff) en términos de métricas de alineación (mIoU, Sketch-$\ell_2$, FaceID-$\ell_2$) y calidad de generación (FID, KID).
• Los resultados cualitativos indicaron que DiffRGD produjo resultados más naturales y visualmente atractivos con una mejor alineación a las condiciones en comparación con DSG, que a menudo sufría de un control inestable.

Fuerza de Guía

Los experimentos sobre la fuerza de guía ($\eta_t$) mostraron que DiffRLED logra un mejor equilibrio entre la alineación (mIoU) y la calidad de generación (FID) a través de un rango más amplio de intensidades en comparación con DSG, que tuvo un desempeño pobre en todas las configuraciones.

Significancia y Reivindicaciones

El artículo afirma que DiffRGD aborda una brecha crítica en la guía en tiempo de inferencia: el fallo de los métodos existentes para preservar la estructura gaussiana latente paso a paso. Al formular la guía como una optimización restringida en un manifold esférico inducido por la distribución, DiffRGD evita los fenómenos de "fuera del manifold" que degradan la calidad de la muestra.

Los autores sostienen que la guía geométricamente consistente es una alternativa viable y superior a la aplicación directa de gradientes en el espacio ambiente o al uso de manifolds restrictivos que no son conscientes de la distribución. El método ofrece una solución plug-and-play sencilla que mejora las métricas de las tareas mientras preserva la calidad de alta fidelidad inherente a los modelos de difusión preentrenados, sin necesidad de un reentrenamiento costoso.