Sintesi Tecnica: Un Approccio Riemanniano al Trasporto Ottimo a Basso Rango

Definizione del Problema

Il Trasporto Ottimo (OT) classico fornisce una metrica rigorosa per confrontare distribuzioni di probabilità, ma soffre di un collo di bottiglia computazionale quadratico ($O(mn)$) dovuto alla necessità di calcolare matrici di accoppiamento dense. Sebbene la regolarizzazione entropica permetta l'uso dell'algoritmo di Sinkhorn, essa mantiene questa scalabilità e introduce un bias. Per mitigare questo problema, il Trasporto Ottimo a Basso Rango (LOT) vincola il piano di trasporto a una struttura a bassa dimensionalità, fattorizzando la matrice di accoppiamento $\Gamma$ in sub-accoppiamenti $(U, V)$ che condividono una marginale latente $g$.

Nonostante l'efficienza strutturale delle formulazioni a basso rango, i solver esistenti (ad es., LOT, LR-GW, UB-LOT) si affidano pesantemente ad aggiornamenti di discesa del gradiente speculare (mirror-descent) del primo ordine accoppiati a loop interni iterativi (proiezioni di Dykstra o Bregman). Tali approcci affrontano diverse limitazioni:

1. Sensibilità agli Iperparametri: Richiedono una calibrazione attenta dei passi (step-size) e dei parametri di regolarizzazione entropica.
2. Curvatura Ignorata: Non riescono a sfruttare la curvatura locale del paesaggio di ottimizzazione, portando a una convergenza più lenta.
3. Overhead Computazionale: I loop iterativi interni per proiezione e ritrazione aumentano la complessità per iterazione, particolarmente nei casi sbilanciati (unbalanced).

Metodologia

Gli autori propongono un quadro geometrico Riemanniano unificato che modella l'insieme dei accoppiamenti fattorizzati positivi di rango-$r$ come sotto-varietà lisce (smooth submanifolds) immerse all'interno dell'ortante positivo. Equipaggiando queste varietà con il prodotto metrico di Fisher-Rao, gli autori derivano formulazioni trattabili per tutti i necessari ingredienti geometrici, consentendo un'ottimizzazione robusta del primo e secondo ordine.

1. Formulazione della Varietà (Manifold)

Il framework definisce due varietà distinte basate sui vincoli marginali:

• Varietà Bilanciata ($\mathcal{M}_{bal}$): Impone la ricostruzione rigorosa delle marginali sorgente e target ($U\mathbf{1}=a, V\mathbf{1}=b$) e la coerenza latente ($U^\top\mathbf{1}=V^\top\mathbf{1}$).
• Varietà Sbilanciata ($\mathcal{M}_{ub}$): Rilassa i vincoli marginali esterni, richiedendo solo la coerenza latente ($U^\top\mathbf{1}=V^\top\mathbf{1}$), con le deviazioni marginali penalizzate tramite divergenza KL nell'obiettivo.

2. Ingredienti Geometrici

Sotto la metrica Fisher-Rao, gli autori derivano espressioni esplicite per:

• Proiettori Ortogonali:
  • Per $\mathcal{M}_{bal}$, la proiezione comporta la risoluzione di un sistema KKT (ridotto a un complemento di Schur) tramite Gradiente Coniugato (CG).
  • Per $\mathcal{M}_{ub}$, la proiezione si semplifica in un'espressione in forma chiusa, eliminando la necessità di solver iterativi.
• Ritrazioni (Retractions):
  • La ritrazione mappa i vettori tangenti nuovamente sulla varietà tramite un passo speculare esponenziale coerente con Fisher seguito da una proiezione KL.
  • Per $\mathcal{M}_{bal}$, ciò richiede proiezioni Bregman cicliche (iterative).
  • Per $\mathcal{M}_{ub}$, la proiezione KL collassa in una scalatura di colonna in forma chiusa, eliminando interamente i loop interni.
