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

Sintesi Tecnica: Differenza di Programmazione Convessa nello Spazio di Wasserstein con Applicazioni all'Ottimizzazione MMD

Enunciato del Problema
L'ottimizzazione di funzionali sugli spazi di misure è centrale per compiti di apprendimento automatico che spaziano dall'inferenza variazionale alla modellazione generativa. Un approccio standard consiste nell'effettuare l'ottimizzazione direttamente sullo spazio di Wasserstein $(P_2(\mathbb{R}^d), W_2)$ discretizzando i flussi di gradiente di Wasserstein. Sebbene efficace per funzionali convessi (es. divergenza KL, divergenze f), molti obiettivi praticamente rilevanti sono non convessi lungo le geodetiche di Wasserstein. Esempi notevoli includono la distanza di Wasserstein al quadrato, la distanza di Sliced-Wasserstein e la Massima Discrepanza Media (MMD). I metodi standard del primo ordine come la Discesa del Gradiente di Wasserstein (WGD) faticano in questi contesti non convessi, richiedendo spesso l'iniezione delicata di rumore o l'ottimizzazione in spazi alternativi (es. spazi Wasserstein-MMD o Wasserstein-Fisher-Rao) che alterano i pesi delle distribuzioni o mescolano passaggi impliciti.

Metodologia
Gli autori propongono un nuovo algoritmo basato su particelle adattato per funzionali non convessi che ammettono una decomposizione di Differenza di Convessi (DC) all'interno dello spazio di Wasserstein. La metodologia centrale consiste nel sollevare il classico Procedimento Convesso-Concavo (CCCP) dagli spazi euclidei allo spazio di Wasserstein, dando origine al Procedimento Convesso-Concavo di Wasserstein (WCCCP).

1. Decomposizione DC nello Spazio di Wasserstein: Il funzionale obiettivo $F$ è decomposto come $F(\mu) = F_+(\mu) - F_-(\mu)$, dove entrambi $F_+$ e $F_-$ sono funzionali totalmente convessi in $P_2(\mathbb{R}^d)$. A differenza dei lavori precedenti che limitavano la parte concava a potenziali energetici (funzionali lineari), questo approccio permette a $F_-$ di essere qualsiasi funzionale differenziabile.
2. Schema Algoritmico: Ad ogni iterazione $k$, data una misura $\mu_k$, l'algoritmo risolve un problema di majorazione-minimizzazione, linearizzando la parte concava $F_-$ attorno a $\mu_k$ per costruire un limite superiore e minimizzando il risultante surrogato convesso sullo spazio delle mappe di trasporto $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)$$
   L'iterata successiva è poi definita come $\mu_{k+1} = (T_{k+1})\#\mu_k$.
3. Implementazione: La minimizzazione interna viene eseguita tramite discesa del gradiente su $L^2(\mu_k)$. Per distribuzioni empiriche (particelle), ciò si traduce nell'aggiornamento delle posizioni delle particelle. Gli autori stabiliscono inoltre connessioni tra WCCCP e Discesa di Mirror Descent/Bregman Proximal nel contesto di Wasserstein.
4. Specifiche MMD: Per la Massima Discrepanza Media, gli autori propongono esplicite decomposizioni DC scindendo il kernel stesso $k(x,y) = \psi(x-y)$ in componenti convesse $\psi = \psi_+ - \psi_-$. Forniscono decomposizioni specifiche per kernel invarianti per traslazione, inclusi i kernel Gaussiani e di Riesz, utilizzando la decomposizione di Jordan o lo splitting algebrico (es. $\cosh/\sinh$).

Contributi Chiave

• Algoritmo WCCCP: Introduzione del Procedimento Convesso-Concavo di Wasserstein, estendendo l'ottimizzazione DC allo spazio di Wasserstein senza limitare la componente concava a potenziali lineari.
• Convergenza Teorica:
  • Dimostrazione che il valore del funzionale $F(\mu_k)$ è non crescente lungo le iterazioni.
  • Stabilimento di risultati di quasi stazionarietà. Sotto ipotesi di regolarità (smoothness) e convessità forte (totalità della convessità) delle componenti, gli autori dimostrano tassi di convergenza sublineari per la distanza al quadrato tra le iterazioni e la norma del gradiente di Wasserstein. Nello specifico, $\min_{k} \|\nabla_{W_2}F(\mu_k)\|^2 = O(1/K)$.
  • Dimostrazione che lo schema è equivalente alla Discesa Bregman Proximal e alla Discesa di Mirror sotto specifiche condizioni.
• Decomposizioni DC Esplicite per MMD: Il documento fornisce un modo sistematico per scomporre i funzionali MMD con kernel invarianti per traslazione in forme DC, soddisfacendo le necessarie condizioni di convessità e regolarità per le garanzie di convergenza.
• Validazione Empirica: Gli esperimenti sulla Distanza di Energia e sulla MMD con kernel Gaussiani dimostrano che decomposizioni DC ben scelte producono una convergenza più veloce e stabile rispetto alla standard WGD e al metodo di Gradiente Prossimale di Wasserstein (Forward-Backward).

Risultati

• Distanza di Energia: Su dati sintetici (forme a spirale e a categoria) e immagini CIFAR10, WCCCP converge significativamente più velocemente di WGD e dello schema Forward-Backward (FB). L'algoritmo minimizza con successo la Distanza di Energia per far corrispondere distribuzioni target complesse.
• MMD con Kernel Gaussiano: Minimizzare la MMD con un kernel Gaussiano è notoriamente sensibile alla larghezza di banda (bandwidth) e all'inizializzazione. Gli autori mostrano che WGD e lo schema FB (utilizzando una decomposizione standard) spesso rimangono intrappolati in minimi locali dove le particelle si allontanano dal target. Al contrario, WCCCP utilizzando specifiche decomposizioni di kernel (Jordan o $\cosh/\sinh$) ottiene una migliore convergenza. La decomposizione di Jordan è notata per essere più fluida e meno soggetta a instabilità numerica, mentre la decomposizione $\cosh/\sinh$ assicura che tutte le particelle convergano infine, sebbene a un ritmo più lento.
• Confronto: I risultati evidenziano come le prestazioni di WCCCP dipendano fortemente dalla scelta della decomposizione DC. Una decomposizione "ben scelta" è critica per superare la discesa del gradiente standard.

Significatività e Rivendicazioni
Il documento sostiene di fornire un framework robusto per l'ottimizzazione di funzionali non convessi nello spazio di Wasserstein, affrontando una lacuna dove i metodi di gradiente standard falliscono a causa della non convessità. Sollevando il CCCP in questo contesto, gli autori offrono un'alternativa teoricamente fondata che evita i problemi di miscelazione dei pesi dei precedenti approcci specifici per MMD. Il lavoro sottolinea che, sebbene l'algoritmo sia generale, il suo successo pratico dipende dalla specifica decomposizione DC del kernel. Gli autori notano con modestia che sono necessari ulteriori lavori per comprendere più profondamente l'impatto di diverse decomposizioni e per progettare strategie automatiche e adattive per la loro selezione. Inoltre, l'analisi teorica attuale assume soluzioni esatte dei problemi di ottimizzazione interna, e la presa in conto di solver imprecisi è identificata come una direzione necessaria per la ricerca futura.