Diffusion-OAMP for Joint Image Compression and Wireless Transmission

1. Declaração do Problema

O artigo aborda o desafio da Compressão de Imagem e Transmissão Sem Fio Conjuntas (JCWT). Diferente da restauração genérica de imagens, a JCWT envolve a recuperação de uma imagem de origem ($s$) a partir de um sinal recebido ($y$) que sofreu duas transformações lineares distintas:

1. Compressão e Modulação: A imagem de origem é comprimida e modulada em um sinal transmitido ($x$) por meio de um operador $F$.
2. Transmissão Sem Fio: O sinal passa por um canal sem fio ($A$) corrompido por Ruído Gaussiano Branco Aditivo (AWGN, $n$).

O sistema é modelado como um problema inverso linear:
$$y = A(Fs) + n = Ax + n$$
onde $A$ representa a matriz do canal e $F$ representa o operador conjunto de compressão/modulação.

Desafios Principais:

• Condicionamento Deficiente: O operador combinado $AF$ frequentemente resulta em um sistema severamente mal-condicionado, especialmente em canais de desvanecimento (por exemplo, 3GPP TDL-A) e baixas taxas de compressão.
• Incompatibilidade de Priori: Algoritmos tradicionais de recuperação, como a Passagem de Mensagens Aproximada Ortogonal (OAMP), dependem de priores analíticos (por exemplo, Bernoulli-Gaussiano) ou denoisers artesanais (por exemplo, BM3D). Estes falham em capturar a distribuição complexa e de alta dimensão de imagens naturais, levando a desempenho subótimo quando o priori assumido não corresponde à verdadeira distribuição da imagem.

2. Metodologia: Diffusion-OAMP

Os autores propõem o Diffusion-OAMP, um framework de reconstrução livre de treinamento que integra modelos generativos de difusão pré-treinados no algoritmo OAMP. A ideia central é utilizar o modelo de difusão como um poderoso Estimador Não-Linear (NLE) dentro do loop iterativo do OAMP.

A. Arquitetura do Sistema

O framework opera no domínio de transmissão ($x$), mas aproveita o domínio de origem ($s$) para priores generativos:

1. Estimador Linear (LE): Um estimador LMMSE processa a observação ruidosa $y$ para produzir uma estimativa a posteriori de $x$.
2. Ortogonalização: A saída é ortogonalizada para desacoplar o estimador linear do estimador não-linear, produzindo "pseudo-observações" extrínsecas.
3. Multiplexação Aleatória Inversa (RM): A estimativa é mapeada do domínio de transmissão de volta ao domínio da imagem de origem ($s$) usando a inversa do operador de compressão ($F^{-1}$).
4. Correspondência de SNR: O algoritmo calcula a variância de ruído efetiva da estimativa atual e mapeia-a para um específico passo de tempo de difusão ($t^*$). Isso garante que o modelo de difusão opere no nível de ruído correto correspondente à incerteza da iteração atual.
5. Estimador Não-Linear (NLE): Um modelo de difusão pré-treinado (seja DDIM ou Flow Matching) atua como um denoiser, projetando a estimativa da imagem ruidosa na variedade de dados aprendida.
6. Correção MMSE: A imagem denoised é remodulada ($F$), e um passo de correção atualiza o priori para a próxima iteração do OAMP usando o Estimador de Risco Não Viesado de Stein (SURE) para garantir ortogonalidade.

B. Componentes Técnicos Principais

• Multiplexação Aleatória (RM): O operador de compressão $F$ combina uma matriz de sinais aleatórios, uma matriz DCT, uma matriz de permutação e um seletor de linhas. Este design garante que a matriz de canal equivalente seja invariante unitariamente à direita, um requisito para a análise de evolução de estado do OAMP.
• Regra de Correspondência de SNR: O passo de tempo de difusão $t^*$ é selecionado dinamicamente com base na variância de erro do OAMP ($v_{post}^{orth}$):
  • Para Flow Matching: $t^* = \frac{1}{1 + \sqrt{v_{post}^{orth}}}$
  • Para DDIM: $t^* = \frac{1}{1 + v_{post}^{orth}}$
  • À medida que as iterações progridem e a variância de erro diminui, $t^*$ move-se em direção à extremidade "limpa" da trajetória de difusão.
