[논문 리뷰] Unbalancedness in Neural Monge Maps Improves Unpaired Domain Translation

Luca Eyring, Dominik Klein외 5명|arXiv (Cornell University)|2023. 11. 25.

Cancer-related molecular mechanisms research인용 수 4

한 줄 요약

이 논문은 신경 몽주 맵 추정기의 어떤 것도에 비균형 최적 운반(UOT)을 이론적으로 탄탄한 방법으로 통합하기 위해 원천 및 대상 측도를 재스케일링하는 방법을 제안한다. 이는 비대응 도메인 번역에서 상당한 성능 향상을 이끌어낸다. 이 방법은 단일세포 데이터와 이미지 번역 모두에서 특징 보존과 궤적 추론을 향상시키며, UOT-FM은 운반 비용을 줄이고 분포 이탈을 더 잘 다루는 바에 있어 OT-FM을 능가한다.

ABSTRACT

In optimal transport (OT), a Monge map is known as a mapping that transports a source distribution to a target distribution in the most cost-efficient way. Recently, multiple neural estimators for Monge maps have been developed and applied in diverse unpaired domain translation tasks, e.g. in single-cell biology and computer vision. However, the classic OT framework enforces mass conservation, which makes it prone to outliers and limits its applicability in real-world scenarios. The latter can be particularly harmful in OT domain translation tasks, where the relative position of a sample within a distribution is explicitly taken into account. While unbalanced OT tackles this challenge in the discrete setting, its integration into neural Monge map estimators has received limited attention. We propose a theoretically grounded method to incorporate unbalancedness into any Monge map estimator. We improve existing estimators to model cell trajectories over time and to predict cellular responses to perturbations. Moreover, our approach seamlessly integrates with the OT flow matching (OT-FM) framework. While we show that OT-FM performs competitively in image translation, we further improve performance by incorporating unbalancedness (UOT-FM), which better preserves relevant features. We hence establish UOT-FM as a principled method for unpaired image translation.

연구 동기 및 목표

• 균형 최적 운반의 한계, 특히 외곽치와 분포 이탈에 대한 민감성 문제를 해결하기 위해.
• 대안적 훈련이 필요 없이 어떤 신경 몽주 맵 추정기에도 비균형성을 통합할 수 있는 일반화 가능한 프레임워크를 개발하기 위해.
• 질량 이탈을 모델링하여 단일세포 데이터의 궤적 추론과 자극에 대한 세포 반응 예측을 향상시키기 위해.
• OT 플로우 매칭(OT-FM) 프레임워크를 비균형 설정(비균형-플로우 매칭, UOT-FM)으로 확장하여 이미지 번역 과업에서 성능을 향상시키기 위해.
• 비균형성이 실제 응용에서 더 생물학적으로 타당하고 정확한 매핑을 이끌어내는지 입증하기 위해.

제안 방법

• 서로 다른 질량을 가진 측도 간의 몽주 맵 계산을 재스케일링된 측도 간의 균형 잡힌 맵으로 재구성하기 위한 재스케일링 기법을 제안한다.
• 이론적으로 비균형성을 측도 재스케일링을 통해 어떤 몽주 맵 추정기에도 통합할 수 있음을 증명하며, 질량 이탈에 대한 처벌을 보장한다.
• OT-ICNN, 몽주 갭, OT-FM 등 여러 추정기들에 이 프레임워크를 적용하여 연속적인 정규화 플로우를 위한 UOT-FM을 가능하게 한다.
• 시뮬레이션 또는 플로우 기반 생성 모델링이 필요 없이, 미니배치 최적 운반 커플링을 사용해 UOT-FM을 훈련시킨다.
• 엔트로피 정규화와 비균형 처벌을 포함한 운반 비용 최소화 목적함수를 사용하여 훈련을 안정화시킨다.
• 합성 데이터, 단일세포 궤적, 이미지 번역(EMNIST)에서 이 방법을 검증하며, 기준 OT 및 기존 방법과 비교한다.

Figure 1: Comparison of balanced and unbalanced Monge map computed on the EMNIST dataset translating digits $\rightarrow$ letters . Source and target distribution are rescaled leveraging the unbalanced OT coupling. The computed balanced mapping includes $8\rightarrow\{O,B\}$ , and $1\rightarrow\{O,I

실험 결과

연구 질문

• RQ1 대안적 훈련에 의존하지 않고 비균형 최적 운반을 신경 몽주 맵 추정기에 체계적으로 통합할 수 있는가?
• RQ2 비균형성을 통합하면 단일세포 데이터에서 세포 궤적 추론의 정확성과 생물학적 타당성이 향상되는가?
• RQ3 비균형성이 비대응 이미지 번역 과업에서 특징 보존과 운반 비용에 어떤 영향을 미치는가?
• RQ4 UOT-FM 프레임워크는 계산 효율성을 유지하면서도 표준 OT-FM을 능가하는가?
• RQ5 비균형성이 도메인 번역에서 분포 이탈과 외곽치의 영향을 어느 정도 완화하는가?

주요 결과

• UOT-FM은 비대응 이미지 번역 과업에서 OT-FM보다 뛰어난 성능을 보이며, 특히 입력 특징 보존과 운반 비용 감소에서 뛰어나다.
• EMNIST 데이터셋에서 비균형성이 숫자를 올바르게 문자로 매핑한다(예: 8→B, 1→I), 균형 잡힌 OT에서 관찰된 잘못된 매핑(예: 8→{O,B})를 방지한다.
• 단일세포 데이터에서 UOT-ICNN은 Ngn3 높은 세포가 췌장내분비 계열로 전환되는 등 더 생물학적으로 타당한 세포 유형 전환을 복원하며 전이 확률 정확도가 향상된다.
• 이 방법은 이미지 번역 과업에서 균형 잡힌 기준 모델 대비 운반 비용을 최대 25%까지 감소시킨다. 이는 질량 운반의 효율성이 높아졌음을 시사한다.
• 자연 이미지 번역 벤치마크에서 경쟁적인 성능을 유지하면서도 분포 이탈 상황에서 더 우수한 일반화 성능을 보인다.
• 이론적 분석을 통해 원천 및 대상 측도의 재스케일링이 비균형 조건 하에서도 유효한 몽주 맵 계산을 가능하게 하며, 이는 본 방법의 원리적인 기반을 제공한다.