• Gradiente e Hessiano Riemanniano:
  • Il framework fornisce una formulazione indipendente dal costo (cost-agnostic) in cui il gradiente e il prodotto Hessiano-vettore dipendono solo dalle derivate euclidee della specifica funzione di costo (OT Lineare, Gromov-Wasserstein, ecc.) e dalla geometria della varietà.
  • Ciò consente l'uso di solver del secondo ordine (Trust-Region, Gradiente Coniugato) che sfruttano l'informazione di curvatura.

3. Certificato di Sufficienza del Rango

Il documento introduce un certificato per la verifica dell'ottimalità globale. Calcolando un "gradiente ridotto" (matrice di slack duale) in un punto critico di rango-$r$, il metodo può verificare se la soluzione è globalmente ottimale per il problema a rango pieno. Se il rango è insufficiente, il certificato identifica una direzione per l'incremento del rango.

Contributi Chiave

1. Formulazione Riemanniana Unificata: Formalizza il trasporto ottimo a basso rango bilanciato e sbilanciato come varietà lisce sotto la metrica Fisher-Rao, fornendo un unico quadro geometrico per l'OT lineare, il Gromov-Wasserstein (GW) e le loro varianti sbilanciate.
2. Operazioni Sbilanciate in Forma Chiusa: Dimostra che per l'OT sbilanciato, i proiettori dello spazio tangente e le ritrazioni si riducono a espressioni in forma chiusa. Ciò elimina i loop interni di Dykstra/Bregman richiesti dai metodi allo stato dell'arte (es. UB-LOT), riducendo la complessità per iterazione a $O(rN)$.
3. Solver del Secondo Ordine: Abilita l'uso di solver sensibili alla curvatura (CG, Trust-Region) che raggiungono una convergenza più rapida rispetto alla discesa del gradiente speculare del primo ordine, senza richiedere regolarizzazione entropica o calibrazione del passo.
4. Verifica dell'Ottimalità Globale: Fornisce un certificato computabile per verificare se una soluzione a basso rango è globalmente ottimale o per guidare le strategie di incremento del rango.

Risultati Sperimentali

Esperimenti estesi in contesti bilanciati e sbilanciati (OT Lineare, GW, Fused GW) con dimensioni dei dataset fino a $N=50.000$ dimostrano:

• Velocità di Convergenza: I solver Riemanniani convergono in un ordine di grandezza in meno di iterazioni esterne rispetto ai baseline di mirror-descent.
• Tempo di Esecuzione (Wall-Clock Time): L'eliminazione dei loop interni nel setting sbilanciato porta a significativi aumenti di velocità. Ad esempio, nell'OT sbilanciato con outlier, i solver Riemanniani convergono in secondi, mentre UB-LOT va in timeout o diverge.
• Performance: I metodi raggiungono valori di obiettivo superiori o comparabili senza il bias entropico inerente ai baseline regolarizzati.
• Vantaggio del Secondo Ordine: Nei problemi non convessi (GW, Fused GW), i metodi Trust-Region sfruttano l'informazione dell'Hessiano per uscire da minimi locali superficiali, superando i metodi del primo ordine.

Significato

L'articolo sostiene che, passando da un paradigma di discesa del gradiente speculare del primo ordine a una prospettiva geometrica Riemanniana, i solver di OT a basso rango possono superare i limiti di ill-conditioning e sensibilità agli iperparametri. Il significato primario risiede nella semplificazione strutturale delle operazioni sbilanciate, che trasforma un processo iterativo computazionalmente costoso in una valutazione in forma chiusa. Ciò consente un'ottimizzazione priva di regolarizzazione e di calibrazione che scala linearmente con la dimensione del dataset, fornendo al contempo garanzie rigorose di convergenza del secondo ordine e verifica dell'ottimalità globale. Il framework è presentato come un approccio generale, indipendente dal costo, applicabile a un ampio spettro di problemi di trasporto ottimo.