[논문 리뷰] The Monge Gap: A Regularizer to Learn All Transport Maps

Théo Uscidda, Marco Cuturi|arXiv (Cornell University)|2023. 02. 09.

Adversarial Robustness in Machine Learning인용 수 4

한 줄 요약

이 논문은 일반적인 신경망을 최적 운반(OT) 사상으로 훈련시킬 수 있도록 해주는 새로운 정규화 항인 몽게 갭을 소개한다. 아키텍처 제약 없이, OT 최적성에서의 이탈 정도를 측정하는 몽게 갭을 정규화 항으로 사용하여 피팅 손실을 최소화함으로써, 합성 및 단일세포 게놈 분석 벤치마크에서 ICNN과 MLP를 모두 능가하는 성능을 달성한다. 특히 고차원에서 뚜렷한 성능 향상을 보인다.

ABSTRACT

Optimal transport (OT) theory has been been used in machine learning to study and characterize maps that can push-forward efficiently a probability measure onto another. Recent works have drawn inspiration from Brenier's theorem, which states that when the ground cost is the squared-Euclidean distance, the ``best'' map to morph a continuous measure in $\mathcal{P}(\Rd)$ into another must be the gradient of a convex function. To exploit that result, [Makkuva+ 2020, Korotin+2020] consider maps $T= abla f_θ$, where $f_θ$ is an input convex neural network (ICNN), as defined by Amos+2017, and fit $θ$ with SGD using samples. Despite their mathematical elegance, fitting OT maps with ICNNs raises many challenges, due notably to the many constraints imposed on $θ$; the need to approximate the conjugate of $f_θ$; or the limitation that they only work for the squared-Euclidean cost. More generally, we question the relevance of using Brenier's result, which only applies to densities, to constrain the architecture of candidate maps fitted on samples. Motivated by these limitations, we propose a radically different approach to estimating OT maps: Given a cost $c$ and a reference measure $ρ$, we introduce a regularizer, the Monge gap $\mathcal{M}^c_ρ(T)$ of a map $T$. That gap quantifies how far a map $T$ deviates from the ideal properties we expect from a $c$-OT map. In practice, we drop all architecture requirements for $T$ and simply minimize a distance (e.g., the Sinkhorn divergence) between $T\sharpμ$ and $ν$, regularized by $\mathcal{M}^c_ρ(T)$. We study $\mathcal{M}^c_ρ$, and show how our simple pipeline outperforms significantly other baselines in practice.

연구 동기 및 목표

• 최적 운반(OT) 학습에서 입력 볼록 신경망(ICNN)의 한계를 해결하고자 하며, 이는 훈련 불안정성, 복잡한 초기화, 비제곱 유클리드 거리 비용과의 호환성 부족을 포함한다.
• 표본 기반 측도를 다룰 때 브레니에 정리(볼록 잠재함수의 기울기)를 강제로 적용하는 것이 이론적으로 타당한지 의문을 제기하고자 한다.
• 모든 비용 함수 $ c $ 에 대해 작동하는 일반적이고 아키텍처에 종속되지 않는 OT 사상 학습 방법을 개발하고자 한다.
• 모든 아키텍처 제약 없이, 사상 $ T $ 가 최적 OT 사상에 얼마나 가까운지를 측정하는 정규화 항인 몽게 갭 $ \mathcal{M}^{c}_{\rho}(T) $ 을 도입하고자 한다.

제안 방법

