TOC-SR: Task-Optimal Compact diffusion for Image Super Resolution

Resumen Técnico de TOC-SR: Difusión Compacta Óptima para Tareas de Superresolución de Imágenes

Planteamiento del Problema
Aunque los modelos de difusión han demostrado un rendimiento superior en tareas de restauración de imágenes como la Superresolución de Imagen Única (SR) al modelar estadísticas de imágenes naturales y proporcionar priores generativos sólidos, enfrentan barreras significativas para su despliegue práctico. Las pipelines de difusión estándar requieren grandes redes base y procedimientos de muestreo iterativo (a menudo decenas de pasos de eliminación de ruido), lo que resulta en alta latencia, uso de memoria y costo computacional. Los intentos existentes de crear modelos de SR compactos de un solo paso a menudo sufren de optimización inestable y calidad degradada porque el modelo debe aprender simultáneamente el prior generativo y el mapeo de restauración dentro de una capacidad limitada.

Metodología
Los autores proponen TOC-SR, un marco diseñado para construir modelos de SR de un solo paso eficientes mediante la descubierta primero de una base de difusión compacta que preserve la fidelidad generativa. La metodología procede en tres etapas distintas:

1. Expansión de Capacidad Latente (de 4 a 16 Canales):
   Reconociendo que la representación latente estándar de cuatro canales en Stable Diffusion crea un cuello de botella de información para los detalles de alta frecuencia requeridos en SR, los autores expanden la dimensionalidad latente a dieciséis canales manteniendo el mismo factor de compresión espacial. Dado que no existe un modelo estándar de Stable Diffusion de 16 canales, construyen una nueva base destilando la capacidad representacional de un VAE de alta capacidad (de la familia de modelos FLUX) en un VAE de Stable Diffusion modificado. Esto implica entrenar un codificador estudiante para igualar la distribución latente de 16 canales del profesor mientras se mantiene la fidelidad de reconstrucción. Posteriormente, la U-Net de difusión se modifica para aceptar y producir tensores de 16 canales, seguida de un ajuste fino conjunto para adaptar el prior de difusión a este espacio expandido.

2. Descubrimiento de Base Compacta mediante Optimización Bayesiana $\epsilon$-Constrained Basada en Sustitutos:
   Para reducir el costo computacional del modelo expandido de 16 canales, los autores emplean una estrategia de búsqueda en lugar del entrenamiento directo de modelos compactos.

   • Construcción de Sustitutos: La U-Net del profesor se descompone en etapas secuenciales. Para cada etapa, se construye una biblioteca de bloques sustitutos eficientes en parámetros mediante poda estructurada y reducción de módulos (por ejemplo, eliminando módulos de atención o simplificando bloques residuales) mientras se preserva estrictamente la dimensionalidad de los tensores de entrada y salida.
   • Destilación Generativa a Nivel de Características (FwKD): Cada bloque sustituto candidato se calibra de forma independiente utilizando una pérdida de destilación a nivel de características contra las características del profesor en un subconjunto de calibración. Esto evita la necesidad de un entrenamiento completo de difusión durante la fase de búsqueda.
   • Optimización Bayesiana $\epsilon$-Constrained (BO): La búsqueda de arquitectura se formula como un problema de optimización con restricciones. El objetivo es minimizar la cantidad de parámetros ($f_{param}$) sujeto a una restricción sobre la fidelidad generativa ($f_{acc} \leq \epsilon$), donde $f_{acc}$ mide la discrepancia en las predicciones de ruido entre el candidato y el profesor. Un modelo sustituto de Proceso Gaussiano guía el proceso de BO para explorar eficientemente el espacio de arquitecturas a través de una cuadrícula de valores de $\epsilon$, identificando arquitecturas que ofrecen el mejor equilibrio entre compresión y fidelidad.

3. Especialización de Tarea y Destilación de Un Paso:
   Una vez identificada la base compacta óptima, esta se adapta a la tarea de SR mediante entrenamiento de difusión condicional, donde el modelo aprende a eliminar ruido de una representación latente condicionada a una entrada de baja resolución (LR). Finalmente, este modelo de difusión SR iterativo se destila en un generador de un solo paso. La red estudiante se entrena para predecir directamente el latente limpio de alta resolución final desde un latente ruidoso y la condición LR en una sola pasada hacia adelante, eliminando la necesidad de eliminación de ruido iterativa.

Contribuciones Clave

• Expansión Latente para SR: El artículo demuestra que expandir el espacio latente de 4 a 16 canales alivia el cuello de botella para modelar detalles de alta frecuencia, proporcionando un punto de partida más sólido para la compresión que las arquitecturas estándar.
• Búsqueda de Arquitectura Basada en Sustitutos: Los autores introducen un marco novedoso que combina la destilación generativa a nivel de características con la Optimización Bayesiana $\epsilon$-constrained para descubrir bases de difusión compactas. Este enfoque reduce sistemáticamente los parámetros mientras asegura que las arquitecturas candidatas permanezcan fieles a la dinámica generativa del profesor sin requerir entrenamiento específico de la tarea durante la búsqueda.
• SR de Un Paso Eficiente: El marco resultante, TOC-SR, integra exitosamente un prior de difusión compacto con un mecanismo de inferencia de un solo paso, logrando una reconstrucción de alta calidad sin muestreo iterativo.

Resultados Experimentales
Se realizaron experimentos en benchmarks estándar (DIV2K, RealSR, DRealSR) comparando TOC-SR con métodos de SR basados en difusión de última generación (por ejemplo, StableSR, DiffBIR, SeeSR) y modelos compactos/de un solo paso existentes (por ejemplo, SinSR, OSEDiff, PocketSR).

• Eficiencia: En comparación con el modelo de difusión expandido de 16 canales, la base TOC-SR descubierta logra una reducción de 6.6× en parámetros y una reducción de 2.8× en GMACs.
• Rendimiento: En el conjunto de validación DIV2K, TOC-SR alcanza un PSNR de 24.26 dB y la mejor puntuación LPIPS (0.242) entre los métodos listados, operando en un solo paso.
• Generalización al Mundo Real: En conjuntos de datos del mundo real (RealSR y DRealSR), TOC-SR supera a otras bases compactas en métricas perceptuales (MUSIQ) y mantiene un PSNR competitivo, demostrando robustez frente a degradaciones del mundo real.
• Análisis Cualitativo: Las comparaciones visuales muestran que TOC-SR preserva mejor las texturas finas y las estructuras delgadas que otros modelos compactos, evitando artefactos comunes como el suavizado excesivo y las microtexturas antinaturales.

Significado y Afirmaciones
El artículo afirma que TOC-SR proporciona una vía estable y efectiva para construir sistemas de restauración eficientes y de alta calidad. Al desacoplar el descubrimiento de un prior generativo compacto de la adaptación específica de la tarea, el marco evita la inestabilidad de optimización a menudo observada en el entrenamiento directo de modelos de SR compactos. Los autores afirman que su enfoque equilibra exitosamente la fuerza generativa de los modelos de difusión con la eficiencia requerida para el despliegue práctico, ofreciendo una reducción significativa en el costo computacional (de decenas de pasos a uno) mientras mantiene una fuerte calidad de reconstrucción. El trabajo sugiere que descubrir priores de difusión compactos mediante optimización con restricciones es una estrategia viable para la restauración de imágenes amigable con dispositivos periféricos (edge).