A convergence rate for the entropic JKO scheme

1. Probleemstelling en Context

Het artikel richt zich op de numerieke benadering van partiële differentiaalvergelijkingen (PDE's) die optreden als gradiëntstromen in de ruimte van kansmaten, uitgerust met de Wasserstein-metriek. De klassieke methode hiervoor is het JKO-scheme (genoemd naar Jordan, Kinderlehrer en Otto), dat een impliciete Euler-discretisatie gebruikt om de evolutie van een dichtheid $\rho$ te berekenen door een variatieprobleem op te lossen dat de Wasserstein-afstand $W_2$ en een functionaal $F$ combineert.

Het centrale probleem:
Hoewel het JKO-scheme theoretisch elegant is, is het numeriek zeer kostbaar omdat het berekenen van de exacte Wasserstein-afstand $W_2$ complex is (het vereist het oplossen van een optimalisatieprobleem voor transportplannen).

De oplossing en de uitdaging:
Een veelgebruikte strategie is het vervangen van de Wasserstein-afstand door een entropische regularisatie, bekend als de Schrödinger-kost ($Sch_{\alpha\tau}$). Dit maakt gebruik van de efficiënte Sinkhorn-algoritme voor berekening. Echter, deze regularisatie introduceert een verstoring in het systeem.

• Als de regularisatieparameter $\varepsilon$ evenredig is met de tijdstap $\tau$ (d.w.z. $\varepsilon = \alpha\tau$), convergeert het entropische JKO-scheme niet naar de oorspronkelijke PDE, maar naar een gewijzigde PDE met een extra lineaire diffusieterm: $\frac{\alpha}{2}\Delta\rho$.
• Bestaande literatuur (bijv. [4]) heeft laten zien dat deze convergentie plaatsvindt, maar geen expliciete convergentie-snelheden heeft geleverd die de afstand kwantificeren tussen het entropische scheme en de oplossing van de oorspronkelijke PDE, noch tussen het klassieke en het entropische scheme.

Het doel van dit artikel is om onder convexiteitsaannames een expliciete convergentiesnelheid af te leiden voor de afstand tussen het klassieke JKO-scheme en het entropische JKO-scheme, en vervolgens een schatting te geven voor de afstand tot de oplossing van de initiële PDE.

2. Methodologie

De auteurs gebruiken een combinatie van variatiële analyse, optimalisatietheorie in de ruimte van kansmaten en discrete dynamische systemen.

Belangrijke aannames (Hypothese 1.1):
De functionaal $F$ moet voldoen aan:

1. Laagste semicontinuiteit ten opzichte van zwakke convergentie in $P_2(\mathbb{R}^d)$.
2. $\lambda$-convexiteit langs gegeneraliseerde geodeten (een versterkte vorm van geodetische convexiteit, essentieel voor stabiliteit).
3. Een groei-aanname gerelateerd aan de warmtestroom (convolutie met de hittekernel $\sigma_t$), wat een schatting garandeert op de verandering van $F$ onder diffusie.

Kernstrategie:
De bewijstechniek is gebaseerd op een iteratieve vergelijking van de iteraties van beide schema's.

1. Decompositie: De afstand tussen twee stappen ($k+1$) wordt opgesplitst in twee delen via de driehoeksongelijkheid:
   • Term (I): De afstand tussen twee iteraties van het klassieke JKO-scheme, startend vanuit verschillende meetpunten. Dit wordt beheerst door de contractie-eigenschap van het klassieke JKO-scheme (gebaseerd op de Discrete Evolution Variational Inequality - EVI).
   • Term (II): De afstand tussen één stap van het klassieke en één stap van het entropische scheme, startend vanuit hetzelfde punt. Dit is de kern van het nieuwe bewijs.
2. Competitor-construction: Om Term (II) te schatten, construeren de auteurs een "goede competitor" voor het entropische probleem. Ze gebruiken het idee dat de stroom van $F + \frac{\alpha}{2}H$ (waarbij $H$ de Boltzmann-entropie is) asymptotisch equivalent is aan eerst de stroom van $F$ te volgen en daarna de stroom van $\frac{\alpha}{2}H$ (de warmtestroom). Ze gebruiken dus $J^0_\tau(\mu) * \sigma_{\alpha\tau}$ als competitor.
3. Discrete Gronwall-lemma: De recursieve ongelijkheden die uit de vergelijking van de stappen voortvloeien, worden opgelost met een variant van het discrete Gronwall-lemma om een globale foutbound voor $n$ stappen te verkrijgen.

3. Belangrijkste Resultaten

Hoofdstelling (Theorem 1.3):
De auteurs leiden een expliciete bovengrens af voor de kwadratische Wasserstein-afstand tussen de iteraties van het klassieke JKO-scheme ($J^0_{n,\tau}$) en het entropische JKO-scheme ($J^\alpha_{n,\tau}$) na $n$ stappen met tijdstap $\tau$.

• Voor $\lambda = 0$:
  $$W_2(J^0_{n,\tau}(\mu_0), J^\alpha_{n,\tau}(\mu_0)) \leq \sqrt{2\tau(F(\mu_0) - F(J^0_{n,\tau}(\mu_0)))} + \sqrt{n\tau\alpha(H(\mu_0) - H(J^\alpha_{n,\tau}(\mu_0)) + Kn\tau)}$$
• Voor $\lambda \neq 0$: Een soortgelijke, maar complexere ongelijkheid die exponentiële factoren bevat afhankelijk van $\lambda$ en $\tau$.

Deze schattingen tonen aan dat de afstand afhangt van $\sqrt{\tau}$ (de discretisatiefout) en $\sqrt{\alpha}$ (de regularisatiefout).

Corollary 1.3.1 (Convergentie naar de PDE-oplossing):
Door de bovengenoemde schatting te combineren met bekende convergentiesnelheden van het klassieke JKO-scheme naar de gradiëntstroom $\rho_0(t)$, leiden de auteurs af dat:
$$W_2(J^\alpha_{n,t/n}(\mu_0), \rho_0(t)) \leq C \left( \sqrt{\alpha} + \sqrt{\frac{1}{n}} \right)$$
Dit betekent dat de totale fout van het entropische scheme lineair afhangt van de wortel van de regularisatieparameter $\alpha$ en de wortel van het aantal stappen (of tijdstap).

Optimaliteit (Sectie 5):
De auteurs onderzoeken of deze schattingen scherp zijn.

• Ze tonen aan dat de schattingen in de continue limiet (Theorem 1.4) scherp zijn voor specifieke modellen (Gaussische verdelingen met een kwadratische potentiaal).
• Ze bewijzen dat hun discrete schattingen (Theorem 5.6) in de limiet $\tau \to 0$ overeenkomen met de continue scherpste grenzen, tot op een term die naar nul gaat. Dit bevestigt dat de $\sqrt{\alpha}$-afhankelijkheid fundamenteel is en niet een artefact van hun bewijstechniek.

4. Significatie en Bijdrage

1. Eerste expliciete snelheid: Dit is een van de eerste werken dat een expliciete convergentiesnelheid biedt voor het entropische JKO-scheme in termen van zowel de tijdstap $\tau$ als de regularisatieparameter $\alpha$.
2. Brug tussen theorie en praktijk: Het werk rechtvaardigt het gebruik van entropische regularisatie (en de Sinkhorn-algoritme) in numerieke simulaties van gradiëntstromen. Het geeft gebruikers een kwantitatief inzicht in de fout die wordt geïntroduceerd door de regularisatie.
3. Technische innovatie: De auteurs omzeilen het gebrek aan een directe EVI-ongelijkheid voor het entropische schema door de stabiliteit van het klassieke schema te gebruiken en het entropische schema als een verstoring te behandelen. Dit is een krachtige nieuwe benadering.
4. Optimaliteitsanalyse: Door het gebruik van een "toy model" (Gaussische verdelingen) tonen ze aan dat hun theoretische grenzen bijna optimaal zijn, wat het vertrouwen in de schattingen versterkt.

Samenvattend biedt dit artikel een rigoureuze theoretische onderbouwing voor het gebruik van entropisch geregulariseerde JKO-schema's, met duidelijke foutmarges die essentieel zijn voor de implementatie en interpretatie van numerieke algoritmen in optimal transport en gerelateerde PDE's.