• Livre de Treinamento: O framework não requer o re-treinamento dos modelos de difusão; ele utiliza modelos pré-treinados existentes (por exemplo, em CelebA) como componentes plug-and-play.

3. Contribuições Principais

1. Framework Inovador: Introdução do Diffusion-OAMP, o primeiro framework a incorporar modelos de difusão especificamente no algoritmo OAMP para a tarefa conjunta de compressão de imagem e transmissão sem fio.
2. Flexibilidade de Priori: O método é compatível com múltiplos priores generativos. Os autores o instanciam com DDIM (Modelos Implícitos de Difusão com Denoising) e Flow Matching, demonstrando versatilidade.
3. Mecanismo de Correspondência de SNR: Uma regra inovadora para alinhar a variância de ruído interna do algoritmo OAMP com o cronograma de ruído do modelo de difusão, permitindo integração perfeita sem ajuste fino.
4. Robustez à Distorção do Canal: O framework lida efetivamente com o condicionamento deficiente causado por canais de desvanecimento sem fio realistas (3GPP TDL-A), onde métodos tradicionais frequentemente falham.

4. Resultados Experimentais

Os autores avaliaram o framework no conjunto de dados CelebA (256x256) sob taxas de compressão variáveis ($\beta \in \{0.1, 0.4, 0.7\}$) e níveis de ruído ($\sigma \in \{0.05, 0.5\}$).

• Compressão Pura de Imagem (AWGN):
  • O Diffusion-OAMP (DDIM) alcançou o maior PSNR, particularmente em regimes de alto ruído ($\sigma=0.5$), superando o OAMP+BM3D e o DDRM.
  • O Flow Matching mostrou superioridade em SSIM (similaridade estrutural) em alguns casos e convergência mais rápida (atingindo PSNR quase ótimo em ~6 NFEs).
• Compressão e Transmissão Sem Fio Conjuntas (Canais de Desvanecimento):
  • Em canais de desvanecimento 3GPP TDL-A realistas, a matriz de observação torna-se severamente mal-condicionada.
  • O Diffusion-OAMP manteve reconstrução estável, superando significativamente as linhas de base. Por exemplo, em baixa compressão ($\beta=0.1$) e alto ruído ($\sigma=0.5$), o Diffusion-OAMP (DDIM) alcançou 21,95 dB PSNR, enquanto o OAMP+BM3D caiu para 16,20 dB.
  • Resultados visuais mostraram preservação superior de detalhes faciais e identidade, evitando o desfoque e artefatos vistos nos métodos de linha de base.
• Convergência:
  • O algoritmo converge rapidamente, estabilizando-se dentro de 3 iterações do OAMP em todos os cenários testados.

5. Significado

• Ponte entre IA Generativa e Comunicações: Este trabalho demonstra uma aplicação prática de modelos generativos em larga escala (Difusão/Flow Matching) em comunicações da camada física, indo além da simples restauração de imagens para problemas complexos de codificação conjunta fonte-canal.
• Tratamento de Problemas Mal-Postos: Ao utilizar um priori generativo aprendido, o framework regulariza o problema inverso de forma mais eficaz do que priores artesanais, tornando-o robusto contra distorções severas do canal e restrições de baixa largura de banda.
• Eficiência: A natureza "livre de treinamento" da abordagem significa que ela pode aproveitar modelos generativos existentes e de última geração sem o custo computacional de re-treiná-los para condições específicas de canal.
• Direção Futura: Abre um novo caminho para "Comunicação Semântica" ou design de "Camada Física baseada em Aprendizado Profundo", onde modelos generativos avançados servem como a lógica central do receptor para sistemas de comunicação inteligentes de próxima geração.