Difference of Convex Programming in the Wasserstein Space with Applications to MMD Optimization

기술 요약: 와서스타인 공간에서의 볼록 차이 프로그래밍과 MMD 최적화에 대한 응용

문제 정의
확률 측도 공간 상에서 범함수(functionals)를 최적화하는 것은 변분 추론(variational inference)에서 생성 모델링에 이르기까지 머신러닝 작업의 핵심입니다. 표준적인 접근 방식은 와서스타인 그래디언트 흐름(Wasserstein gradient flows)을 이산화하여 와서스타인 공간 $(P_2(\mathbb{R}^d), W_2)$ 위에서 직접 최적화를 수행하는 것입니다. 이는 볼록 범함수(예: KL 발산, f-발산)에 대해서는 효과적이지만, 많은 실질적으로 중요한 목적 함수들은 와서스타인 측지선(geodesics)을 따라 비볼록(non-convex)한 특성을 보입니다. 제곱 와서스타인 거리(squared Wasserstein distance), 슬라이스드 와서스타인 거리(Sliced-Wasserstein distance), 최대 평균 불일치(Maximum Mean Discrepancy, MMD)가 대표적인 예입니다. WGD(Wasserstein Gradient Descent)와 같은 표준 1차 방법론은 이러한 비볼록 설정에서 어려움을 겪으며, 종종 정교한 노이즈 주입을 요구하거나 분포의 가중치를 변경하거나 암시적 단계를 혼합하는 다른 공간(예: Wasserstein-MMD 또는 Wasserstein-Fisher-Rao 공간)에서의 최적화를 필요로 합니다.

방법론
저자들은 비볼록 범함수에 특화된 새로운 입자 기반 알고리즘을 제안합니다. 이 알고리즘은 와서스타인 공간 내에서 볼록 차이(Difference-of-Convex, DC) 분해가 가능한 범함수를 대상으로 합니다. 핵심 방법론은 고전적인 볼록-오목 절차(Convex-Concave Procedure, CCCP)를 와서스타인 공간으로 들어 올리는 것이며, 그 결과 **와서스타인 볼록-오목 절차(WCCCP)**가 도출됩니다.

