UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems

Resumen Técnico: UOTIP – Mapa de Transporte Óptimo Desbalanceado para Problemas Inversos No Emparejados

Planteamiento del Problema
El artículo aborda problemas inversos de imágenes no emparejados, un escenario desafiante donde el objetivo es recuperar una señal limpia desconocida $x$ a partir de mediciones ruidosas $y$ (modeladas como $y = A(x) + n$), pero los datos de entrenamiento consisten en conjuntos independientes y no emparejados de mediciones ruidosas y señales objetivo limpias. Los métodos supervisados tradicionales requieren datos emparejados $(y, x)$, los cuales a menudo no están disponibles en escenarios del mundo real. Si bien se han propuesto modelos generativos y métodos basados en Transporte Óptimo (OT), las formulaciones estándar de OT imponen restricciones estrictas de coincidencia de marginales ($T_\#\mu = \nu$), lo que puede ser frágil al tratar con valores atípicos, desequilibrios de clases entre conjuntos de datos o niveles de ruido variables. Además, los problemas inversos suelen estar mal planteados, lo que significa que el mapeo $y \to x$ puede no tener una solución única, complicando el aprendizaje de un mapa de transporte determinista.

Metodología
Los autores proponen UOTIP (Mapa de Transporte Óptimo Desbalanceado para Problemas Inversos), un solucionador novedoso que formula la tarea de reconstrucción como el aprendizaje de un mapa de Transporte Óptimo Desbalanceado (UOT) desde la distribución de mediciones ruidosas ($\mu$) hasta la distribución de señales limpias ($\nu$).

1. Formulación UOT: A diferencia del OT estándar, el UOT relaja la restricción de coincidencia exacta de marginales introduciendo penalizaciones de divergencia $f$ ($D_{\Psi_1}, D_{\Psi_2}$) sobre las marginales. Esto permite que el plan de transporte maneje desviaciones en la masa entre las distribuciones de origen y destino, proporcionando robustez frente a valores atípicos y desequilibrios de clases.
2. Función de Costo Basada en Verosimilitud: Para alinear el mapa de transporte con la estimación de Máxima A Posteriori (MAP), los autores introducen una función de costo basada en verosimilitud:
   $$c_l(y, x) = \|A(x) - y\|_2^2$$
   Bajo supuestos de ruido gaussiano, esto es proporcional al logaritmo negativo de la verosimilitud. Minimizar este costo asegura la fidelidad de los datos (la señal reconstruida produce la medición observada).
3. Regularización Cuadrática para Bien Planteamiento: Para abordar la naturaleza mal planteada de los problemas inversos (donde el costo de verosimilitud por sí solo puede violar la "condición de giro" requerida para la existencia y unicidad de un mapa OT), los autores incorporan un término de costo cuadrático:
   $$c_q(y, x) = \|y - x\|_2^2$$
   El costo total es una suma ponderada: $c(y, x) = \tau(c_l(y, x) + c_q(y, x))$. Teóricamente, los autores demuestran que añadir este término cuadrático asegura que el costo total satisfaga la condición de giro, garantizando la existencia y unicidad del mapa de transporte incluso para problemas mal planteados (por ejemplo, desenfoque).
4. Implementación Neuronal: El modelo aprende el mapa UOT $T_\theta$ y una función potencial $v_\phi$ utilizando una formulación semi-dual de UOT. El objetivo de aprendizaje minimiza una pérdida derivada de los conjugados convexos de las funciones de penalización, actuando efectivamente como un estimador MAP global y no emparejado.

Contribuciones Clave

• Primer Solucionador No Emparejado Basado en UOT: El artículo introduce el primer modelo para problemas inversos no emparejados basado en Transporte Óptimo Desbalanceado, integrando específicamente una función de costo basada en verosimilitud.
• Garantías Teóricas: Los autores proporcionan una demostración teórica de que incorporar un término de costo cuadrático asegura la existencia y unicidad del mapa de transporte para problemas inversos mal planteados al satisfacer la condición de giro.
• Robustez ante Restricciones del Mundo Real: Al relajar las restricciones de marginales, el modelo demuestra una robustez inherente ante:
  • Ruido de observación multinivel: Manejo de escenarios donde los niveles de ruido varían entre muestras.
  • Desequilibrio de clases: Adaptación a situaciones donde los conjuntos de datos de origen y destino tienen proporciones de clases diferentes (por ejemplo, diferentes modos).
  • Tipos de ruido diversos: Generalización a distribuciones de ruido no gaussianas (Laplace, Poisson) a pesar de ser entrenado con un costo de verosimilitud basado en Gaussiana.
• Rendimiento de Vanguardia: El modelo logra un rendimiento superior en benchmarks de problemas inversos tanto lineales (desenfoque gaussiano, superresolución) como no lineales (reconstrucción HDR, desenfoque no lineal).

Resultados Experimentales
Los autores evaluaron UOTIP en los conjuntos de datos FFHQ y AFHQ en cuatro tareas.

• Rendimiento Cuantitativo: UOTIP logró consistentemente las puntuaciones mejores o segundas mejores en PSNR, SSIM y FID en todas las tareas y conjuntos de datos, superando a líneas base existentes como NOT (OT Neuronal), OTUR (OT basado en GAN) y RCOT.
• Ruido Multinivel: En experimentos con niveles de ruido mixtos ($\sigma_y \in \{0.025, 0.05, 0.1, 0.2\}$), UOTIP superó significativamente a las variantes estándar de OT y otras líneas base, demostrando que la formulación desbalanceada maneja eficazmente distribuciones de ruido heterogéneas.
• Desequilibrio de Clases: Cuando las distribuciones de origen y destino tenían ratios de clases desequilibrados (por ejemplo, imágenes de gatos vs. perros), UOTIP mantuvo un rendimiento superior en comparación con las líneas base, las cuales a menudo se degradaban a medida que aumentaba el desequilibrio.
• Estudios de Ablación: Eliminar el costo de verosimilitud (usando solo costo cuadrático) aún produjo resultados competitivos, sugiriendo utilidad como solucionador ciego. Sin embargo, el modelo completo (verosimilitud + cuadrático) logró el mejor rendimiento general. El estudio también mostró que el modelo generaliza a tipos de ruido Laplace y Poisson incluso cuando se entrena con un costo de verosimilitud gaussiano.

Significado y Afirmaciones
El artículo afirma que UOTIP ofrece un avance significativo para los problemas inversos no emparejados al cerrar la brecha entre la teoría del transporte óptimo y la resolución práctica de problemas inversos. El significado radica en:

1. Relajar Restricciones Estrictas: La capacidad del marco UOT para relajar la coincidencia de marginales hace que el solucionador sea adaptable a datos del mundo real donde la alineación estricta entre conjuntos de datos ruidosos y limpios rara vez es alcanzable.
2. Solidez Teórica: El trabajo proporciona una base teórica rigurosa para resolver problemas inversos mal planteados mediante mapas de transporte neuronales, asegurando la unicidad de la solución a través de la condición de giro.
3. Robustez Práctica: La capacidad del modelo para manejar ruido multinivel y desequilibrio de clases sin reentrenamiento ni cambios arquitectónicos complejos lo hace particularmente efectivo para aplicaciones del mundo real donde las condiciones de los datos son variables e imperfectas.

Los autores señalan una limitación: el método actualmente depende de funciones de costo de forma fija sin capacidad de entrenamiento para el costo en sí mismo, aunque la verosimilitud basada en Gaussiana mostró una generalización prometedora a otros tipos de ruido.