An Unconditional Representation of the Conditional Score in Infinite-Dimensional Linear Inverse Problems

Resumo Técnico: Uma Representação Incondicional da Pontuação Condicional em Problemas de Inversão Linear em Dimensão Infinita

Enunciação do Problema
Modelos de difusão baseados em pontuação (SDMs) tornaram-se uma ferramenta poderosa para amostragem de distribuições posteriores em problemas de inversão bayesianos, oferecendo priores flexíveis que superam regularizadores tradicionais construídos manualmente. No entanto, os métodos existentes para amostragem posterior enfrentam um compromisso crítico entre eficiência computacional e escalabilidade.

1. Métodos baseados em correção (por exemplo, adicionar gradientes de verossimilhança à pontuação ou impor consistência com os dados) exigem avaliações repetidas do operador direto $A$ durante o processo de amostragem. Isso cria um gargalo computacional proibitivo para problemas de grande escala e alta dimensão.
2. Métodos de treinamento condicional direto aprendem uma função de pontuação condicionada às medições ($s(x, t; \mu_y)$). Embora isso evite avaliações diretas durante a amostragem, aumenta drasticamente a dimensionalidade de entrada da rede neural (de $x$ para $(x, y)$), exigindo dados de treinamento extensos e arquiteturas complexas para aproximar com precisão a pontuação condicional, particularmente em configurações de dimensão infinita.

O artigo aborda o desafio de realizar amostragem posterior eficiente e escalável em problemas de inversão linear em dimensão infinita, sem incorrer no alto custo online de avaliações do modelo direto ou na alta complexidade de amostragem do treinamento condicional direto.

Metodologia: Representação Incondicional da Pontuação Condicional (UCoS)
Os autores propõem o UCoS, um framework que desloca o ônus computacional do condicionamento da fase online de amostragem para uma fase offline de treinamento. A ideia central é que, para problemas de inversão linear com ruído gaussiano, a função de pontuação condicional pode ser derivada exatamente a partir de uma função de pontuação incondicional de um processo de difusão modificado e dependente da tarefa, por meio de transformações afins.

• Fundamentação Teórica:

  • O método opera em um espaço de Hilbert separável $\mathcal{H}$, garantindo invariância à discretização.
  • Em vez de treinar uma rede para aproximar $s(x, t; \mu_y)$, os autores definem um novo processo de difusão $\tilde{X}_t$ cuja distribuição é uma mistura do prior e de um termo gaussiano dependente do operador direto $A$.
  • Eles provam (Teorema 3.7) que a verdadeira pontuação condicional $s(x, t; \mu_y)$ está relacionada à pontuação incondicional $\tilde{s}$ deste novo processo pela identidade:
    $$s(x, t; \mu_y) = \lambda(t) \left( r(\xi_t(x, y)) - e^{-t/2}x \right)$$
    onde $r$ é uma versão transformada da pontuação incondicional, $\lambda(t)$ é um escalar dependente do tempo e $\xi_t(x, y)$ é uma transformação afim envolvendo a medição $y$ e o operador adjunto $A^*$.
  • Crucialmente, o termo $A^*\Gamma^{-1}y$ (onde $\Gamma$ é a covariância do ruído) pode ser computado uma única vez antes do início da amostragem. O processo de amostragem subsequente requer apenas a avaliação da rede treinada $r_\theta$ e adições simples de vetores, eliminando a necessidade de quaisquer aplicações adicionais de $A$ ou $A^*$ durante as etapas de difusão reversa.

• Implementação:

  • Fase Offline: Uma rede neural (especificamente um Operador Neural de Fourier, FNO) é treinada para aproximar a função $r$. O treinamento utiliza um objetivo de correspondência de pontuação com remoção de ruído no processo modificado $\tilde{X}_t$. Esta fase requer acesso a $A$ e $A^*$, mas nenhum dado de medição específico $y$.
  • Fase Online: Amostras posteriores são geradas resolvendo a EDE (Equação Diferencial Estocástica) de tempo reverso usando a pontuação transformada. O processo é "livre de matrizes", exigindo apenas o deslocamento pré-computado $A^*\Gamma^{-1}y$ e a avaliação da rede.

Principais Contribuições

1. Identidade Exata em Dimensões Infinitas: O artigo estabelece uma identidade teórica rigorosa (Teorema 3.7) conectando a pontuação condicional a uma pontuação incondicional dependente da tarefa por meio de transformações afins em espaços de Hilbert de dimensão infinita. Isso generaliza o princípio além de discretizações específicas de dimensão finita.
2. Análise de Convergência: Os autores fornecem um limite de erro quantitativo (Teorema 4.1) para o processo generativo, caracterizando explicitamente as contribuições de erro provenientes da perda de treinamento, discretização numérica e inicialização. Isso justifica a convergência do método para a medida posterior verdadeira.
3. Escalabilidade e Eficiência: Ao desacoplar o condicionamento do loop de amostragem, o UCoS alcança a eficiência computacional de métodos incondicionais (sem avaliações diretas durante a amostragem) enquanto mantém a precisão dos métodos condicionais. Ele evita a maldição da dimensionalidade associada ao treinamento de redes no espaço conjunto $(x, y)$.

Resultados Experimentais
Os autores validam o UCoS em três problemas de inversão: preenchimento de imagens (inpainting), tomografia computadorizada (TC) de visão esparsa e desfocagem de imagens.

• Preenchimento de Imagens (Prior Gaussiano): Em um cenário onde a posterior é gaussiana e as pontuações podem ser computadas em forma fechada, o UCoS desempenha de forma comparável ao método condicional exato e a outras bases, confirmando a identidade teórica.
• Imagem de TC e Desfocagem (Priors Aprendidos): Utilizando os conjuntos de dados LIDC-IDRI e CelebA, respectivamente, o UCoS é comparado contra SDE ALD, Amostragem Posterior de Difusão (DPS), métodos baseados em projeção (Proj) e treinamento Condicional direto.
  • Precisão: O UCoS produz amostras com viés e desvio padrão comparáveis ou melhores que métodos incondicionais e significativamente melhores que o método Condicional ao usar arquiteturas de baixa complexidade (por exemplo, 32 nós por camada).
  • Eficiência: O UCoS iguala a velocidade de amostragem do método Condicional (ambos evitam avaliações diretas), mas supera significativamente os métodos baseados em correção (DPS, Proj, SDE ALD), que exigem passagens diretas repetidas.
  • Robustez: O desempenho do método Condicional degrada-se notavelmente com menos parâmetros de rede devido ao aumento da complexidade de aproximar a distribuição conjunta, enquanto o UCoS mantém alta qualidade com arquiteturas mais simples.

Significado e Alegações
O artigo alega que o UCoS oferece uma solução "o melhor dos dois mundos": elimina o gargalo computacional de avaliações de modelos diretos durante a amostragem (uma limitação dos métodos baseados em correção) enquanto evita a alta complexidade de amostragem e as demandas arquitetônicas do treinamento condicional direto.

Os autores enfatizam que sua abordagem é invariante à discretização, tornando-a adequada para problemas de inversão em dimensão infinita onde garantias de desempenho frequentemente falham em aproximações de dimensão finita de alta dimensão. Eles reconhecem que o método é dependente da tarefa, exigindo re-treinamento se o operador direto $A$ mudar, e que a correspondência exata depende da linearidade do modelo direto e da gaussianidade da verossimilhança (permitindo o passo de "completar o quadrado"). O trabalho é apresentado como uma prova de conceito para uma nova classe de amostradores baseados em difusão que deslocam a complexidade do condicionamento para a fase de treinamento, sem erro de aproximação na formulação teórica.