A McKean-Pontrygin maximum principle for entropic-regularized optimal transport

1. Problema di Riferimento

Il lavoro si concentra sul problema del trasporto ottimo dinamico tra due distribuzioni di probabilità date, $\pi_0$ e $\pi_T$, definite su uno spazio $x \in \mathbb{R}^{d_x}$. L'obiettivo è trovare un controllo $U_t$ che guidi un sistema stocastico da $\pi_0$ a $\pi_T$ minimizzando un costo energetico.

Il sistema è governato dall'equazione differenziale stocastica (SDE) controllata:
$$d\tilde{X}_t = U_t dt + \sqrt{2}\Sigma^{1/2} d\tilde{B}_t$$
dove $\tilde{B}_t$ è un moto browniano. Il funzionale di costo da minimizzare è:
$$J(U) = \frac{1}{2} \int_0^T \mathbb{E}[\|U_t\|_R^2] dt$$
con $\|u\|_R^2 = u^T R^{-1} u$, dove $R$ è una matrice simmetrica definita positiva.

Il problema generalizza due casi noti:

• Trasporto Ottimo Classico: quando $\Sigma = 0$ e $R = I$.
• Ponte di Schrödinger: quando $\Sigma = R = I$ (corrisponde a un trasporto regolarizzato entropicamente).

L'obiettivo specifico dell'autore è formulare una soluzione mean-field (di campo medio) basata su un'estensione del principio del massimo di Pontryagin, noto come principio di McKean-Pontryagin.

2. Metodologia: Il Principio di McKean-Pontryagin

L'approccio proposto evita il campionamento diretto su percorsi stocastici, adottando invece una formulazione variazionale pura.

A. Formulazione Variazionale

L'autore introduce stati $X_t(a)$ e co-stati $P_t(a)$, dove $a \in \mathbb{R}^{d_x}$ sono etichette (o variabili indipendenti) che rimangono costanti nel tempo. La distribuzione di $X_t(a)$ è denotata da $\rho_t$.
Viene definito un funzionale d'azione $S$ che include:

• Termini cinetici e di controllo.
• Vincoli di accoppiamento tra stati e co-stati tramite potenziali $\psi_t$.
• Termini di penalizzazione per i vincoli ai bordi (distribuzioni iniziali e finali).

Il funzionale è costruito in modo tale che le condizioni al contorno siano soddisfatte in forma debole, garantendo che le densità marginali $\rho_0$ e $\rho_T$ coincidano con $\pi_0$ e $\pi_T$.

B. Equazioni del Moto

Applicando il calcolo delle variazioni al funzionale d'azione rispetto alle variabili $X_t, P_t, U_t, \beta_t, \psi_t$, si ottengono le equazioni di evoluzione Hamiltoniane:

1. Equazione di stato: $\dot{X}_t = U_t + \beta_t$
2. Equazione di co-stato: $-\dot{P}_t = \nabla_x(\Sigma : D^2_x \psi_t(X_t)) - D^2_x \psi_t(X_t) \beta_t$
3. Vincoli ottimali:
   • $P_t = \nabla_x \psi_t(X_t)$ (relazione tra co-stato e gradiente del potenziale)
   • $U_t = R \nabla_x \psi_t(X_t)$ (legge di controllo ottima)
   • Un'equazione di continuità per la densità $\rho_t$ che coinvolge un termine di drift $\beta_t$.

Un risultato cruciale è che il termine $\beta_t$ non è univocamente determinato dall'equazione variazionale. L'autore mostra che $\beta_t$ può essere scelto arbitrariamente (ad esempio $\beta_t = -\Sigma \nabla_x \log \rho_t$) senza influenzare la legge di controllo ottima $u_t(x) = R \nabla_x \psi_t(x)$. Questa libertà è una conseguenza della simmetria di ridenominazione (relabeling symmetry) della formulazione mean-field.

