Méthodes d'optimisation pour le Transport Optimal¶
- Opérateurs proximaux et optimisation non-lisse:
- "A Fast and Accurate Splitting Method for Optimal Transport" - Vien V. Mai et al.
- Optimisation sous contraintes et multiplicateurs de lagrange:
- "Efficient and Scalable Spatial Regularization of Optimal Transport" Lucas Brifault, David Cohen-Steiner & Mathieu Desbrun. (Siggasia 25)
Le transport optimal ? Jamais entendu parler...¶
Cas discret, deux mesures:
$$ \mu = \sum_{i=1}^n \delta_{x_i} \mu_i, \quad \nu = \sum_{i=1}^n \delta_{y_i} \nu_i $$
Nuances avec un nuage de points:
- Pas d'ordonnancement
- Comportement additif sur la position
- interprétation proba claire
Dans le cas $n=m, \mu_i = \nu_j = \frac{1}{n}$Assignation vs coupling¶
Coupling (Cas général)¶
On cherche un couplage qui minimise un cout de transport
$$
\pi_{i,j} : \text{Masse envoyée de $x_i$ à $y_j$}\\
s.t. \pi_{ij} > 0, \sum_{i} \pi_{ij} = \nu_j \text{ et } \sum_{j} \pi_{ij} = \mu_i
$$
$$ \min_{\pi \in \Pi(\mu,\nu)} \sum_{i,j} \pi_{ij}c(x_i,y_i) $$
Interprétation probabiliste des couplages¶
On définit les variables aléatoires $X\sim \mu, Y \sim \nu$.
_
Un couplage donne la loi de l'observation du couple $(X,Y) \sim \pi$.
_
Le problème de transport se réécrit simplement: $$ \min_{\pi\in\Pi(\mu,\nu)} \mathbb{E}_{(X,Y)\sim\pi}\left[c(X,Y)\right] $$
Régularisation la plus connue¶
$$ \min_{\pi>0} \langle C,\pi\rangle + \epsilon \text{KL}(\pi| \mu \otimes \nu) \\ \text{s.t. } \pi 1_n = \mu \\ \pi^\top 1_m = \nu $$
Rend le problème strictement convexe (gradient plus informatif) et donne des itérations explicites et parallélisables.
Optimisation sous contraintes¶
On cherche à résoudre: $$ \min_{\pi>0} \langle C,\pi\rangle \\ \text{s.t. } \pi 1_n = \mu \\ \pi^\top 1_m = \nu $$
C'est un problème d'optimisation, linéaire (donc convexe), sous contraintes (aussi linéaires), énorme littérature.
_
Deux cas de figures:
- Contraintes absolues $\implies $ optimisation Riemannienne (on en parlera pas ici)
- Contraintes cool raoul: Pas la mort si pas satisfaites exactement en pratiques
L'OT est clairement cool raoul.
Première approche: splitting et opérateurs proximaux¶
Réduction à un problème ""sans contrainte"": $$ \min_{\pi} \langle C,\pi\rangle + \iota_{\mathbf{R}^+}(\pi) + \iota_{\Pi(\mu,\nu)}(\pi) $$
où $\iota$ est une indicatrice au sens de l'optimisation:
$$ \iota_{A}(x) = \begin{cases} 0, & \text{si } x\in A \\ +\infty, & \text{si } x \notin A \end{cases} $$
_
Première idée: Descente de gradient?¶
Impossible car la fonction objectif n'est pas dérivable (entre autres)¶
Descente de gradient¶
$$ x_{i+1} = x_i - \alpha \nabla f(x_i) $$
Garantie si $\alpha < \frac{2}{L}$ où $L$ est la constante de Lipschitz de $\nabla f$.
Donc $f$ doit avoir un certain degré de régularité!
Echoue par exemple sur $$ f(x) = \|x\|_1 $$ Plus qu'un limitation théorique ici: on a une oscilation infinie en 0.
Idée de l'opérateur proximal¶
$$ x_{i+1} = x_i - \alpha \nabla f(x_{\color{red}{i+1}}) $$ Idée similaire aux intégrateurs implicites!
Opérateur proximal et optimisation non lisse¶
L'opérateur proximal d'une function $f$ est défini comme: $$ \text{prox}_{\rho f}(x) = \arg\min_y \frac{1}{2}\|x-y\|_2^2 + \rho f(y) $$
_
Hypothèses requises sur $f$ extrêmement faibles!
On peut remplacer dans la descente de gradient: $$ x_{i+1} = x_i - \alpha \nabla f(x_i) \implies x_{i+1} = \text{prox}_{\alpha f}(x_i) $$
Rendre ça utile par splitting:¶
Par exemple, si fonction de la forme: $$ \min_x f(x) = g(x) + h(x) $$ où $g$ est différentiable mais pas $h$ mais on sait calculer $\text{prox}_h$.
Algorithme du gradient proximal¶
$$ x_{i+1} = \text{prox}_{\frac{\alpha}{2} h} (x_i - \frac{\alpha}{2}\nabla g(x_i)) $$
Exemple, le problème du LASSO:¶
$$ \min_x \|Ax-b\|^2_2 + \lambda\|x\|_1 $$
Le terme quadratique attache aux données et le terme $L_1$ impose de la sparsité.
Intérêt : on sait calculer le prox pour plein de fonctions simples¶
Exemple: $$ \text{prox}_{\gamma\|\cdot\|_1}(x) = \left(\text{Soft}_\gamma(x_i)\right)_i \\ $$
où $$ \text{Soft}_\gamma(x) = \begin{cases} x - \gamma &, x > \gamma \\ 0 &, x\in [-\gamma,\gamma] \\ x + \gamma &, x < - \gamma \end{cases} $$
Prox d'une indicatrice¶
Si $f$ est une "indicatrice" d'un ensemble $A$, on a:
$$
\text{prox}_{\iota_A}(x) = \min_y \|x-y\|_2^2 + \iota_A(y) \\
= \min_{y\in A} \|x-y\|_2^2
$$
i.e. $\text{prox}_{\iota_A}(x) = \text{proj}_A(x)$
L'algo du gradient proximal donne l'algo du gradient projeté: $$ x_{i+1} = \text{proj}_A(x_i-\rho \nabla f(x_i)) $$
Le calcul de prox est un sport national dans certaines communautés https://proximity-operator.net/
Retour à l'OT¶
Les auteurs ré-écrivent l'objectif initial comme: $$ \min_{\pi} \underbrace{\langle C,\pi\rangle + \iota_{\mathbf{R}^+}(\pi)}_{g(\pi)} + \underbrace{\iota_{\Pi(\mu,\nu)}(\pi)}_{h(\pi)} $$
Et calculent le prox de $g$ et $h$ (forme close et pas trop compliquée), par exemple $$\text{prox}_{\rho g}(\pi)= \left[\pi- \rho C\right]_+$$
Douglas-Rachford¶
$$ \begin{cases} x_{i+1} = \text{prox}_{\rho g}(y_k) \\ y_{i+1} = y_i + \text{prox}_{\rho h}(2x_{i+1} - y_i) - x_{i+1} \end{cases} $$ Convergence pour tout $\rho > 0$ (vitesse variable).
Une partie de la contrib consiste à optimiser les appels mémoires et une implem clean et efficace en CUDA.
Résultats:¶
En gros bien meilleur que Sinkhorn:
- itérations un poil plus rapide
- plus robuste numériquement
- résultats beaucoup plus sparses
Convergence du coût de transport en $O(\frac{1}{k})$ au lieu de $O(\frac{1}{\sqrt{k}})$ pour Sinkhorn.
Complexité reste quadratique en le nombre de variables car plan de transport complet.
Transfert de couleur¶
Couleurs quantisées sur des clusters
Questions avant la suite ?¶
Régularisation spatiale du Transport Optimal¶
L'OT c'est très chouette mais pour du matching de forme c'est super moche parce qu'aucune notion de topologie.
Approche naïve à la Dirichlet¶
$$ \sum_{i,j}^n \pi_{ij}c(x_i,y_j) + \lambda \sum_{i,j,k,l}^n \pi_{i,j}\pi_{k,l}\frac{\|y_j-y_l\|^2}{\|x_i - x_k\|^2} $$
Problèmes:
- Problème quadratique donc bcp plus long à calculer
- mauvaises propriétés théoriques, ne donne pas une distance entre des probas
Autre approche:
Gromov-Wasserstein:¶
$$ \sum_{i,j,k,l}^n \pi_{i,j}\pi_{k,l}|d(x_i,x_k) - d(y_j,y_l)|^2 $$
- Très élégant et général, marche très bien
- NP-Hard
Approche du papier: régularisation de la mean map¶
Forward mean map (aussi Brenier Map dans la littérature) $$ \vec{m_\pi}(x) = \mathbb{E}_{(X,Y)\sim \pi}\left[Y|X = x\right] $$
I.e. moyenne pondérée de là où chaque particule est envoyée:
$$ \vec{m_\pi}(x) = \frac{\sum_j \pi_{ij}y_j}{\sum_j \pi_{ij}} $$
Régularisation convexe¶
Previous work de Mathieu sur de la régularisation spatiale du transport mais pas convexe, ici l'idée:
$$ \min_{\pi \in \Pi(\mu,\nu)} \sum_{i,j}^n \pi_{ij}c(x_i,y_j) + \epsilon H(\pi) + J_X(\vec{m_\pi}(x)) + J_Y(\overleftarrow{m}_{\pi}(x)) $$
Convexe si $J$ convexes, ils se concentrent sur des potentiels $J$ quadratiques $\vec{m_\pi} \cdot A \vec{m_\pi} = Q_A(\vec{m_\pi})$, par exemple:
$$ \min_{\pi \in \Pi(\mu,\nu)} \sum_{i,j}^n \pi_{ij}c(x_i,y_j) + \epsilon H(\pi) +Q_A(\vec{m_\pi}) $$
Par exemple, pour du matching entre surfaces, $A$ est lié à la discrétisation de la géométrie (ici hessienne sur un maillage $M$ par exemple)
$$ Q_A(\vec{m_\pi}) = \|H_M\vec{m_\pi} \|^2 $$
Qu'a t-on gagné ?¶
Pas grand chose ici parce que la relation entre $\vec{m}_{\pi}(x)$ et $\pi$ n'est pas linéaire...
Essayons de découpler les variables¶
$$ \min_{\pi \in \Pi(\mu,\nu)} \sum_{i,j}^n \pi_{ij}c(x_i,y_j) + \epsilon H(\pi)+ Q_A(\phi) \quad \text{s.t. } \phi= \vec{m_\pi} $$
On a juste ajouté une contrainte...
Les multiplicateurs de Lagrange¶
$$ \min_{g(x)=c} f(x) $$
equivalent à $$ \min_x \max_\lambda f(x) + \lambda (g(x) - c) $$ Problème "non-contraint"!
On cherche donc un point-selle du lagrangien: $$ \mathcal{L}(x,\lambda) = f(x) + \lambda (g(x) - c) $$
Qu'on trouve toujours par $$ \nabla_{x,\lambda} \mathcal{L} = 0 $$
Qu'est ce que ça donne:¶
$$ \mathcal{L}(x,\lambda) = f(x) + \lambda (g(x) - c) $$
$$ \nabla_{x,\lambda} \mathcal{L} = \begin{pmatrix} \nabla_x \mathcal{L} \\ \nabla_\lambda \mathcal{L} \end{pmatrix} = 0 $$
On obtient
$$ \nabla_x \mathcal{L} = \nabla f(x) + \lambda \nabla g(x) \\ \nabla_\lambda \mathcal{L} = g(x) - c $$
Donc la condition de point selle $\nabla_{x,\lambda} \mathcal{L} = 0$ donne $$ \nabla f(x) = \lambda \nabla g(x) \\ \underbrace{g(x) = c}_{\text{contrainte satisfaite}} $$
$$
\nabla f(x) = \lambda \nabla g(x)
$$Application pratique : ADMM¶
ADMM (ou ADA en français Algorithme des Directions Alternées)
Algo ULTRA important et générique, pour résoudre: $$ \min_{Ax+By=c} f(x) + g(y) $$
Passe par un Lagrangien augmenté:
$$ \mathcal{L}(x,y,\lambda,\rho) = f(x) + g(y) + \lambda\cdot( Ax+By - c) + \rho \| Ax + By - c \|^2 $$
En théorie ne change rien à l'optimum, mais rend le problème plus convexe en pratique.
_
Algo de minimisation alternée:
$$ x_{i+1}= \arg \min_x \mathcal{L}(x,y_i,\lambda_i) \\ y_{i+1}= \arg \min_x \mathcal{L}(x_{i+1},y,\lambda_i) \\ \lambda_{i+1}= \lambda_{i+1} + \rho_i (Ax_{i+1} + By_{i+1} - c) $$
Lagrangien du papier¶
$$ \min_{\pi \in \Pi(\mu,\nu)} \sum_{i,j}^n \pi_{ij}c(x_i,y_j) + \epsilon H(\pi) +Q_A(\phi) \quad \text{s.t. } \phi= \vec{m_\pi} $$
$$ \mathcal{L}(\pi,\phi,\lambda) = \sum_{i,j}^n \pi_{ij}c(x_i,y_j) + \epsilon H(\pi) +Q_A(\phi) + \langle \lambda, \phi - \vec{m_\pi}\rangle_\mu \\ = \underbrace{\sum_{i,j}^n \pi_{ij}(\underbrace{c(x_i,y_j) - \lambda_i y_i)}_{c_\lambda} + \epsilon H(\pi)}_{f(\pi,\lambda)} + \underbrace{Q_A(\phi) + \langle \lambda, \phi\rangle_\mu}_{g(\phi,\lambda)} $$
Plan de transport et mean map découplés, intéraction seulement par $\lambda$!
Tout est convexe/s'exprime en forme close, soit par résolution de système linéaire, soit par algo sinkhorn sur coût modifié $c_\lambda$.
Résultat: Matching de forme très efficace avec peu de landmarks¶
Résumé: Minimisation sous contraintes pour le transport optimal¶
On a vu deux types de gestion des contraintes:
- Optimisation non lisses : contraintes dures, opérateurs proximaux qui agissent comme des projecteurs sur les contraintes
- Optimisation par lagragien : variables duales, minimisation alternée, ADMM