Fixed-Point Neural Optimal Transport without Implicit Differentiation

Sintesi Tecnica: Trasporto Ottimo Neurale a Punto Fisso senza Differenziazione Implicita

Enunciato del Problema
Il Trasporto Ottimo (OT) ricerca una mappatura tra distribuzioni di probabilità che minimizzi un costo di trasporto prescritto. Sebbene recenti approcci di deep learning abbiano reso possibili solutori OT scalabili, essi fanno spesso affidamento su ottimizzazioni min–max avversarie (ad esempio, WGAN) o su architetture multi-rete che coinvolgono mappe di trasporto e potenziali duali. Questi metodi soffrono spesso di instabilità nell'addestramento, sensibilità agli iperparametri e difficoltà di convergenza, specialmente in contesti ad alta dimensionalità. Approcci alternativi basati su flussi continui richiedono la risoluzione di equazioni differenziali ordinarie (ODE), comportando un significativo sovraccarico computazionale, mentre le formulazioni Hamilton–Jacobi–Bellman (HJB) faticano a scalare in dimensioni elevate. Inoltre, molti metodi esistenti richiedono reti ausiliarie o procedure di ottimizzazione alternata per imporre la fattibilità duale.

Metodologia
Gli autori propongono un framework a rete singola per il trasporto ottimo quadratico che elimina l'addestramento avversario e la differenziazione implicita. Il cuore del metodo è una formulazione neurale implicita derivata dal problema duale di Kantorovich.

1. Parametrizzazione del Potenziale Singolo: Il potenziale duale $\phi$ è parametrizzato come il negativo di una funzione convessa $g_\theta$ rappresentata da una singola rete neurale ($\phi_\theta(x) = -g_\theta(x)$). Questo sfrutta il teorema di Brenier, che afferma che la mappa di trasporto ottima è il gradiente di un potenziale convesso.
2. Riformulazione a Punto Fisso Prossimale: Sotto questa parametrizzazione, la $c$-trasformata (un componente chiave dell'obiettivo duale) si riduce a un problema di ottimizzazione prossimale fortemente convesso:
   $$(-g_\theta)^c(z) = \inf_{y \in \mathbb{R}^d} \left\{ \frac{1}{2}\|y - z\|^2 + g_\theta(y) \right\}$$
   Il minimizzatore $y^*_\theta(z)$ di questo problema definisce la mappa di trasporto inversa ($T_{\nu \to \mu}(z) = y^*_\theta(z)$), mentre la mappa diretta è data da $T_{\mu \to \nu}(x) = x + \nabla g_\theta(x)$.
3. Addestramento tramite Iterazione a Punto Fisso: L'obiettivo di addestramento consiste nel minimizzare la perdita duale negativa. Crucialmente, la minimizzazione interna (l'operatore prossimale) viene risolta utilizzando iterazioni standard a punto fisso (ad esempio, passi di discesa del gradiente) piuttosto che ottimizzazione srotolata o differenziazione implicita.
4. Calcolo del Gradiente senza Differenziazione Implicita: Un'insight teorico chiave è che il gradiente della perdita rispetto ai parametri della rete $\theta$ non richiede la retropropagazione attraverso le iterazioni a punto fisso. A causa delle condizioni di ottimalità del primo ordine dell'operatore prossimale, i termini che coinvolgono la derivata del punto fisso rispetto a $\theta$ ($dy^*/d\theta$) si cancellano esattamente. Di conseguenza, il gradiente dipende solo dalle derivate esplicite di $g_\theta$ valutate sui campioni di input e sulle soluzioni approssimate a punto fisso. Ciò consente un addestramento efficiente utilizzando la retropropagazione standard.
5. Convergenza: Gli autori dimostrano che la Discesa Stocastica del Gradiente (SGD) converge a un punto stazionario anche quando il problema interno a punto fisso viene risolto in modo approssimato (entro una tolleranza specificata), purché l'errore di approssimazione sia limitato.
6. Estensione Condizionata alla Classe: Il framework si estende naturalmente all'OT condizionato alla classe condizionando il potenziale $g_\theta(x, k)$ sulle etichette di classe $k$, permettendo a una singola rete di gestire simultaneamente più compiti di trasporto.

