PythonOT / POT

POT : Python Optimal Transport
https://PythonOT.github.io/
MIT License
2.41k stars 499 forks source link

Sinkhorn_debiaised_barycenter seems too diffuse #458

Open tlacombe opened 1 year ago

tlacombe commented 1 year ago

Describe the bug

Theory (Theorem 3 in this paper ) tells us that the Sinkhorn barycenter between two Gaussian distribution with the same std $\sigma$ should be a Gaussian with std $\sigma$.

However, when computing the barycenter between two Dirac-ish measures (Eulerian representation : measures are supported on a grid, with the mass concentrated on a single pixel), the barycentric interpolation using convolutional_barycenter2d_debiased returns a quite spread measure, which seems quite surprising.

To Reproduce

Steps to reproduce the behavior:

  1. Define two measures on a grid of size (say) $20 \times 20$.
  2. Compute there barycenter using convolutional_barycenter2d_debiased with parameter $t \in (0,1)$
  3. Plot the result : it is quite blurry.

Screenshots

The "interpolation" we get by computing the Sinkhorn debiaised barycenter as a minimizer of $\mu \mapsto (1-t) S_\epsilon(\mu, \mu0) + t S\epsilon(\mu, \mu_1)$, with $\mu_0, \mu_1$ being two Dirac/very sharp Gaussian measures. The regularization parameter $\epsilon$ is set to reg=0.1.

image

Note : a similar behavior occurs when considering 1D measures, suggesting that the issue is not intrinsic to the convolutional-2D approach (did not investigated this in detail). The following screenshot show the returned barycenter for two 1D-Gaussian distributions, with reg=0.1 . The Gaussian distribution midway is too wide.

image

Code sample

To reproduce the plot in 2D :

import numpy as np
import matplotlib.pyplot as plt
from ot.bregman import convolutional_barycenter2d_debiased

def F(mu0, mu1, t, reg=1.):
    mu_t = convolutional_barycenter2d_debiased(np.array([mu0, mu1]), weights=np.array([1-t, t]), reg=reg, method='sinkhorn_log')
    return mu_t

def plot_bary(mu0, mu1, ts, reg=1.):
    n, m = mu0.shape
    vmax = max(np.max(mu0), np.max(mu1))
    fig, axs = plt.subplots(1, len(ts)+2, figsize=(10 * len(ts), 10))
    ax = axs[0]
    ax.imshow(mu0, cmap='Reds', vmin=0, vmax=vmax)
    ax.set_title("Source", fontsize=24)

    for i,t in enumerate(ts):
        ax = axs[i+1]
        mu_t = F(mu0, mu1, t, reg=reg)
        ax.imshow(mu_t * vmax/mu_t.max(), cmap='Purples', vmin=0, vmax=vmax)
        ax.set_title("Interp., t=%s" %t, fontsize=24)

    ax = axs[-1]
    ax.imshow(mu1, cmap='Blues', vmin=0, vmax=vmax)
    ax.set_title("Target", fontsize=24)

    fig.suptitle("Sinkhorn barycenter interpolation with reg=%s" %reg, fontsize=40)

n = 20
x, y = 2 * n//10, 8 * n//10

ts = (0.01, 0.25, 0.5, 0.75, 0.99)

mu0 = np.zeros((n, n))
mu1 = np.zeros((n, n))
mu0[x, x] = 1
mu1[y, y] = 1

plot_bary(mu0, mu1, ts, reg=0.1)

Expected behavior

The barycenter of two Dirac-ish measures should be (close to, up to numerical approximation) a Dirac-ish measure supported midway. More generally, the barycenter of two gaussian with std $\sigma$ should have std $\sigma$.

Environment (please complete the following information):

>>> import platform; print(platform.platform())
Linux-5.19.0-38-generic-x86_64-with-glibc2.10
>>> import sys; print("Python", sys.version)
Python 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:10) 
[GCC 10.3.0]
>>> import numpy; print("NumPy", numpy.__version__)
NumPy 1.20.3
>>> import scipy; print("SciPy", scipy.__version__)
SciPy 1.5.2
>>> import ot; print("POT", ot.__version__)
POT 0.8.2

Additional context

Has been discussed with @hichamjanati (who agrees on the issue). We plan to investigate this later. This issue is for documentation---discussion is welcome!

rflamary commented 1 year ago

very good point, I was a bit underwhelemd by the debiased barycenters that did not look that much bebiased in practice in practice, maybe it is an implementation problem?

tlacombe commented 1 year ago

Yes it really looks like an implementation issue (or maybe something deeper, like the IBP not converging?). I have no clue on what is happening yet. I did not investigate in details; I am quite busy these weeks (so is @hichamjanati ), but we may look at it closer later. In the meantime, if people have suggestions/remarks, there are welcome.

rflamary commented 1 year ago

Of course we are all sawmped ;). thanks for looking into it !