A Noise Constrained Diffusion (NC-Diffusion) Framework for High Fidelity Image Compression

Resumen Técnico: NC-Diffusion para Compresión de Imágenes de Alta Fidelidad

1. Planteamiento del Problema

La compresión de imágenes aprendida (basada en redes neuronales) ha superado a los métodos tradicionales, pero a menudo produce resultados borrosos debido a la función de pérdida basada en el Error Cuadrático Medio (MSE) y al proceso de cuantización. Para mejorar la calidad perceptual, se han explorado modelos generativos, específicamente los Modelos de Difusión Probabilística (DDPM).

Sin embargo, los métodos existentes de compresión basados en difusión presentan dos problemas críticos:

1. Desajuste de Ruido (Noise Mismatch): Los modelos de difusión tradicionales inician el proceso de generación desde ruido gaussiano aleatorio puro. En la compresión de imágenes, el "ruido" introducido es ruido de cuantización (determinista y dependiente de la imagen), no ruido gaussiano aleatorio. Aplicar difusión con ruido gaussiano a un problema de compresión introduce incertidumbre innecesaria y artefactos, desviando la reconstrucción de la imagen original.
2. Ineficiencia en la Inferencia: Los métodos actuales requieren múltiples pasos de inferencia partiendo de ruido aleatorio para propagar la señal hasta la imagen deseada, lo que es computacionalmente costoso y lento.

2. Metodología Propuesta: NC-Diffusion

Los autores proponen un marco de Difusión Constrained por Ruido (NC-Diffusion) que alinea el proceso de difusión con el proceso de compresión real.

• Reformulación del Proceso de Difusión:

  • En lugar de tratar el ruido como una variable aleatoria gaussiana, el marco NC-Diffusion modela el ruido de cuantización (la diferencia entre la imagen original $I_0$ y el resultado inicial de compresión $I_T$) como el ruido añadido en el proceso forward de la difusión.
  • Proceso Forward: Se construye un proceso que difunde la imagen original hacia el resultado de compresión inicial (ruidoso) utilizando el ruido de cuantización real ($\epsilon_n$) en lugar de ruido aleatorio.
  • Proceso Reverse (Inferencia): La inferencia comienza directamente desde la imagen comprimida inicial ($I_T$) y se denoisa paso a paso hacia la imagen original. Esto elimina la necesidad de generar ruido aleatorio, asegurando que la reconstrucción sea fiel a la imagen original y eliminando solo el ruido de cuantización.

• Arquitectura y Módulos Clave:

  • Base U-Net: Utiliza una arquitectura U-Net estándar como red de predicción de ruido.
  • Filtrado Adaptativo en el Dominio de la Frecuencia (AFF): Se introduce un módulo para mejorar las conexiones de salto (skip connections) en la U-Net. Transforma las características al dominio de Fourier, aplica una máscara aprendible para realzar selectivamente las altas frecuencias (detalles finos) y luego las invierte. Esto ayuda a recuperar los detalles de textura perdidos durante la compresión.
  • Pérdida de Preservación de Detalles de Alta Frecuencia ($L_{high}$): Se añade una función de pérdida basada en la transformada wavelet para supervisar y minimizar el error específicamente en las sub-bandas de alta frecuencia.
  • Mejora Guiada por Muestras (Zero-Shot): Durante la inferencia, se utiliza una pérdida perceptual basada en CLIP para guiar el proceso de muestreo. Esto ajusta el gradiente para mantener la fidelidad semántica y perceptual sin necesidad de entrenar un clasificador adicional.

3. Contribuciones Principales

1. Marco NC-Diffusion: La primera propuesta que resuelve el problema del desajuste de distribución de ruido al equiparar el ruido de cuantización con el ruido añadido en la difusión. Esto permite una reconstrucción determinista y de alta fidelidad sin introducir ruido aleatorio extra.
2. Eficiencia de Inferencia: Al iniciar la inferencia desde la imagen comprimida y no desde ruido gaussiano puro, el método logra resultados de alta calidad en menos pasos (incluso en un solo paso), mejorando drásticamente la velocidad de decodificación.
3. Módulo de Filtrado de Frecuencia (AFF): Un módulo "plug-and-play" compatible con cualquier arquitectura U-Net que mejora la recuperación de detalles de alta frecuencia.
4. Rendimiento Equilibrado: Logra un compromiso óptimo entre métricas de distorsión (PSNR) y métricas perceptuales (LPIPS, FID), evitando el problema común de los métodos generativos que sacrifican la fidelidad estructural por una apariencia "bonita" pero incorrecta.

4. Resultados Experimentales

Las pruebas se realizaron en los conjuntos de datos Kodak y CLIC2020, comparando con métodos state-of-the-art como HiFiC, MS-ILLM, DiffEIC y CDC.

• Calidad de Reconstrucción: NC-Diffusion supera a los métodos existentes tanto en métricas de distorsión (PSNR) como en métricas perceptuales (FID, LPIPS).
  • En Kodak, logra ahorros de tasa de bits (BD-Rate) del 12.04% solo con el módulo NC-Diffusion, y hasta un 28.01% con todos los componentes (AFF + Pérdida de alta frecuencia).
  • Mantiene una fidelidad estructural superior: a diferencia de HiFiC o MS-ILLM, que a veces alteran texturas o eliminan detalles (ej. letras en imágenes), NC-Diffusion preserva la textura original fielmente.
• Velocidad de Inferencia:
  • La velocidad de decodificación es significativamente mayor que otros métodos basados en difusión (ej. 0.102s por imagen frente a 1.837s de DiffEIC). Esto se debe a que no necesita propagar ruido desde cero.
• Compromiso Distorsión-Percepción: El método demuestra una mejora significativa en la calidad perceptual con una pérdida mínima en la calidad objetiva (PSNR), superando la curva de compromiso tradicional.

5. Significado e Impacto

Este trabajo es significativo porque cierra la brecha teórica entre los modelos de difusión generativa y la compresión de imágenes práctica. Al reconocer y modelar correctamente la naturaleza del ruido de cuantización, el método elimina la aleatoriedad indeseada que ha limitado la aplicación de la difusión en compresión de alta fidelidad.

El marco NC-Diffusion demuestra que es posible utilizar la capacidad de generación de los modelos de difusión para refinar imágenes comprimidas (eliminación de ruido) en lugar de generar imágenes desde la nada. Esto abre la puerta a sistemas de compresión de próxima generación que ofrecen una calidad visual superior y una fidelidad estructural estricta, esenciales para aplicaciones donde la precisión de la imagen es crítica (como medicina, vigilancia o archivado), manteniendo al mismo tiempo tasas de bits bajas.