C. Connessione alle Equazioni di Hamilton-Jacobi-Bellman (HJB)

Il sistema derivato è coerente con le equazioni classiche del controllo ottimo stocastico:

• La densità $\rho_t$ soddisfa l'equazione di Fokker-Planck.
• Il potenziale $\psi_t$ soddisfa l'equazione di Hamilton-Jacobi-Bellman (HJB):
  $$-\partial_t \psi_t = \mathcal{L}_{u_t} \psi_t - \frac{1}{2} \|u_t\|_R^2$$
  dove $\mathcal{L}$ è il generatore del processo.

3. Contributi Chiave e Risultati

1. Unificazione Deterministica e Stocastica: La metodologia offre un trattamento unificato che copre sia il trasporto ottimo deterministico ($\Sigma=0$) sia i ponti di Schrödinger stocastici, tutto all'interno dello stesso quadro variazionale.
2. Evitamento del Campionamento Stocastico: A differenza degli approcci basati su equazioni differenziali stocastiche avanti-indietro (FBSDE), questa formulazione mean-field non richiede il campionamento di percorsi stocastici per la soluzione, portando potenzialmente a stime meno rumorose.
3. Generalizzazione a Dinamiche Complesse: La formulazione è estendibile a:
   • Dinamiche di Langevin sottodsmorzate (underdamped).
   • Rumore moltiplicativo $\Sigma(x)$.
   • Flussi di gradiente di Wasserstein-Kalman.
4. Connessione con il Ponte di Schrödinger: Viene dimostrata la relazione esplicita tra i potenziali $\psi_t$ della formulazione proposta e i potenziali di Schrödinger classici ($\phi_t, \hat{\phi}_t$) tramite la trasformazione esponenziale $\hat{\phi}_t = e^{\psi_t}$.
5. Estensioni Non Variazionali e Stocastiche:
   • Viene discussa una connessione con approcci non variazionali (Lagrangiani) proposti in letteratura precedente.
   • Viene esplorata una versione stocastica del principio di McKean-Pontryagin, dove $\beta_t$ è interpretato come un rumore browniano, collegando la formulazione alle FBSDE classiche. Tuttavia, l'autore nota che la versione variazionale deterministica (mean-field) è preferibile per la stabilità numerica.
6. Controllo della Matrice di Diffusione: L'autore propone un'estensione in cui la matrice di diffusione $\Sigma$ è trattata come un termine di controllo aggiuntivo, introducendo una penalità sulla sua norma di Frobenius. Questo porta a una regolarizzazione sulla "liscezza" del potenziale $\psi_t$.

4. Significato e Implicazioni

Questo lavoro è significativo perché fornisce un quadro teorico unificato per il trasporto ottimo regolarizzato entropicamente, basato su un principio del massimo generalizzato (McKean-Pontryagin).

• Vantaggi Computazionali: La natura mean-field e variazionale suggerisce che gli algoritmi numerici basati su questo principio (discretizzando lo spazio delle etichette e il tempo) potrebbero essere più efficienti e meno soggetti a rumore statistico rispetto ai metodi Monte Carlo basati su FBSDE.
• Flessibilità: La capacità di gestire dinamiche complesse (come quelle di Langevin) e di trattare la matrice di diffusione come variabile di controllo apre nuove strade per l'ottimizzazione in fisica statistica e nell'apprendimento automatico (es. flussi normalizzanti, generazione di dati).
• Teoria: Colma un gap tra la teoria del controllo ottimo stocastico classica e le moderne formulazioni di trasporto ottimo, offrendo una prospettiva geometrica e variazionale che preserva le simmetrie intrinseche del problema.

In sintesi, il paper propone un potente strumento analitico e numerico per risolvere problemi di trasporto ottimo dinamico, spostando il focus dal campionamento stocastico alla risoluzione di equazioni differenziali ordinarie (ODE) in uno spazio di funzioni di densità e potenziali.