기술 요약: 저계수 최적 운송을 위한 리만 기하학적 접근법

문제 정의

고전적 최적 운송(Optimal Transport, OT)은 확률 분포를 비교하는 엄격한 메트릭을 제공하지만, 조밀한 결합 행렬(coupling matrix)을 계산해야 하기 때문에 이차적인 계산 병목 현상($O(mn)$)을 겪습니다. 엔트로피 정규화(entropic regularization)를 통해 싱크혼(Sinkhorn) 알고리즘을 사용할 수 있지만, 여전히 이러한 스케일링 문제를 유지하며 편향(bias)을 유발합니다. 이를 완화하기 위해, **저계수 최적 운송(Low-Rank Optimal Transport, LOT)**은 운송 계획을 저차원 구조로 제한하며, 결합 행렬 $\Gamma$를 잠재 마진(latent marginal) $g$를 공유하는 하위 결합들 $(U, V)$로 인수분해합니다.

저계수 구조의 효율성에도 불구하고, 기존 솔버(예: LOT, LR-GW, UB-LOT)는 1차 미러 경사 하강법(first-order mirror-descent) 업데이트와 반복적인 내부 루프(Dykstra 또는 Bregman 투영)에 크게 의존합니다. 이러한 접근 방식은 다음과 같은 몇 가지 한계에 직면합니다:

1. 하이퍼파라미터 민감도: 스텝 사이즈(step-size)와 엔트로피 정규화 파라미터에 대한 세심한 튜닝이 필요합니다.
2. 곡률 무시: 국소적 곡률(local curvature)을 활용하지 못하여 수렴 속도가 느려집니다.
3. 계산 오버헤드: 투영 및 리트랙션(retraction)을 위한 내부 반복 루프가 특히 불균형(unbalanced) 설정에서 반복당 복잡도를 증가시킵니다.

방법론

저자들은 계수 $r$인 양의 인수 분해된 결합(positive factored couplings)의 집합을 양의 직교 영역(positive orthant) 내의 매끄러운 임베디드 부분 다양체(smooth embedded submanifolds)로 모델링하는 통합 리만 기하학 프레임워크를 제안합니다. 이 다양체들에 피셔-라우(Fisher-Rao) 곱 제품 메트릭을 부여함으로써, 저자들은 견고한 1차 및 2차 최적화를 가능하게 하는 다루기 쉬운 기하학적 구성 요소들을 도출합니다.

1. 다양체 공식화

프레임워크는 마진 제약 조건에 따라 두 가지 별개의 다양체를 정의합니다:

• 균형 다양체 ($\mathcal{M}_{bal}$): 소스 및 타겟 마진의 엄격한 재구성($U\mathbf{1}=a, V\mathbf{1}=b$)과 잠재 일관성($U^\top\mathbf{1}=V^\top\mathbf{1}$)을 강제합니다.
• 불균형 다양체 ($\mathcal{M}_{ub}$): 외부 마진 제약을 완화하며, 목적 함수 내에서 KL 발산을 통해 마진 편차를 페널티로 부과하고 잠재 일관성($U^\top\mathbf{1}=V^\top\mathbf{1}$)만을 요구합니다.

2. 기하학적 구성 요소

피셔-라우 메트릭 하에서, 저자들은 다음의 명시적인 표현식을 도출합니다:

• 직교 투영(Orthogonal Projectors):
  • $\mathcal{M}_{bal}$의 경우, 투영은 공액 경사법(Conjugate Gradients, CG)을 통한 슈르 보충 행렬(Schur complement)로 축소된 KKT 시스템을 푸는 과정을 포함합니다.
  • $\mathcal{M}_{ub}$의 경우, 투영은 **폐쇄형 표현식(closed-form expression)**으로 단순화되어 반복적인 솔버를 제거합니다.
• 리트랙션(Retractions):
  • 리트랙션은 탄젠트 벡터를 피셔 일관적인 지수 미러 단계(exponential mirror step)를 거친 후 KL 투영을 통해 다양체로 다시 매핑합니다.
  • $\mathcal{M}_{bal}$의 경우, 이는 순환적 브레그만 투영(cyclic Bregman projections)을 필요로 합니다(반복적).
  • $\mathcal{M}_{ub}$의 경우, KL 투영이 **폐쇄형 열 스케일링(closed-form column scaling)**으로 붕괴되어 내부 루프를 완전히 제거합니다.
