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

Título: Um Framework de Difusão Confinada por Ruído (NC-Diffusion) para Compressão de Imagem de Alta Fidelidade

1. O Problema

A compressão de imagem baseada em aprendizado (learned image compression) utiliza redes neurais para codificar e decodificar imagens, mas frequentemente sofre com o processo de quantização, que introduz ruído e resulta em imagens reconstruídas borradas ou ambíguas.

Para melhorar a qualidade perceptual, modelos de difusão (como DDPMs) foram aplicados à compressão. No entanto, os métodos existentes apresentam duas falhas críticas:

1. Incompatibilidade de Distribuição de Ruído: Os modelos de difusão tradicionais adicionam ruído gaussiano aleatório durante o processo de difusão. Em contraste, o ruído introduzido na compressão de imagem é ruído de quantização, que possui uma distribuição diferente (geralmente uniforme ou dependente do transformador de decodificação) e padrões estruturados (especialmente em bordas).
2. Incerteza e Ineficiência: Os métodos atuais começam a geração a partir de ruído gaussiano puro. Isso introduz incerteza inerente, levando a reconstruções que se desviam da imagem original (perda de fidelidade), o que é catastrófico para tarefas de compressão onde a consistência é vital. Além disso, o processo de inferência a partir do ruído puro exige muitas etapas, tornando-o lento.

2. Metodologia Proposta: NC-Diffusion

Os autores propõem o NC-Diffusion, um framework que alinha o processo de difusão com a realidade da compressão de imagem, tratando o ruído de quantização como o ruído a ser removido, em vez de adicionar ruído aleatório.

A. Difusão Confinada por Ruído (Noise Constrained Diffusion)

• Conceito Central: Em vez de modelar a difusão a partir de ruído gaussiano, o framework formula o processo de difusão forward (adição de ruído) usando o próprio ruído de quantização ($\epsilon_n$) gerado entre a imagem original ($I_0$) e o resultado inicial da compressão ($I_T$).
• Processo Forward: A imagem original é difusa até o resultado da compressão inicial adicionando escalas controladas do ruído de quantização real, em vez de ruído aleatório.
• Processo Reverse (Inferência): A inferência começa diretamente da imagem comprimida inicial ($I_T$) e remove o ruído de quantização passo a passo para recuperar a imagem original.
• Vantagem: Isso elimina a necessidade de adicionar ruído externo, garantindo que a reconstrução seja fiel à imagem original e eliminando a incerteza gerada por amostragens aleatórias.

B. Módulo de Filtragem Adaptativa no Domínio da Frequência (AFF)

• Para recuperar detalhes de alta frequência perdidos na compressão inicial, os autores desenvolveram um módulo AFF (Adaptive Frequency-domain Filtering).
• Este módulo é inserido nas conexões de "skip" (pulos) da arquitetura U-Net do modelo de difusão.
• Ele transforma as características do codificador para o domínio de Fourier, aplica uma máscara aprendível para realçar seletivamente as componentes de alta frequência e as transforma de volta. Isso melhora a reconstrução de texturas finas sem afetar as informações de baixa frequência.

C. Método de Melhoria Guiada por Amostra (Zero-Shot Sample-Guided)

• Durante a inferência, um método de otimização sem treinamento prévio (zero-shot) é utilizado.
• Utiliza-se uma função de perda baseada no CLIP (Contrastive Language-Image Pre-training) para guiar o processo de amostragem.
• O gradiente dessa perda é aplicado para ajustar a média da distribuição de difusão, forçando a imagem gerada a manter a fidelidade semântica e perceptual da imagem comprimida inicial, evitando que o modelo "alucine" detalhes que não estavam presentes.

3. Principais Contribuições

1. Resolução da Incompatibilidade de Ruído: O NC-Diffusion é a primeira abordagem a modelar explicitamente o ruído de quantização como o ruído de difusão, resolvendo o problema de distribuição e permitindo reconstruções fiéis.
2. Eficiência de Inferência: Ao iniciar a inferência a partir da imagem comprimida (e não do ruído puro), o método atinge alta qualidade em menos etapas, superando a eficiência de métodos baseados em Stable Diffusion ou DDPMs tradicionais.
3. Módulo AFF Plug-and-Play: Um módulo de filtragem no domínio da frequência que pode ser integrado a qualquer arquitetura U-Net para melhorar a preservação de detalhes de alta frequência.
4. Equilíbrio Distorsão-Percepção: O método consegue melhorar a qualidade perceptual sem sacrificar significativamente as métricas de distorção objetiva (como PSNR), algo raro em métodos generativos.

4. Resultados Experimentais

Os experimentos foram realizados nos conjuntos de dados Kodak e CLIC2020, comparando o método com o estado da arte (SOTA), incluindo HiFiC, MS-ILLM, DiffEIC, DIRAC e métodos tradicionais como BPG.

• Desempenho R-D (Taxa-Distorção) e R-P (Taxa-Percepção): O NC-Diffusion superou todos os métodos concorrentes, alcançando o melhor equilíbrio entre PSNR (fidelidade) e métricas perceptuais (LPIPS, FID, MS-SSIM).
• Qualidade Visual: As imagens reconstruídas apresentaram texturas mais nítidas e menos artefatos de borrão em comparação com métodos baseados em VAE ou GANs.
• Fidelidade: Diferente de métodos como MS-ILLM e HiFiC, que às vezes alteram detalhes críticos (ex: letras ou padrões específicos) para melhorar a pontuação perceptual, o NC-Diffusion manteve a estrutura original da imagem com alta fidelidade.
• Velocidade: O método demonstrou uma velocidade de decodificação significativamente superior a outros métodos baseados em difusão (como CDC e DiffEIC), pois não requer a geração a partir do ruído zero.
• Estudos de Ablação:
  • A remoção do NC-Diffusion (usando ruído gaussiano) causou uma queda drástica de ~12% na eficiência de bits (BD-Rate).
  • O módulo AFF e a perda de alta frequência ($L_{high}$) contribuíram adicionalmente para ganhos de ~10% e ~6% respectivamente.

5. Significado e Impacto

Este trabalho representa um avanço significativo na interseção entre modelos generativos e compressão de dados.

• Mudança de Paradigma: Demonstra que modelos de difusão não precisam ser usados apenas para "gerar" imagens a partir do nada, mas podem ser adaptados para "refinar" dados existentes de forma determinística e eficiente.
• Aplicabilidade Prática: Ao resolver o problema da incerteza e da ineficiência, o NC-Diffusion torna viável o uso de modelos de difusão em cenários reais de compressão onde a fidelidade da imagem original é obrigatória (ex: imagens médicas, satélite, arquivamento), algo que métodos generativos anteriores não conseguiam garantir.
• Eficiência Computacional: A capacidade de obter resultados de alta qualidade com menos etapas de inferência torna a tecnologia mais acessível para dispositivos com recursos limitados.

Em resumo, o NC-Diffusion estabelece um novo estado da arte ao alinhar a teoria da difusão com a física da quantização, oferecendo compressão de imagem de alta fidelidade que é simultaneamente fiel à origem e perceptualmente superior.