An Expectation-Maximization Algorithm for Training Clean Diffusion Models from Corrupted Observations

Definição do Problema
Modelos de difusão (DMs) demonstraram uma capacidade excepcional em resolver problemas inversos de imagem computacional ao modelar priors complexos de imagem. No entanto, sua aplicação prática é severamente limitada pela dependência de conjuntos de dados de grande escala, de alta qualidade e limpos para o pré-treinamento. Em muitos cenários do mundo real, particularmente em imagens científicas e biomédicas (ex: TC de raios-X, criomicroscopia eletrônica), adquirir tais dados limpos é inviável ou proibitivamente caro. Por outro outro lado, observações corrompidas (com ruído, desfoque ou mascaramento) com modelos de propagação diferenciáveis são abundantes. Isso cria um dilema de "ovo e galinha": treinar um modelo de difusão preciso requer imagens limpas, mas reconstruir imagens limpas a partir de observações corrompidas requer um modelo de difusão bem treinado. Métodos existentes para aprender a partir de dados corrompidos, como AmbientDiffusion ou SURE-Score, frequentemente lutam contra tipos específicos de corrupção (ex: inpainting de larga escala ou ruído significativo) ou falham em generalizar através de diversos problemas inversos.

Metodologia: EMDiffusion
Os autores propõem o EMDiffusion, um novo framework de Expectativa-Maximização (EM) projetado para resolver conjuntamente problemas inversos de imagem e treinar um modelo de difusão diretamente a partir de observações corrompidas. O método opera através de um ciclo iterativo de duas etapas:

1. Inicialização: O processo começa com um prior de difusão "vago" treinado em um conjunto muito pequeno de imagens limpas (tão poucas quanto 10–50 imagens). Esta inicialização é crucial para evitar que o modelo colapse em uma distribuição definida unicamente pelos inputs corrompidos. Os autores observam que essas imagens limpas iniciais não precisam ser da mesma distribuição da tarefa alvo (dados fora da distribuição podem servir como inicializadores eficazes).

2. E-Step (Expectativa): Utilizando o modelo de difusão atual como um prior, o algoritmo realiza a Amostragem de Posterior de Difusão Adaptativa para reconstruir imagens limpas a partir das observações corrompidas.

   • A amostragem de posterior padrão muitas vezes falha nas primeiras iterações porque o prior é fraco, levando ao "colapso de modo" onde as amostras ignoram a observação e revertem para a distribuição de treinamento.
   • Para mitigar isso, o EMDiffusion introduz um hiperparâmetro $\lambda$ que reescala o termo de verossimilhança na equação diferencial estocástica (SDE) reversa.
   • $\lambda$ é ajustado adaptativamente em cada E-step para minimizar a perda de dados ($\|y - A(\hat{x}_0)\|_2^2$), equilibrando a influência do prior imperfeito contra os dados observados. Isso permite a geração de reconstruções mais limpas mesmo quando o prior é inicialmente ruim.

3. M-Step (Maximização): O algoritmo atualiza os pesos do modelo de difusão usando as imagens reconstruídas (amostras posteriores) geradas no E-step.

   • Esta etapa se assemelha ao pareamento de escore de denoising padrão, onde o modelo aprende a aproximar a função de escore dos dados reconstruídos "limpos".
   • Para acelerar o treinamento e evitar o overfitting para reconstruções ruidosas iniciais, o método emprega uma estratégia de treinamento de duas fases: inicialmente, ele faz o ajuste fino dos pesos herdados da iteração anterior; uma vez que a qualidade da reconstrução melhora (indicada pelo $\lambda$ ótimo caindo abaixo de 1), o modelo é reinicializado e retreinado do zero sobre o conjunto completo de amostras posteriores.

Principais Contribuições

• Framework EM para Dados Corrompidos: O artigo introduz o primeiro framework geral para treinar modelos de difusão baseados em escore a partir de observações arbitrariamente corrompidas usando uma abordagem EM, quebrando efetivamente a dependência de grandes conjuntos de dados limpos.
• Amostragem de Posterior Adaptativa: A proposta de um fator de escala adaptativo ($\lambda$) para o termo de verossimilhança na SDE reversa aborda o problema de colapso de modo inerente às iterações iniciais do EM, permitindo uma reconstrução estável mesmo com priors fracos.
• Refinamento Iterativo: O método demonstra que uma pequena quantidade de dados limpos pode atuar como um regularizador implícito, permitindo que o modelo convirja progressivamente para a verdadeira distribuição de dados limpos através de etapas alternadas de reconstrução e treinamento.
• Generalizabilidade: Ao contrário de trabalhos anteriores direcionados a corrupções específicas (ex: mascaramento), o EMDiffusion fornece uma abordagem unificada aplicável a inpainting aleatório, denoising e deblurring.

Resultados Experimentais
Os autores validam o EMDiffusion nos datasets CIFAR-10 e CelebA através de três tarefas: inpainting aleatório (mascaramento de 60% dos pixels), denoising Gaussiano e deblurring Gaussiano.

• Inpainting: O EMDiffusion supera significativamente os baselines como AmbientDiffusion e SURE-Score, alcançando uma qualidade de reconstrução (PSNR/LPIPS) comparável ao Diffusion Posterior Sampling (DPS) que utiliza um prior treinado no conjunto de dados limpos completo.
• Denoising & Deblurring: O método supera o SURE-Score e baselines autossupervisionados (Noise2Self), recuperando detalhes finos (ex: rodas de carros, características faciais) que outros métodos suavizam.
• Qualidade Generativa: Os modelos de difusão aprendidos alcançam pontuações FID de estado da arte em comparação com outros métodos treinados em dados corrompidos. Notavelmente, um modelo aprendido de 50.000 imagens corrompidas via EM performa melhor do que um modelo treinado em aproximadamente 15.000 imagens limpas.
• Eficiência: O método alcança estes resultados utilizando uma arquitetura DDPM padrão (35.7M de parâmetros), que é significativamente menor do que as arquiteturas usadas pelos baselines concorrentes (ex: os 113M+ parâmetros do AmbientDiffusion).

Significância e Alegações
O artigo afirma que o EMDiffusion demonstra que é informação-teoricamente possível aprender a distribuição subjacente de dados limpos primariamente a partir de medições corrompidas. A significância reside em:

1. Superar a Escassez de Dados: Oferece uma solução prática para domínios onde dados limpos não estão disponíveis, aproveitando a abundância de observações corrompidas.
2. Inicialização Robusta: Estabelece que um conjunto mínimo de imagens limpas (mesmo fora da distribuição) é suficiente para inicializar o manifold, evitando que o modelo colapse em uma distribuição enviesada de inputs corrompidos.
3. Desempenho de Estado da Arte: O método atinge um novo desempenho de estado da arte em tarefas de imagem inversa sem exigir modelos pré-treinados em grandes conjuntos de dados limpos, diminuindo a lacuna entre priors teóricos e restrições práticas de dados.

Os autores mantêm a modéstia quanto ao escopo, reconhecendo que o método depende da premissa de que o modelo de propagação é conhecido e diferenciável, e que trabalhos futuros poderiam se estender para aprender modelos de propagação desconhecidos ou estatísticas de ruído.