UOTIP: Unbalanced Optimal Transport Map for Unpaired Inverse Problems

기술 요약: UOTIP – 비정합 역문제용 불균형 최적 수송 지도

문제 정의
본 논문은 비정합 (unpaired) 이미지 역문제를 다루며, 이는 잡음이 섞인 측정값 $y$ (모델: $y = A(x) + n$) 로부터 알려지지 않은 깨끗한 신호 $x$를 복원하는 것을 목표로 하지만, 훈련 데이터는 잡음이 섞인 측정값과 깨끗한 목표 신호의 독립적이고 짝지어지지 않은 (non-paired) 집합으로 구성된 까다로운 설정을 의미합니다. 기존의 지도 학습 방법은 실제 시나리오에서 종종 구할 수 없는 짝지어진 $(y, x)$ 데이터를 요구합니다. 생성 모델과 최적 수송 (OT) 기반 방법들이 제안되었으나, 표준 OT 공식은 엄격한 주변부 (marginal) 일치 제약 ($T_\#\mu = \nu$) 을 부과하여, 이상치, 데이터셋 간의 클래스 불균형, 또는 다양한 잡음 수준을 다룰 때 취약할 수 있습니다. 또한, 역문제는 종종 잘 정의되지 않아 (ill-posed), $y \to x$ 매핑이 고유한 해를 가지지 않을 수 있으므로 결정론적 수송 지도의 학습을 복잡하게 만듭니다.

방법론
저자들은 역문제용 불균형 최적 수송 (UOT) 지도를 잡음이 섞인 측정값의 분포 ($\mu$) 에서 깨끗한 신호의 분포 ($\nu$) 로 학습하는 새로운 솔버인 UOTIP(Unbalanced Optimal Transport Map for Inverse Problems) 을 제안합니다.

1. UOT 공식화: 표준 OT 와 달리, UOT 는 $f$-발산 패널티 ($D_{\Psi_1}, D_{\Psi_2}$) 를 주변부에 도입하여 정확한 주변부 일치 제약을 완화합니다. 이를 통해 수송 계획은 소스 및 목표 분포 간의 질량 편차를 처리할 수 있어 이상치와 클래스 불균형에 대한 강인성을 제공합니다.
2. 가능도 기반 비용 함수: 수송 지도를 최대 사후 확률 (MAP) 추정치와 정렬시키기 위해, 저자들은 가능도 비용 함수를 도입합니다:
   $$c_l(y, x) = \|A(x) - y\|_2^2$$
   가우시안 잡음 가정 하에서 이는 음의 로그 가능도에 비례합니다. 이 비용을 최소화하면 데이터 충실도 (복원된 신호가 관측된 측정값을 생성함) 를 보장합니다.
3. 잘 정의됨을 위한 2 차 정규화: 역문제의 잘 정의되지 않은 특성 (가능도 비용만으로는 OT 지도의 존재성과 고유성을 보장하는 'twist condition'을 위반할 수 있음) 을 해결하기 위해, 저자들은 2 차 비용 항을 통합합니다:
   $$c_q(y, x) = \|y - x\|_2^2$$
   전체 비용은 가중 합입니다: $c(y, x) = \tau(c_l(y, x) + c_q(y, x))$. 이론적으로 저자들은 이 2 차 항을 추가함으로써 전체 비용이 twist 조건을 만족하여, 비정합 문제 (예: 디블러링) 에 대해서도 수송 지도의 존재성과 고유성을 보장함을 증명합니다.
4. 신경망 구현: 모델은 UOT 의 반-이중 (semi-dual) 공식을 사용하여 UOT 지도 $T_\theta$와 잠재 함수 $v_\phi$를 학습합니다. 학습 목적 함수는 패널티 함수의 볼록 켤레 (convex conjugates) 로부터 유도된 손실 함수를 최소화하며, 이는 전역적 비정합 MAP 추정기로 작용합니다.

주요 기여

• 최초의 UOT 기반 비정합 솔버: 본 논문은 특히 가능도 기반 비용 함수를 통합하여, Unbalanced Optimal Transport 를 기반으로 한 비정합 역문제용 최초의 모델을 소개합니다.
• 이론적 보장: 저자들은 2 차 비용 항을 통합함으로써 twist 조건을 만족시켜 비정합 역문제에 대한 수송 지도의 존재성과 고유성을 보장한다는 이론적 증명을 제공합니다.
• 실제 제약에 대한 강인성: 주변부 제약을 완화함으로써, 모델은 다음과 같은 실제 세계 제약에 내재된 강인성을 입증합니다:
  • 다단계 관측 잡음: 샘플 간 잡음 수준이 변하는 시나리오 처리.
  • 클래스 불균형: 소스 및 목표 데이터셋의 클래스 비율이 다른 상황 (예: 다른 모드) 에 적응.
  • 다양한 잡음 유형: 가우시안 기반 가능도 비용으로 훈련되었음에도 불구하고 비가우시안 잡음 분포 (라플라스, 푸아송) 로의 일반화.
• 최첨단 성능: 모델은 선형 (가우시안 디블러링, 초해상도) 및 비선형 (HDR 복원, 비선형 디블러링) 역문제 벤치마크 모두에서 우수한 성능을 달성합니다.

실험 결과
저자들은 FFHQ 및 AFHQ 데이터셋을 대상으로 네 가지 작업에서 UOTIP 를 평가했습니다.

• 정량적 성능: UOTIP 는 모든 작업과 데이터셋에서 PSNR, SSIM, FID 점수에서 일관되게 최고 또는 차고 성능을 달성하여, NOT(Neural OT), OTUR(GAN 기반 OT), RCOT 와 같은 기존 베이스라인을 능가했습니다.
• 다단계 잡음: 혼합 잡음 수준 ($\sigma_y \in \{0.025, 0.05, 0.1, 0.2\}$) 을 가진 실험에서 UOTIP 는 표준 OT 변형 및 기타 베이스라인을 크게 능가하여, 불균형 공식이 이질적인 잡음 분포를 효과적으로 처리함을 입증했습니다.
• 클래스 불균형: 소스 및 목표 분포가 불균형한 클래스 비율 (예: 고양이 vs 개 이미지) 을 가질 때, UOTIP 는 불균형이 증가함에 따라 종종 성능이 저하되는 베이스라인에 비해 우수한 성능을 유지했습니다.
• 애블레이션 연구: 가능도 비용을 제거하고 (2 차 비용만 사용) 경쟁력 있는 결과를 여전히 산출하여 블라인드 솔버로서의 유용성을 시사했습니다. 그러나 전체 모델 (가능도 + 2 차) 이 전반적으로 가장 좋은 성능을 달성했습니다. 연구는 또한 가우시안 가능도 비용으로 훈련되었을 때 모델이 라플라스 및 푸아송 잡음 유형으로도 일반화됨을 보여주었습니다.

의의 및 주장
본 논문은 UOTIP 가 최적 수송 이론과 실제 역문제 해결 간의 간극을 메움으로써 비정합 역문제에 대한 중요한 진전을 제공한다고 주장합니다. 그 의의는 다음과 같습니다:

1. 엄격한 제약 완화: UOT 프레임워크의 주변부 일치 완화 능력은 잡음이 섞인 데이터와 깨끗한 데이터셋 간의 엄격한 정렬이 거의 달성되지 않는 실제 데이터에 솔버를 적응 가능하게 만듭니다.
2. 이론적 건전성: 이 작업은 신경 수송 지도를 통한 비정합 역문제 해결을 위한 엄격한 이론적 기반을 제공하여, twist 조건을 통해 해의 고유성을 보장합니다.
3. 실용적 강인성: 모델은 재훈련이나 복잡한 아키텍처 변경 없이 다단계 잡음과 클래스 불균형을 처리할 수 있어, 데이터 조건이 가변적이고 불완전한 실제 응용에 특히 효과적입니다.

저자들은 한 가지 제한 사항을 지적합니다: 현재 방법은 비용 자체에 대한 훈련 용량 없이 고정된 형태의 비용 함수에 의존하지만, 가우시안 기반 가능도는 다른 잡음 유형으로의 유망한 일반화를 보여주었습니다.