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

Technische Samenvatting: Difference-of-Convex Programming in de Wasserstein-ruimte met Toepassingen op MMD-optimalisatie

Probleemstelling
Het optimaliseren van functionalen over de ruimte van kansverdelingen staat centraal bij machine learning-taken die variëren van variatieve inferentie tot generatieve modellering. Een standaardbenadering is het direct uitvoeren van optimalisatie over de Wasserstein-ruimte $(P_2(\mathbb{R}^d), W_2)$ door de discretisering van Wasserstein-gradiëntstromen. Hoewel dit effectief is voor convexe functionalen (bijv. KL-divergentie, f-divergenties), zijn veel praktisch relevante doelstellingen niet-convex langs Wasserstein-geodesieken. Prominente voorbeelden hiervan zijn de gekwadrateerde Wasserstein-afstand, de Sliced-Wasserstein-afstand en de Maximum Mean Discrepancy (MMD). Standaard eerste-orde methoden zoals Wasserstein Gradient Descent (WGD) worstelen in deze niet-convexe settings en vereisen vaak delicate ruisinjectie of optimalisatie in alternatieve ruimten (bijv. Wasserstein-MMD of Wasserstein-Fisher-Rao ruimten) die de gewichten van de distributies wijzigen of impliciete stappen mengen.

Methodologie
De auteurs stellen een nieuw deeltjesgebaseerd algoritme voor dat is afgestemd op niet-convexe functionalen die een Difference-of-Convex (DC) decompositie bezitten binnen de Wasserstein-ruimte. De kernmethodologie behelst het tillen van de klassieke Convex-Concave Procedure (CCCP) vanuit Euclidische ruimten naar de Wasserstein-ruimte, wat resulteert in de Wasserstein Convex-Concave Procedure (WCCCP).

1. DC Decompositie in de Wasserstein-ruimte: Het doelfunctional $F$ wordt gedecomposeerd als $F(\mu) = F_+(\mu) - F_-(\mu)$, waarbij zowel $F_+$ als $F_-$ totaal convexe functionalen zijn op $P_2(\mathbb{R}^d)$. In tegenstelling tot eerder werk, dat de concave component beperkte tot potentiële energieën (lineaire functionalen), staat deze benadering toe dat $F_-$ elke differentiabele functionaal is.
2. Algoritmisch Schema: In elke iteratie $k$, gegeven een maat $\mu_k$, lost het algoritme een majorisatie-minimalisatieprobleem op. Het lineariseert de concave component $F_-$ rond $\mu_k$ om een bovengrens te construeren en minimaliseert de resulterende convexe surrogaat over de ruimte van transportkaarten $T \in L^2(\mu_k)$:
   $$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)$$
   De volgende iteratie wordt vervolgens gedefinieerd als $\mu_{k+1} = (T_{k+1})\#\mu_k$.
3. Implementatie: De innerlijke minimalisatie wordt uitgevoerd via gradiëntafdaling op $L^2(\mu_k)$. Voor empirische distributies (deeltjes) vertaalt dit zich naar het updaten van de posities van de deeltjes. De auteurs leggen ook verbanden tussen WCCCP en Mirror Descent/Bregman Proximal Descent in de Wasserstein-ruimte.
4. MMD Specificaties: Voor de Maximum Mean Discrepancy stelt de auteur expliciete DC-decompositie voor door de kernel $k(x,y) = \psi(x-y)$ zelf op te splitsen in convexe componenten $\psi = \psi_+ - \psi_-$. Zij bieden specifieke decomposities voor translatie-invariante kernels, inclusief Gaussische en Riesz-kernels, gebruikmakend van de Jordan-decompositie of algebraïsche splitsing (bijv. $\cosh/\sinh$).

Belangrijkste Bijdragen

• WCCCP Algoritme: Introductie van de Wasserstein Convex-Concave Procedure, die DC-optimalisatie uitbreidt naar de Wasserstein-ruimte zonder de concave component te beperken tot lineaire potentialen.
• Theoretische Convergentie:
  • Bewijs dat de doelwaarde $F(\mu_k)$ niet-toenemend is langs de iteraties.
  • Vaststelling van bijna-stationariteit resultaten. Onder aannames van gladheid en sterke convexiteit (totale convexiteit) van de componenten, bewijzen de auteurs sublineaire convergentiesnelheden voor de gekwadrateerde afstand tussen iteraties en de norm van de Wasserstein-gradiënt. Specifiek, $\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$.
  • Demonstratie dat het schema equivalent is aan Bregman Proximal Descent en Mirror Descent onder specifieke condities.
• Expliciete DC Decomposities voor MMD: Het artikel biedt een systematische manier om MMD-functionalen met translatie-invariante kernels in DC-vormen te decomponeren, waarbij aan de noodzakelijke convexiteits- en gladheidseisen voor convergentiegaranties wordt voldaan.
• Empirische Validatie: Experimenten met Energy Distance en MMD met Gaussische kernels laten zien dat goed gekozen DC-decomposities aanzienlijk sneller en stabieler convergeren dan standaard WGD en de Wasserstein Proximal Gradient (Forward-Backward) methode.

Resultaten

• Energy Distance: Op synthetische data (spiraal- en katvormen) en CIFAR10-beelden convergeert WCCCP aanzienlijk sneller dan WGD en de Forward-Backward (FB) methode. Het algoritme convergeert succesvol de Energy Distance om complexe doelverdelingen te matchen.
• MMD met Gaussische Kernel: Het minimaliseren van MMD met een Gaussische kernel is berucht gevoelig voor bandbreedte en initialisatie. De auteurs tonen aan dat WGD en de FB-methode (gebruikmakend van een standaard decompositie) vaak vastlopen in lokale minima waar deeltjes wegdriften van het doel. In contrast hiermee bereikt WCCCP, gebruikmakend van specifieke kernel-decomposities (Jordan of $\cosh/\sinh$), een betere convergentie. De Jordan-decompositie wordt opgemerkt als zijnde gladder en minder gevoelig voor numerieke instabiliteit, terwijl de $\cosh/\sinh$-decompositie ervoor zorgt dat alle deeltjes uiteindelijk convergeren, zij het met een lagere snelheid.
• Vergelijking: De resultaten benadrukken dat de prestaties van WCCCP sterk afhankelijk zijn van de keuze van de DC-decompositie. Een "goed gekozen" decompositie is cruciaal om standaard gradiëntafdaling te overtreffen.

Betekenis en Claims
Het artikel claimt een robuust kader te bieden voor het optimaliseren van niet-convexe functionalen in de Wasserstein-ruimte, waarmee een gat wordt gedicht waar standaard gradiëntmethoden falen door niet-convexiteit. Door de CCCP naar deze setting te tillen, bieden de auteurs een theoretisch gefundeerd alternatief dat de problemen met het mengen van gewichten van eerdere MMD-specifieke benaderingen vermijdt. Het werk benadrukt dat hoewel het algoritme algemeen is, het praktische succes sterk afhangt van de specifieke DC-decompositie van de kernel. De auteurs merken bescheiden op dat toekomstig werk nodig is om de impact van verschillende decomposities dieper te begrijpen en om automatische, adaptieve strategieën voor de selectie ervan te ontwerpen. Daarnaast gaat de theoretische analyse momenteel uit van exacte oplossingen voor de innerlijke optimalisatieproblemen, en het meenemen van inexacte solvers is geïdentificeerd als een noodzakelijke richting voor toekomstig onderzoek.