Fixed-Point Neural Optimal Transport without Implicit Differentiation

Technische Samenvatting: Vaste-Punt Neuraal Optimale Transport zonder Impliciete Differentiatie

Probleemstelling
Optimaal Transport (OT) zoekt een afbeelding tussen kansverdelingen die een voorgeschreven transportkost minimaliseert. Hoewel recente deep learning-aanpakken schaalbare OT-oplossers mogelijk hebben gemaakt, vertrouwen ze vaak op adversariele min-max optimalisatie (bijv. WGANs) of multi-netwerkarchitecturen die transportafbeeldingen en duale potentialen omvatten. Deze methoden lijden vaak aan trainingsinstabiliteit, gevoeligheid voor hyperparameters en convergentieproblemen, vooral in hoogdimensionale settings. Alternatieve methoden gebaseerd op continue stromen vereisen het oplossen van gewone differentiaalvergelijkingen (ODE's), wat aanzienlijke rekenkundige overhead met zich meebrengt, terwijl Hamilton-Jacobi-Bellman (HJB)-formuleringen moeite hebben met schaalbaarheid in hoge dimensies. Bovendien vereisen veel bestaande methoden hulpnetwerken of afwisselende optimalisatieprocedures om duale haalbaarheid af te dwingen.

Methodologie
De auteurs stellen een enkel-netwerkframework voor voor kwadratisch optimaal transport dat adversariele training en impliciete differentiatie elimineert. De kern van de methode is een impliciete neurale formulering afgeleid van het Kantorovich-duale probleem.

1. Enkele Potentialparametrisatie: Het duale potentieel $\phi$ wordt geparametriserd als het negatieve van een convexe functie $g_\theta$ die wordt weergegeven door een enkel neurale netwerk ($\phi_\theta(x) = -g_\theta(x)$). Dit maakt gebruik van Breniers stelling, die stelt dat de optimale transportafbeelding de gradiënt is van een convex potentieel.
2. Proximale Vaste-Punt Herformulering: Onder deze parametrisatie reduceert de $c$-transformatie (een kerncomponent van het duale doel) tot een sterk convex proximale optimalisatieprobleem:
   $$(-g_\theta)^c(z) = \inf_{y \in \mathbb{R}^d} \left\{ \frac{1}{2}\|y - z\|^2 + g_\theta(y) \right\}$$
   De minimalizer $y^*_\theta(z)$ van dit probleem definieert de achterwaartse transportafbeelding ($T_{\nu \to \mu}(z) = y^*_\theta(z)$), terwijl de voorwaartse afbeelding gegeven wordt door $T_{\mu \to \nu}(x) = x + \nabla g_\theta(x)$.
3. Training via Vaste-Punt Iteratie: Het trainingsdoel omvat het minimaliseren van het negatieve duale verlies. Cruciaal wordt de innerlijke minimalisatie (de proximale operator) opgelost met standaard vaste-puntiteraties (bijv. gradiëntafdaalstappen) in plaats van uitgerolde optimalisatie of impliciete differentiatie.
4. Gradiëntberekening zonder Impliciete Differentiatie: Een belangrijk theoretisch inzicht is dat de gradiënt van het verlies met betrekking tot de netwerkparameters $\theta$ geen backpropagatie door de vaste-puntiteraties vereist. Door de optimaliteitsvoorwaarden van de eerste orde van de proximale operator, heffen de termen die de afgeleide van het vaste punt met betrekking tot $\theta$ bevatten ($dy^*/d\theta$) exact op. Bijgevolg hangt de gradiënt alleen af van de expliciete afgeleiden van $g_\theta$ geëvalueerd bij de invoermonsters en de benaderde vaste-puntoplossingen. Dit maakt efficiënte training met standaard backpropagatie mogelijk.
5. Convergentie: De auteurs stellen vast dat Stochastische Gradiëntafdaal (SGD) convergeert naar een stationair punt, zelfs wanneer het innerlijke vaste-puntprobleem bij benadering wordt opgelost (binnen een gespecificeerde tolerantie), mits de benaderingsfout begrensd is.
6. Klasse-Conditionele Uitbreiding: Het framework breidt zich natuurlijk uit tot klasse-conditioneel OT door het potentieel $g_\theta(x, k)$ te conditioneren op klasselabels $k$, waardoor een enkel netwerk meerdere transporttaken gelijktijdig kan afhandelen.

Belangrijkste Bijdragen

• Enkel-netwerk Herformulering: De auteurs leiden een formulering van kwadratisch OT af die slechts één neurale netwerk vereist om zowel voorwaartse als achterwaartse transportafbeeldingen te herstellen, waardoor de behoefte aan adversariele min-max-opstellingen of hulpgeconjugeerde netwerken wordt verwijderd.
• Training zonder Impliciete Differentiatie: Zij tonen aan dat het doel kan worden gedifferentieerd zonder impliciete differentiatie, door te vertrouwen op de exacte opheffing van termen veroorzaakt door de optimaliteitsvoorwaarden van de proximale operator.
• Convergentiegaranties: Het artikel levert theoretische bewijzen voor de convergentie van SGD naar een stationair punt onder onnauwkeurige vaste-puntbeoordelingen, waarmee de praktische realiteit wordt aangepakt dat de innerlijke lus niet tot machineprecisie wordt opgelost.
• Unificerend Framework: De methode herstelt gelijktijdig voorwaartse en achterwaartse afbeeldingen en breidt zich natuurlijk uit tot klasse-conditionele settings zonder architecturale wijzigingen.

Experimentele Resultaten
De voorgestelde methode werd geëvalueerd op hoogdimensionale Gaussische benchmarks, realistische fysica-datasets (UCI) en beeldvertalingstaken (MNIST naar Fashion-MNIST).

• Hoogdimensionale Gaussische Verdelingen: Op Gaussische benchmarks met dimensies tot $d=64$ bereikte de methode een aanzienlijk lager Onverklaarde Variantiepercentage (UVP) in vergelijking met baselines zoals NOT, MM-v1, MM:R, NCF en WGAN-QC. Het behield stabiele nauwkeurigheid naarmate de dimensionaliteit toenam, terwijl baselines verslechtering vertoonden.
• Rekenkundige Efficiëntie: De methode vertoonde gunstige rekenkundige en geheugenefficiëntie, vereiste minder parameters en minder trainingtijd dan op continue stromen gebaseerde methoden (FFJORD, RNODE) en andere neurale OT-benaderingen.
• Realistische Data: Op diverse UCI-datasets (Power, Gas, HEPMASS, MINIBOONE) en BSDS300-beeldfragmenten bereikte de methode concurrerende of superieure prestaties in het afstemmen van verdelingen (gemeten door Maximum Mean Discrepancy) met aanzienlijk kleinere modelgroottes.
• Beeldvertaling: Bij klasse-conditionele beeldvertalingstaken bereikte de methode de laagste Fréchet Inception Distance (FID) en de hoogste classificienauwkeurigheid onder baselines, wat wijst op superieure behoud van klassestructuur en semantische inhoud.
• Ablatiestudies: Experimenten toonden aan dat de methode robuust is voor de stop-tolerantie van de vaste-puntoptoplosser; matige onnauwkeurigheid in de innerlijke lus levert een gunstige afweging op tussen nauwkeurigheid en snelheid. Bovendien, hoewel het afdwingen van strikte convexiteit via gewichtsprojectie de vaste-puntiteraties stabiliseerde, verslechterde dit de algehele transportnauwkeurigheid, wat suggereert dat een niet-convexe parametrisatie (zonder harde projectie) een betere balans biedt tussen expressiviteit en stabiliteit.

Betekenis en Beweringen
Het artikel beweert dat dit framework een stabiel, efficiënt en schaalbaar alternatief biedt voor bestaande neurale OT-methoden. Door adversariele training en de rekenkundige overhead van ODE-oplossers of het oplossen van HJB-vergelijkingen over het hele domein te elimineren, biedt de methode een "causaliteitsvrij" mechanisme voor het afdwingen van optimaliteitsvoorwaarden. De auteurs stellen dat hun aanpak breed toepasbaar is, in staat om hoogdimensionale en klasse-conditionele problemen te hanteren terwijl nauwkeurige transportafbeeldingen in beide richtingen worden hersteld. Het werk positioneert zichzelf als een stap in de richting van het vereenvoudigen van het optimalisatielandschap van OT, weg van complexe min-max-formuleringen naar een enkel-objectief minimalisatieprobleem dat theoretisch onderbouwd en praktisch efficiënt is.