[WIP] Translation Invariant Sinkhorn for UOT

README.md

@@ -370,4 +370,6 @@ distances between Gaussian distributions](https://hal.science/hal-03197398v2/fil
 
 [68] Chowdhury, S., Miller, D., & Needham, T. (2021). [Quantized gromov-wasserstein](https://link.springer.com/chapter/10.1007/978-3-030-86523-8_49). ECML PKDD 2021. Springer International Publishing.
 
-[69] Delon, J., & Desolneux, A. (2020). [A Wasserstein-type distance in the space of Gaussian mixture models](https://epubs.siam.org/doi/abs/10.1137/19M1301047). SIAM Journal on Imaging Sciences, 13(2), 936-970.
+[69] Delon, J., & Desolneux, A. (2020). [A Wasserstein-type distance in the space of Gaussian mixture models](https://epubs.siam.org/doi/abs/10.1137/19M1301047). SIAM Journal on Imaging Sciences, 13(2), 936-970.
+
+[70] Séjourné, T., Vialard, F. X., & Peyré, G. (2022). [Faster Unbalanced Optimal Transport: Translation Invariant Sinkhorn and 1-D Frank-Wolfe](https://proceedings.mlr.press/v151/sejourne22a.html). In International Conference on Artificial Intelligence and Statistics (pp. 4995-5021). PMLR.

examples/unbalanced-partial/plot_conv_sinkhorn_ti.py

@@ -0,0 +1,101 @@
+# -*- coding: utf-8 -*-
+"""
+===============================================================
+Translation Invariant Sinkhorn for Unbalanced Optimal Transport
+===============================================================
+
+This examples illustrates the better convergence of the translation
+invariance Sinkhorn algorithm proposed in [70] compared to the classical
+Sinkhorn algorithm.
+
+[70] Séjourné, T., Vialard, F. X., & Peyré, G. (2022).
+Faster unbalanced optimal transport: Translation invariant sinkhorn and 1-d frank-wolfe.
+In International Conference on Artificial Intelligence and Statistics (pp. 4995-5021). PMLR.
+
+"""
+
+# Author: Clément Bonet <clement.bonet@ensae.fr>
+# License: MIT License
+
+import numpy as np
+import matplotlib.pylab as pl
+import ot
+
+##############################################################################
+# Setting parameters
+# -------------
+
+# %% parameters
+
+n_iter = 50  # nb iters
+n = 40  # nb samples
+
+num_iter_max = 100
+n_noise = 10
+
+reg = 0.005
+reg_m_kl = 0.05
+
+mu_s = np.array([-1, -1])
+cov_s = np.array([[1, 0], [0, 1]])
+
+mu_t = np.array([4, 4])
+cov_t = np.array([[1, -.8], [-.8, 1]])
+
+
+##############################################################################
+# Compute entropic kl-regularized UOT with Sinkhorn and Translation Invariant Sinkhorn
+# -----------
+
+err_sinkhorn_uot = np.empty((n_iter, num_iter_max))
+err_sinkhorn_uot_ti = np.empty((n_iter, num_iter_max))
+
+
+for seed in range(n_iter):
+    np.random.seed(seed)
+    xs = ot.datasets.make_2D_samples_gauss(n, mu_s, cov_s)
+    xt = ot.datasets.make_2D_samples_gauss(n, mu_t, cov_t)
+
+    xs = np.concatenate((xs, ((np.random.rand(n_noise, 2) - 4))), axis=0)
+    xt = np.concatenate((xt, ((np.random.rand(n_noise, 2) + 6))), axis=0)
+
+    n = n + n_noise
+
+    a, b = np.ones((n,)) / n, np.ones((n,)) / n  # uniform distribution on samples
+
+    # loss matrix
+    M = ot.dist(xs, xt)
+    M /= M.max()
+
+    entropic_kl_uot, log_uot = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg, reg_m_kl, reg_type="kl", log=True, numItermax=num_iter_max, stopThr=0)
+    entropic_kl_uot_ti, log_uot_ti = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg, reg_m_kl, reg_type="kl",
+                                                                       method="sinkhorn_translation_invariant", log=True,
+                                                                       numItermax=num_iter_max, stopThr=0)
+
+    err_sinkhorn_uot[seed] = log_uot["err"]
+    err_sinkhorn_uot_ti[seed] = log_uot_ti["err"]
+
+##############################################################################
+# Plot the results
+# ----------------
+
+mean_sinkh = np.mean(err_sinkhorn_uot, axis=0)
+std_sinkh = np.std(err_sinkhorn_uot, axis=0)
+
+mean_sinkh_ti = np.mean(err_sinkhorn_uot_ti, axis=0)
+std_sinkh_ti = np.std(err_sinkhorn_uot_ti, axis=0)
+
+absc = list(range(num_iter_max))
+
+pl.plot(absc, mean_sinkh, label="Sinkhorn")
+pl.fill_between(absc, mean_sinkh - 2 * std_sinkh, mean_sinkh + 2 * std_sinkh, alpha=0.5)
+
+pl.plot(absc, mean_sinkh_ti, label="Translation Invariant Sinkhorn")
+pl.fill_between(absc, mean_sinkh_ti - 2 * std_sinkh_ti, mean_sinkh_ti + 2 * std_sinkh_ti, alpha=0.5)
+
+pl.yscale("log")
+pl.legend()
+pl.xlabel("Number of Iterations")
+pl.ylabel(r"$\|u-v\|_\infty$")
+pl.grid(True)
+pl.show()

examples/unbalanced-partial/plot_unbalanced_OT.py

@@ -73,7 +73,7 @@
 reg_m_l2 = 5
 mass = 0.7
 
-entropic_kl_uot = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg, reg_m_kl)
+entropic_kl_uot = ot.unbalanced.sinkhorn_unbalanced(a, b, M, reg, reg_m_kl, reg_type="kl")
 kl_uot = ot.unbalanced.mm_unbalanced(a, b, M, reg_m_kl, div='kl')
 l2_uot = ot.unbalanced.mm_unbalanced(a, b, M, reg_m_l2, div='l2')
 partial_ot = ot.partial.partial_wasserstein(a, b, M, m=mass)

ot/unbalanced/__init__.py

@@ -12,6 +12,7 @@
 from ._sinkhorn import (sinkhorn_knopp_unbalanced,
                         sinkhorn_unbalanced,
                         sinkhorn_stabilized_unbalanced,
+                        sinkhorn_unbalanced_translation_invariant,
                         sinkhorn_unbalanced2,
                         barycenter_unbalanced_sinkhorn,
                         barycenter_unbalanced_stabilized,
@@ -22,6 +23,7 @@
 from ._lbfgs import (lbfgsb_unbalanced, lbfgsb_unbalanced2)
 
 __all__ = ['sinkhorn_knopp_unbalanced', 'sinkhorn_unbalanced', 'sinkhorn_stabilized_unbalanced',
-           'sinkhorn_unbalanced2', 'barycenter_unbalanced_sinkhorn', 'barycenter_unbalanced_stabilized',
+           'sinkhorn_unbalanced_translation_invariant', 'sinkhorn_unbalanced2',
+           'barycenter_unbalanced_sinkhorn', 'barycenter_unbalanced_stabilized',
            'barycenter_unbalanced', 'mm_unbalanced', 'mm_unbalanced2', '_get_loss_unbalanced',
            'lbfgsb_unbalanced', 'lbfgsb_unbalanced2']