1. 와서스타인 공간에서의 DC 분해: 목적 범함수 $F$를 $F(\mu) = F_+(\mu) - F_-(\mu)$로 분해하며, 여기서 $F_+$와 $F_-$는 모두 $P_2(\mathbb{R}^d)$ 상의 완전 볼록(totally convex) 범함수입니다. 오목한 부분(concave part)을 포텐셜 에너지(선형 범함수)로 제한했던 이전 연구들과 달리, 본 접근 방식은 $F_-$가 임의의 미분 가능한 범함수가 될 수 있도록 허용합니다.
2. 알고리즘 스킴: 각 반복 단계 $k$에서, 주어진 측도 $\mu_k$에 대해 알고리즘은 주요화-최소화(majorization-minimization) 문제를 해결합니다. 이는 $F_-$를 $\mu_k$ 주변에서 선형화하여 상한(upper bound)을 구성하고, $L^2(\mu_k)$ 상의 이동 맵(transport maps) $T \in L^2(\mu_k)$에 대해 다음과 같이 볼록 대리 함수(convex surrogate)를 최소화합니다:
   $$T_{k+1} = \arg\min_{T \in L^2(\mu_k)} \left( F_+(T\#\mu_k) - \langle \nabla_{W_2}F_-(\mu_k), T - \text{Id} \rangle_{L^2(\mu_k)} \right)$$
   다음 반복값은 $\mu_{k+1} = (T_{k+1})\#\mu_k$로 정의됩니다.
3. 구현: 내부 최소화 문제는 $L^2(\mu_k)$ 상의 그래디언트 하강법을 통해 수행됩니다. 입자(particles) 형태의 경험적 분포의 경우, 이는 입자 위치의 업데이트로 변환됩니다. 저자들은 WCCCP와 와서스타인 공간에서의 브레그만 프록시멀 디센트(Bregman Proximal Descent) 및 미러 디센트(Mirror Descent) 사이의 연결 고리를 확립합니다.
4. MMD 특이 사항: MMD를 위해, 저자들은 커널 $k(x,y) = \psi(x-y)$ 자체를 볼록 성분 $\psi = \psi_+ - \psi_-$로 분할함으로써 명시적인 DC 분해를 제안합니다. 이들은 조던 분해(Jordan decomposition)나 대수적 분할(예: $\cosh/\sinh$)을 사용하여 가우시안 및 리스(Riesz) 커널을 포함한 이동 불변 커널(translation-invariant kernels)에 대한 구체적인 분해를 제공합니다.

주요 기여

• WCCCP 알고리즘: 오목한 부분을 선형 포텐셜로 제한하지 않고 와서스타인 공간으로 DC 최적화를 확장한 와서스타인 볼록-오목 절차를 도입했습니다.
• 이론적 수렴성:
  • 반복 과정 동안 목적 값 $F(\mu_k)$가 감소함을 증명했습니다.
  • 거의 정체(almost stationarity) 결과를 확립했습니다. 성분들의 매끄러움(smoothness)과 강볼록성(strong convexity, total convexity) 가정 하에, 저자들은 반복 간의 거리 제곱과 와서스타인 그래디언트 노름 사이의 아선형(sublinear) 수렴율을 증명했습니다. 즉, $\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$입니다.
  • 이 스킴이 특정 조건 하에서 브레그만 프록시멀 디센트 및 미러 디센트와 동등함을 입증했습니다.
• MMD를 위한 명시적 DC 분해: 논문은 이동 불변 커널을 가진 MMD 범함수를 수렴 보장을 위한 DC 형태로 체계적으로 분해하는 방법을 제공합니다.
• 경험적 검증: 에너지 거리(Energy Distance) 및 가우시안 커널을 사용한 MMD 실험을 통해, 잘 선택된 DC 분해가 표준 WGD 및 와서스타인 프록시멀 그래디언트(Forward-Backward) 방법보다 더 빠른 그리고 안정적인 수렴을 가져온다는 것을 보여줍니다.

결과

• 에너지 거리: 합성 데이터(나선형 및 고양이 모양)와 CIFAR10 이미지에 대해, WCCCP는 WGD 및 Forward-Backward(FB) 스킴보다 현저히 빠르게 수렴합니다. 이 알고리즘은 복잡한 타겟 분포에 맞추기 위해 에너지 거리를 성공적으로 최소화합니다.
• 가우시안 커널을 사용한 MMD: 가우시안 커널을 사용한 MMD 최소화는 대역폭(bandwidth)과 초기화에 매우 민감합니다. 저자들은 WGD와 FB 스킴(표준 분해를 사용하는 경우)이 입자들이 타겟으로부터 멀어지는 국소 최솟값(local minima)에 갇히는 경우가 많음을 보여줍니다. 반면, 특정 커널 분해(조던 또는 $\cosh/\sinh$)를 사용한 WCCCP는 더 나은 수렴 성능을 달tes합니다. 조던 분해는 수치적 불안정성이 적고 더 매끄러운 것으로 언급되며, $\cosh/\sinh$ 분해는 모든 입자가 결국 수렴하도록 보장하지만 속도는 더 느립니다.
• 비교: 결과는 WCCCP의 성능이 DC 분해의 선택에 크게 의존한다는 점을 강조합니다. "잘 선택된" 분해는 표준 그래디언트 디센트를 능가하기 위해 매우 중요합니다.

의의 및 주장
본 논문은 표준 그래디언트 방법들이 비볼록성으로 인해 실패하는 지점을 해결함으로써, 와서스타인 공간에서 비볼록 범함수를 최적화하기 위한 견고한 프레임워크를 제공한다고 주장합니다. CCCP를 이 환경으로 들어 올림으로써, 저자들은 기존의 MMD 전용 접근 방식들이 가진 가중치 혼합 문제를 피하는 이론적으로 근거 있는 대안을 제시합니다. 이 연구는 알고리즘이 일반적이긴 하지만, 실질적인 성공은 커널의 구체적인 DC 분해에 달려 있음을 강조합니다. 저자들은 DC 분해의 영향력을 더 깊이 이해하고, 이를 위한 자동화된 적응형 전략을 설계하기 위해 향로 연구가 필요하다고 겸허히 언급합니다. 또한, 현재의 이론적 분석은 내부 최적화 문제의 정확한 해를 가정하고 있으며, 부정확한 솔버(inexact solvers)를 고려하는 것이 필요한 연구 방향임을 밝히고 있습니다.