• 리만 기울기(Riemannian Gradient) 및 헤시안(Hessian):
  • 프레임워크는 비용 함수(선형 OT, Gromov-Wasserstein 등)의 유클리드 미분과 다양체 기하학에만 의존하는 비용 불가지론적(cost-agnostic) 공식을 제공합니다.
  • 이를 통해 곡률 정보를 활용하는 2차 솔버(Trust-Region, Conjugate Gradient)를 사용할 수 있습니다.

3. 계수 충분성 인증(Rank-Sufficiency Certificate)

논문은 전역 최적성 검증을 위한 인증서를 도입합니다. 계수 $r$인 임계점(critical point)에서 "축소된 기울기"(dual slack matrix)를 계산함으로써, 이 방법은 해당 솔루션이 전체 계수 문제에 대해 전역적으로 최적인지 검증할 수 있습니다. 만약 계수가 불충분하다면, 인증서는 계수 증가를 위한 방향을 식별합니다.

주요 기여

1. 통합 리만 공식화: 피셔-라우 메트릭 하에서 균형 및 불균형 저계수 OT를 매끄러운 다양체로 정형화하여, 선형 OT, Gromov-Wasserstein(GW), 그리고 이들의 불균형 변형을 아우르는 단일 기하학적 프레임워크를 제공합니다.
2. 폐쇄형 불균형 연산: 불균형 OT의 경우, 탄젠트 공간 투영기와 리트랙션이 폐쇄형 표현식으로 축소됨을 입증했습니다. 이는 UB-LOT과 같은 최신 방법들이 요구하는 내부 Dykstra/Bregman 루프를 제거하여, 반복당 복잡도를 $O(rN)$으로 줄입니다.
3. 2차 솔버: 곡률 인지 솔버(CG, Trust-Region)를 가능하게 하여, 엔트로피 정규화나 스텝 사이즈 튜닝 없이도 미러 경사 하강법보다 빠른 수렴을 달세대합니다.
4. 전역 최적성 검증: 저계수 솔루션이 전역적으로 최적인지 확인하거나 계수 증가 전략을 가이드할 수 있는 계산 가능한 인증서를 제공합니다.

실험 결과

최대 $N=50,000$의 데이터셋 크기에 대해 균형 및 불균형 설정(Linear OT, GW, Fused GW) 전반에 걸친 광범위한 실험을 통해 다음을 입증했습니다:

• 수렴 속도: 리만 솔버는 미러 경사 하강법 베이스라인에 비해 현저히 적은 횟수의 외부 반복(outer iterations)만으로 수렴합니다.
• 실제 실행 시간(Wall-Clock Time): 불균형 설정에서의 내부 루프 제거는 상당한 속도 향상을 가져옵니다. 예를 들어, 이상치(outliers)가 있는 불균형 OT 환경에서 리만 솔버는 몇 초 만에 수렴하는 반면, UB-LOT은 시간이 초과되거나 발산합니다.
• 성능: 정규화된 베이스라인에 내재된 엔트로피 편향 없이 우수하거나 대등한 목적 함수 값을 달성합니다.
• 2차의 이점: 비볼록(non-convex) 문제(GW, Fused GW)에서 Trust-Region 방법은 헤시안 정보를 활용하여 얕은 로컬 미니마를 탈출함으로써 1차 방법보다 뛰어난 성능을 보입니다.

의의

본 논문은 리만 기하학적 관점으로 전환함으로써 저계수 OT 솔버가 일조 조건(ill-conditioning)과 하이퍼파라미터 민감도 문제를 극복할 수 있다고 주장합니다. 핵심적인 의의는 불균형 연산의 구조적 단순화에 있으며, 이는 계산 비용이 많이 드는 반복 프로세스를 폐쇄형 평가로 변환합니다. 이를 통해 데이터 크기에 따라 선형적으로 확장되면서도 엄격한 2차 수렴 보장과 전역 최적성 검증을 제공하는 정규화 및 튜닝이 필요 없는 최적화를 가능하게 합니다. 이 프레임워크는 광범한 최적 운송 문제에 적용 가능한 일반적이고 비용 불가지론적인 접근 방식으로 제시됩니다.