Contributi Chiave

• Riformulazione a Rete Singola: Gli autori derivano una formulazione dell'OT quadratico che richiede una sola rete neurale per recuperare sia le mappe di trasporto dirette che inverse, eliminando la necessità di configurazioni min–max avversarie o reti coniugate ausiliarie.
• Addestramento senza Differenziazione Implicita: Dimostrano che l'obiettivo può essere differenziato senza differenziazione implicita, affidandosi alla cancellazione esatta dei termini indotta dalle condizioni di ottimalità dell'operatore prossimale.
• Garanzie di Convergenza: Il lavoro fornisce prove teoriche per la convergenza della SGD a un punto stazionario in presenza di valutazioni a punto fisso inesatte, affrontando la realtà pratica secondo cui il ciclo interno non viene risolto con precisione di macchina.
• Framework Unificato: Il metodo recupera simultaneamente le mappe dirette e inverse e si estende naturalmente a contesti condizionati alla classe senza modifiche architetturali.

Risultati Sperimentali
Il metodo proposto è stato valutato su benchmark gaussiani ad alta dimensionalità, dataset fisici reali (UCI) e compiti di traduzione di immagini (da MNIST a Fashion-MNIST).

• Gaussiane ad Alta Dimensionalità: Su benchmark gaussiani con dimensioni fino a $d=64$, il metodo ha ottenuto una Percentuale di Varianza Spiegata (UVP) significativamente inferiore rispetto a baseline come NOT, MM-v1, MM:R, NCF e WGAN-QC. Ha mantenuto un'accuratezza stabile all'aumentare della dimensionalità, mentre le baseline hanno mostrato un degrado.
• Efficienza Computazionale: Il metodo ha dimostrato un'efficienza computazionale e di memoria favorevole, richiedendo meno parametri e meno tempo di addestramento rispetto ai metodi basati su flussi continui (FFJORD, RNODE) e ad altri approcci neurali OT.
• Dati Reali: Su diversi dataset UCI (Power, Gas, HEPMASS, MINIBOONE) e patch di immagini BSDS300, il metodo ha ottenuto prestazioni di matching delle distribuzioni competitive o superiori (misurate tramite Discrepanza Media Massima) con dimensioni del modello sostanzialmente più piccole.
• Traduzione di Immagini: Nei compiti di traduzione di immagini condizionati alla classe, il metodo ha ottenuto la più bassa Distanza Fréchet Inception (FID) e la più alta accuratezza di classificazione tra le baseline, indicando una conservazione superiore della struttura di classe e del contenuto semantico.
• Studi di Ablazione: Gli esperimenti hanno mostrato che il metodo è robusto alla tolleranza di arresto del solutore a punto fisso; un'inexactezza moderata nel ciclo interno produce un compromesso favorevole tra accuratezza e velocità. Inoltre, sebbene l'imposizione di una convessità rigorosa tramite proiezione dei pesi stabilizzasse le iterazioni a punto fisso, ciò ha degradato l'accuratezza complessiva del trasporto, suggerendo che una parametrizzazione non convessa (senza proiezione rigida) offre un migliore equilibrio tra espressività e stabilità.

Significato e Affermazioni
Il lavoro afferma che questo framework offre un'alternativa stabile, efficiente e scalabile ai metodi neurali OT esistenti. Eliminando l'addestramento avversario e il sovraccarico computazionale dei solutori ODE o della risoluzione di equazioni HJB su tutto il dominio, il metodo fornisce un meccanismo "senza causalità" per imporre le condizioni di ottimalità. Gli autori affermano che il loro approccio è ampiamente applicabile, capace di gestire problemi ad alta dimensionalità e condizionati alla classe, recuperando al contempo mappe di trasporto accurate in entrambe le direzioni. Il lavoro si posiziona come un passo verso la semplificazione del paesaggio di ottimizzazione dell'OT, allontanandosi da formulazioni min–max complesse verso un problema di minimizzazione a obiettivo singolo che è teoricamente fondato e praticamente efficiente.