• 사상 $ T $ 가 최적 OT 행동에서 얼마나 이탈되어 있는지를 측정하는 몽게 갭 $ \mathcal{M}^{c}_{\rho}(T) = \mathbb{E}_{X \sim \rho}[c(X, T(X))] - W_c^2(\rho, T\sharp\rho) $ 을 제안한다.
• 몽게 갭을 훈련 목표 함수에 정규화 항으로 사용한다: $ T\sharp\mu $ 와 $ \nu $ 사이의 피팅 손실(예: 싱크본 발산)을 최소화하고, $ \lambda_{\mathrm{MG}} \cdot \mathcal{M}^{c}_{\rho}(T) $ 를 추가한다.
• 비용 함수가 트와이스트 조건을 만족할 경우, 최적 사상의 구조적 성질을 활용하여 $ T $ 를 잠재함수의 기울기로 매개변수화한다.
• 기준 측도 $ \rho $ 를 사용하며, 이는 원천 측도 $ \mu $ 와 독립적으로 선택 가능하므로 유연성과 강건성을 확보한다.
• 유한 표본 설정에서 몽게 갭을 미분 가능하고 계산 가능하게 하기 위해 $ c $-워샤르슈타인 거리에 대해 싱크본 추정기를 사용한다.
• 일반적인 비용 함수와 제곱 유클리드 비용의 경우 모두 적용 가능하며, 특히 제곱 유클리드 비용의 경우 몽게 갭이 볼록임을 보이며 이 성질이 유한 표본 추정에서도 유지되어 최적화의 안정성을 높인다.

실험 결과

연구 질문

• RQ1 입력 볼록 신경망(ICNN)의 아키텍처 제약을 제거함으로써, 훈련 불안정성과 복잡한 초기화 문제 없이 최적 운반 사상을 학습할 수 있는가?
• RQ2 연속 밀도가 아닌 표본 기반 측도를 다룰 때, 브레니에 정리(볼록 잠재함수의 기울기)를 강제 적용하는 것은 이론적으로나 실용적으로 타당한가?
• RQ3 아키텍처 제약 없이 OT 최적성에서의 이탈 정도를 측정하는 정규화 항을 사용할 경우, ICNN과 일반 MLP보다 OT 사상 학습에서 더 뛰어난 성능을 낼 수 있는가?
• RQ4 다양한 비용 함수에 대해 몽게 갭 정규화 항은 어떻게 행동하는가? 고차원 설정에서 하이퍼파rameter 선택에 대해 강건한가?
• RQ5 제안된 방법은 단일세포 옴리크스와 같은 실제 생물학적 데이터에 일반화되는가? 특히 세포가 자극에 대한 반응을 모델링하는 데 사용되는 운반 사상에 대해.

주요 결과

• 만약 $ T\sharp\mu = \nu $ 이고 $ \mathrm{Spt}(\mu) \subset \mathrm{Spt}(\rho) $ 라면, 몽게 갭 $ \mathcal{M}^{c}_{\rho}(T) $ 는 정확히 $ T $ 가 최적 OT 사상일 때에만 0이 되며, 이는 이 정규화 항의 이론적 타당성을 입증한다.
• 제곱 유클리드 비용의 경우 몽게 갭은 볼록이며, 이 성질이 싱크본 유한 표본 추정에서도 유지되어 안정적인 최적화를 가능하게 한다.
• Korotin 등(2021)의 벤치마크에서, 제안된 방법은 일반 MLP와 ICNN을 모두 능가하며, $ d \geq 16 $ 에서 특히 뚜렷한 성능 향상을 보이며, $ d \geq 64 $ 에서는 성능 격차가 더욱 커진다.
• 4i 및 scRNA-seq 데이터에서의 단일세포 게놈 실험에서는, 일반 MLP와 ICNN보다 더 뛰어난 예측 성능을 달성하였으며, 34개 및 9개의 치료 조건 전반에서 일관된 향상이 관찰되었다.
• 하이퍼파rameter 선택에 대해 강건하다: $ (\lambda_{\mathrm{MG}}, \lambda_{\mathrm{cons}}) $ 의 넓은 영역에서 양호한 성능을 보이며, $ d \leq 64 $ 에서는 최적 가중치가 $ (1, 0.01) $, $ d \geq 128 $ 에서는 $ (10, 0.1) $ 이다.
• ICNN는 가우시안 초기화에도 불구하고 불안정하고 평균 수준의 성능을 보이며, OT 학습에서 아키텍처 제약의 실용적 한계를 드러낸다.