UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems

Resumo Técnico: UOTIP – Mapa de Transporte Ótimo Desbalanceado para Problemas Inversos Não Pareados

Enunciado do Problema
O artigo aborda problemas inversos de imagem não pareados, um cenário desafiador onde o objetivo é recuperar um sinal limpo desconhecido $x$ a partir de medições ruidosas $y$ (modeladas como $y = A(x) + n$), mas os dados de treinamento consistem em conjuntos independentes e não pareados de medições ruidosas e sinais alvo limpos. Métodos supervisionados tradicionais exigem dados pareados $(y, x)$, que frequentemente não estão disponíveis em cenários do mundo real. Embora modelos generativos e métodos baseados em Transporte Ótimo (OT) tenham sido propostos, formulações padrão de OT impõem restrições estritas de correspondência marginal ($T_\#\mu = \nu$), que podem ser frágeis ao lidar com valores atípicos, desequilíbrios de classes entre conjuntos de dados ou níveis de ruído variados. Além disso, problemas inversos são frequentemente mal-postos, o que significa que o mapeamento $y \to x$ pode não ter uma solução única, complicando a aprendizagem de um mapa de transporte determinístico.

Metodologia
Os autores propõem UOTIP (Mapa de Transporte Ótimo Desbalanceado para Problemas Inversos), um solucionador inovador que formula a tarefa de reconstrução como a aprendizagem de um mapa de Transporte Ótimo Desbalanceado (UOT) da distribuição de medições ruidosas ($\mu$) para a distribuição de sinais limpos ($\nu$).

1. Formulação UOT: Diferentemente do OT padrão, o UOT relaxa a restrição exata de correspondência marginal introduzindo penalidades de divergência $f$ ($D_{\Psi_1}, D_{\Psi_2}$) nas marginais. Isso permite que o plano de transporte lide com desvios de massa entre as distribuições de origem e destino, proporcionando robustez a valores atípicos e desequilíbrios de classes.
2. Função de Custo Baseada em Verossimilhança: Para alinhar o mapa de transporte com a estimativa de Máxima A Posteriori (MAP), os autores introduzem uma função de custo baseada em verossimilhança:
   $$c_l(y, x) = \|A(x) - y\|_2^2$$
   Sob suposições de ruído gaussiano, isso é proporcional ao log-verossimilhança negativo. Minimizar esse custo garante fidelidade aos dados (o sinal reconstruído produz a medição observada).
3. Regularização Quadrática para Bem-Postura: Para abordar a natureza mal-posta de problemas inversos (onde o custo de verossimilhança sozinho pode violar a "condição de torção" necessária para a existência e unicidade de um mapa de OT), os autores incorporam um termo de custo quadrático:
   $$c_q(y, x) = \|y - x\|_2^2$$
   O custo total é uma soma ponderada: $c(y, x) = \tau(c_l(y, x) + c_q(y, x))$. Teoricamente, os autores provam que adicionar esse termo quadrático garante que o custo total satisfaça a condição de torção, garantindo a existência e unicidade do mapa de transporte mesmo para problemas mal-postos (por exemplo, desbluramento).
4. Implementação Neural: O modelo aprende o mapa UOT $T_\theta$ e uma função potencial $v_\phi$ usando uma formulação semi-dual de UOT. O objetivo de aprendizagem minimiza uma perda derivada das conjugadas convexas das funções de penalidade, atuando efetivamente como um estimador MAP global e não pareado.

Principais Contribuições

• Primeiro Solucionador Não Pareado Baseado em UOT: O artigo introduz o primeiro modelo para problemas inversos não pareados baseado em Transporte Ótimo Desbalanceado, integrando especificamente uma função de custo baseada em verossimilhança.
• Garantias Teóricas: Os autores fornecem uma prova teórica de que incorporar um termo de custo quadrático garante a existência e unicidade do mapa de transporte para problemas inversos mal-postos ao satisfazer a condição de torção.
• Robustez a Restrições do Mundo Real: Ao relaxar restrições marginais, o modelo demonstra robustez inerente a:
  • Ruído de observação multinível: Lidando com cenários onde os níveis de ruído variam entre as amostras.
  • Desequilíbrio de classes: Adaptando-se a situações onde os conjuntos de dados de origem e destino têm proporções de classes diferentes (por exemplo, modos diferentes).
  • Tipos diversos de ruído: Generalizando para distribuições de ruído não gaussianas (Laplace, Poisson) apesar de ser treinado com um custo de verossimilhança baseado em Gaussiana.
• Desempenho de Estado da Arte: O modelo alcança desempenho superior em benchmarks de problemas inversos tanto lineares (desbluramento gaussiano, super-resolução) quanto não lineares (reconstrução HDR, desbluramento não linear).

Resultados Experimentais
Os autores avaliaram o UOTIP nos conjuntos de dados FFHQ e AFHQ em quatro tarefas.

• Desempenho Quantitativo: O UOTIP consistentemente alcançou as melhores ou segundas melhores pontuações em PSNR, SSIM e FID em todas as tarefas e conjuntos de dados, superando bases existentes como NOT (OT Neural), OTUR (OT baseado em GAN) e RCOT.
• Ruído Multinível: Em experimentos com níveis de ruído mistos ($\sigma_y \in \{0.025, 0.05, 0.1, 0.2\}$), o UOTIP superou significativamente variantes padrão de OT e outras bases, demonstrando que a formulação desbalanceada lida efetivamente com distribuições de ruído heterogêneas.
• Desequilíbrio de Classes: Quando as distribuições de origem e destino tinham proporções de classes desequilibradas (por exemplo, imagens de gatos vs. cães), o UOTIP manteve desempenho superior em comparação com as bases, que frequentemente degradavam conforme o desequilíbrio aumentava.
• Estudos de Ablação: Remover o custo de verossimilhança (usando apenas custo quadrático) ainda produziu resultados competitivos, sugerindo utilidade como solucionador cego. No entanto, o modelo completo (verossimilhança + quadrático) alcançou o melhor desempenho geral. O estudo também mostrou que o modelo generaliza para tipos de ruído Laplace e Poisson mesmo quando treinado com um custo de verossimilhança gaussiano.

Significado e Alegações
O artigo alega que o UOTIP oferece um avanço significativo para problemas inversos não pareados ao preencher a lacuna entre a teoria do transporte ótimo e a resolução prática de problemas inversos. O significado reside em:

1. Relaxamento de Restrições Estritas: A capacidade do framework UOT de relaxar a correspondência marginal torna o solucionador adaptável a dados do mundo real onde o alinhamento estrito entre conjuntos de dados ruidosos e limpos raramente é alcançável.
2. Fundamentação Teórica Sólida: O trabalho fornece uma base teórica rigorosa para resolver problemas inversos mal-postos via mapas de transporte neurais, garantindo unicidade da solução através da condição de torção.
3. Robustez Prática: A capacidade do modelo de lidar com ruído multinível e desequilíbrio de classes sem retreinamento ou mudanças complexas de arquitetura torna-o particularmente eficaz para aplicações do mundo real onde as condições de dados são variáveis e imperfeitas.

Os autores notam uma limitação: o método atualmente depende de funções de custo de forma fixa sem capacidade de treinamento para o próprio custo, embora a verossimilhança baseada em Gaussiana tenha mostrado generalização promissora para outros tipos de ruído.