Wasserstein loss layer/criterion

AjayTalati: <p>Hello <img src="//discuss.pytorch.org/images/emoji/apple/smile.png?v=5" title=":smile:" class="emoji" alt=":smile:"></p>
<p>Are there any plans for an (approximate) Wasserstein loss layer to be implemented - or maybe its already out there?</p>
<p>It’s been in <code>mocha</code> for quite a while,</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pluskid/Mocha.jl/blob/master/src/layers/wasserstein-loss.jl" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pluskid/Mocha.jl/blob/master/src/layers/wasserstein-loss.jl" target="_blank" rel="nofollow noopener">pluskid/Mocha.jl/blob/master/src/layers/wasserstein-loss.jl</a></h4>
<pre><code class="lang-jl">@defstruct WassersteinLossLayer Layer (
  name :: AbstractString = "wasserstein-loss",
  (M :: Matrix = [], !isempty(M)),
  (lambda :: AbstractFloat = 0, lambda &gt; 0),
  (sinkhorn_iter :: Int = 50, sinkhorn_iter &gt; 0),
  (bottoms :: Vector{Symbol} = Symbol[], length(bottoms) == 2)
)
@characterize_layer(WassersteinLossLayer,
  has_loss  =&gt; true,
  can_do_bp =&gt; true,
  is_sink   =&gt; true,
)

type WassersteinLossLayerState{T} &lt;: LayerState
  layer :: WassersteinLossLayer
  loss  :: T

  K     :: Blob
  KM    :: Blob
  alpha :: Blob
</code></pre>

  This file has been truncated. <a href="https://github.com/pluskid/Mocha.jl/blob/master/src/layers/wasserstein-loss.jl" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>The theory’s and implementation is a little bit beyond my superficial understanding, (Appendix D), but it seems quite impressive!</p>
<aside class="onebox pdf">
  <header class="source">
      <a href="https://people.csail.mit.edu/hmobahi/pubs/wasserstein.pdf" target="_blank" rel="nofollow noopener">people.csail.mit.edu</a>
  </header>
  <article class="onebox-body">
    <a href="https://people.csail.mit.edu/hmobahi/pubs/wasserstein.pdf" target="_blank" rel="nofollow noopener"><span class="pdf-onebox-logo"></span></a>
<h3><a href="https://people.csail.mit.edu/hmobahi/pubs/wasserstein.pdf" target="_blank" rel="nofollow noopener">wasserstein.pdf</a></h3>

<p class="filesize">2.55 MB</p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


AjayTalati: <p>Yes I think that’s their particular application in the paper - but it could be more general than that?</p>
<p>The layer <code>wasserstein-loss.jl</code> coded in Mocha, seems to be fairly flexible.</p>
<p>You might have seen this Python optimal transport library,</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/rflamary/POT/blob/master/ot/bregman.py" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/rflamary/POT/blob/master/ot/bregman.py" target="_blank" rel="nofollow noopener">rflamary/POT/blob/master/ot/bregman.py</a></h4>
<pre><code class="lang-py"># -*- coding: utf-8 -*-
"""
Bregman projections for regularized OT
"""

# Author: Remi Flamary &lt;remi.flamary@unice.fr&gt;
#         Nicolas Courty &lt;ncourty@irisa.fr&gt;
#
# License: MIT License

import numpy as np


def sinkhorn(a, b, M, reg, method='sinkhorn', numItermax=1000, stopThr=1e-9, verbose=False, log=False, **kwargs):
    u"""
    Solve the entropic regularization optimal transport problem and return the OT matrix

    The function solves the following optimization problem:

    .. math::
</code></pre>

  This file has been truncated. <a href="https://github.com/rflamary/POT/blob/master/ot/bregman.py" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>I’m guessing that’s what they’ve implemented as their layer? Seems to be a solid piece of theory?</p>

tom: <p>From what I understand, the POT library solves 4.1 (Entropic regularization of the Wasserstein distance, say W(p,q) ), deriving the gradient in 4.2 and the relaxation in 4.3 (first going to W(p_approx,q_approx)+DKL(p_approx,p)+DKL(q_approx,q) and then generalising DKL to allow p/q approx to not be distributions seems to go beyond that.</p>
<p>Funny that they use the difference between MNIST class labels as a metric for the target. That reminded me of your regression approach.</p>

AjayTalati: <p>To be honest, I’m not too sure how to use the POT library yet - but if you want to play around in Mocha, here’s the test of the Wasserstein layer,</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pluskid/Mocha.jl/blob/master/examples/test-wasserstein.jl" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pluskid/Mocha.jl/blob/master/examples/test-wasserstein.jl" target="_blank" rel="nofollow noopener">pluskid/Mocha.jl/blob/master/examples/test-wasserstein.jl</a></h4>
<pre><code class="lang-jl">data_type = Float32

using Mocha
srand(123)

N = 1000
d = 4
lambda = 0.1
sinkhorn_iter = 10

n_hidden_layer = 2
n_hidden_unit = 64

test_wasserstein = true

X = rand(data_type, d, N)
Y_onehot = zeros(data_type, 1, N)
Y = zeros(data_type, d, N)
for i = 1:N
  Y[indmax(X[:,i]),i] = 1
</code></pre>

  This file has been truncated. <a href="https://github.com/pluskid/Mocha.jl/blob/master/examples/test-wasserstein.jl" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>and just for the sake of completness, here’s the code to go with the original paper “Sinkhorn Scaling for Optimal Transport”</p>
<p><a href="http://marcocuturi.net/SI.html" class="onebox" target="_blank" rel="nofollow noopener">http://marcocuturi.net/SI.html</a></p>
<p>That’s probably the best place to start! Hopefully I’ll be able to make some sense of it all soon?</p>

tom: <p>Hi.<br>
The .jl shows</p>
<pre><code>Performance on Validation Set after 10000 iterations
Accuracy (avg over 1000) = 95.9%" 
</code></pre>
<p>or so?</p>
<p>float32 does not seem to provide the precision necessary to implement unmodified sinkhorn algorithm, at least in the Python Optimal Transport’s 1-d-OT example.</p>
<p>The log-stabilized sinkhorn algorithm seems to work better at first sight. In the example, the difference between unregularized and regularized is ~1e6 and the difference between numpy@float64 and pytorch@float32 is ~1e-7. Interestingly, the mocha code seems to implement the unstabilized algorithm (unless they are doing the stabilization elsewhere).</p>
<p>I’ll see about making it into a layer on another day.</p>

AjayTalati: <p>Wow, that was quick!</p>
<p>Yes I get about the same,</p>
<pre><code class="lang-auto">22-Mar 21:36:03:INFO:root:## Performance on Validation Set after 10000 iterations
22-Mar 21:36:03:INFO:root:---------------------------------------------------------
22-Mar 21:36:03:INFO:root:  Accuracy (avg over 1000) = 95.8000%
22-Mar 21:36:03:INFO:root:---------------------------------------------------------
</code></pre>
<p>I’m still going over Marco Cuturi’s Matlab code - implementing a Wasserstein layers quite a long term project - unless you’ve got an immediate application in mind?</p>
<p>For the time being I’m content with just understand it mathematically.</p>
<p>It does seem to have a lot of potential if you want to train a network to give fast approximations to - existing slow simulation algorithms, or algorithms that currently calculate a Wasserstein metric using a linear program - there’s a lot of things like this in scientific computing - that’s actually the application I’ve got in mind.</p>
<p>Chiyuan Zhang (pluskid) has already applied this to computationally intensive geological simulations,</p>
<p><a href="http://pluskid.org/papers/TLE2017-seismic.pdf" class="onebox" target="_blank" rel="nofollow noopener">http://pluskid.org/papers/TLE2017-seismic.pdf</a></p>
<p>So basically - once your nets trained - you can run it forward in batch mode, and get output in second, which previously would have taken hours. Here’s a quote from the peper,</p>
<p>Recent work (Zhang et al., 2014; Frogner et al., 2015;<br>
Dahlke et al., 2016) demonstrates a new approach that builds<br>
and uses a deep neural network (DNN) statistical model to<br>
transform raw-input seismic data directly to the final mapping<br>
of faults in 2D. In this case, fault locations were chosen as the<br>
output because of their relevance to optimizing production in<br>
existing fields. DNNs are built on the premise that they can<br>
be used to replicate any function (in theory, even a nonlinear<br>
one like acoustic wave propagation). We show that DNNs can<br>
be used to identify fault structure in 3D volumes with reasonable<br>
accuracy. The greater promise is that as computational tools<br>
improve, we can use even more complex neural networks to<br>
improve accuracy. The costs are all computational, mostly in<br>
the form of training incurred only once up front. Once the<br>
neural network is trained, predictions can be produced in a<br>
fraction of the training time and of the time needed for traditional<br>
physics modeling with partial differential equations.<br>
This means that nearly instantaneous earth models could be<br>
created as acquisition progresses.</p>

AjayTalati: <p>Hi <a class="mention" href="/u/tom">@tom</a>, I just saw one of Marco Cuturi’s recent papers, and if I understood correctly, it gives a method to calculate the  divergence between distributions using something like SGD, rather than Sinkhorn’s algorithm</p>
<p>That’s much more understandable, but I’m not sure how easy it is to implement? Basically, it would involve constructing a layer which itself would involve a sgd loop! Pretty funky <img src="//discuss.pytorch.org/images/emoji/apple/smile.png?v=5" title=":smile:" class="emoji" alt=":smile:"></p>
<p>I know that there’s already a leanrable quadratic programming layer that’s been implemented, <a href="https://github.com/locuslab/qpth" rel="nofollow noopener">https://github.com/locuslab/qpth</a> but this seems more general than that</p>
<p>Anyway here’s the link, "Stochastic Optimization for Large-scale Optimal Transport "</p>
<p><a href="https://arxiv.org/abs/1605.08527" class="onebox" target="_blank" rel="nofollow noopener">https://arxiv.org/abs/1605.08527</a></p>

AjayTalati: <p>I found the code for “Stochastic Optimization for Large-scale Optimal Transport”</p>
<p><a href="https://github.com/audeg/StochasticOT" class="onebox" target="_blank" rel="nofollow noopener">https://github.com/audeg/StochasticOT</a></p>
<p>It doesn’t look too confusing <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=9" title=":smile:" class="emoji" alt=":smile:">, I think I’m starting to understand what’s going on!</p>

smth: <p>you can optimize divergence between distributions using the Wasserstein metric using a reformulated GAN.<br>
Here is our paper on that:</p>
<aside class="onebox whitelistedgeneric">
  <header class="source">
      <img src="https://static.arxiv.org/static/browse/0.3.0/images/icons/favicon.ico" class="site-icon" width="32" height="32">
      <a href="https://arxiv.org/abs/1701.07875" target="_blank">arXiv.org</a>
  </header>
  <article class="onebox-body">
    <img src="" class="thumbnail" width="" height="">

<h3><a href="https://arxiv.org/abs/1701.07875" target="_blank">Wasserstein GAN</a></h3>

<p>We introduce a new algorithm named WGAN, an alternative to traditional GAN
training. In this new model, we show that we can improve the stability of
learning, get rid of problems like mode collapse, and provide meaningful
learning curves useful for...</p>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>and the corresponding code (which is awfully simple):</p>
<p><a href="https://github.com/martinarjovsky/WassersteinGAN" class="onebox" target="_blank">https://github.com/martinarjovsky/WassersteinGAN</a></p>

AjayTalati: <p>Thanks <a class="mention" href="/u/smth">@smth</a>  - seems like there’s quite a few ways of doing the same thing?</p>
<ol start="0">
<li>Exact calculation using the <a href="https://github.com/wmayner/pyemd" rel="nofollow noopener">pyemd</a> module - which is really slow because it does a linear program<br>
i ) If I understand correctly, the <code>wasserstein.jl</code> layer in Mocha uses Sinkhorn’s algorithm to approximate the Wasserstein distance<br>
ii) The code in the repo above which uses something similar to <code>sgd</code>,<br>
iii) Your code which I need to have a good look at.</li>
</ol>
<p>They should all roughly give the same value!</p>
<p>It should be pretty simple to do the test between the different methods - here’s the example code from case 0)</p>
<pre><code class="lang-auto">&gt;&gt;&gt; from pyemd import emd
&gt;&gt;&gt; import numpy as np
&gt;&gt;&gt; first_signature = np.array([0.0, 1.0])
&gt;&gt;&gt; second_signature = np.array([5.0, 3.0])
&gt;&gt;&gt; distance_matrix = np.array([[0.0, 0.5], [0.5, 0.0]])
&gt;&gt;&gt; emd(first_signature, second_signature, distance_matrix)
3.5
</code></pre>

tom: <p>Thanks you, <a class="mention" href="/u/ajaytalati">@AjayTalati</a> and <a class="mention" href="/u/smth">@smth</a> for the links. (I had seen the Wasserstein GAN paper and code, but did not dig into the code much yet). The Sinkhorn algorithm is iterative, too, but as Genevay et al point out, its “batch” nature may make it prohibitive for large scale applications.</p>
<p>Python Optimal Transport also has an exact solver and compares to the entropy regularized version here: <a href="https://github.com/rflamary/POT/blob/master/examples/Demo_1D_OT.ipynb" rel="nofollow noopener">https://github.com/rflamary/POT/blob/master/examples/Demo_1D_OT.ipynb</a>.<br>
(Though in 1D you can just use the CDF_1^(-1)(CDF(x)) transport plan.)</p>
<p>In some applications Wasserstein distances beyond W_1=EMD might be interesting (e.g. the W_2 inner product might be handy sometimes).</p>
<p>The Wasserstein GAN paper and method is awesome, but I am not quite certain that the GAN distance does actually approximate the W_1 distance:</p>
<ul>
<li>As the authors point out there is the issue whether the supremum is actually attained in the test set of the maximization, (not sure how that compares with the discretization you have to do before using Sinkhorn etc., the linked paper Genevay et al paper kernelizes for the continuous case).</li>
<li>The other thing about the Wasserstein GAN is that the maximizer of equation (3), even if the maximizer of equation (2) is contained in W, will not be a scaled version of the latter. This is not terribly relevant to the ends of the article as you still get the a good norm, so the authors do well to only briefly mention it. On the other hand it would be relevant if you use the W_1 distance for something where you need the W_1 distance itself, so you would need to compute the Lipschitz constant in the maximization procedure and divide by it in the quantity maximized in (3), i.e. for the candidates and not just the maximum.</li>
</ul>
<p>Well, I better stop the talking get the code in shape to share. <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

AjayTalati: <p>Hi <a class="mention" href="/u/tom">@tom</a>, it’s really cool that you’re getting interested in this problem. I just want to avoid you going down blind alley’s.</p>
<p>The paper "Stochastic Optimization for Large-scale Optimal Transport " <a href="https://arxiv.org/abs/1605.08527" rel="nofollow noopener">https://arxiv.org/abs/1605.08527</a>, is a conference paper, and they’re usually a bit of a <strong>wild card</strong>. I spent a while on the code, audeg/StochasticOT, and it’s probably more suited to information retrieval rather than actually training a network like <a class="mention" href="/u/smth">@smth</a>’s GAN code, or as a layer in one. So I think I made a mistake, and it’s perhaps not such a good idea implementing it as a layer.</p>
<p>I don’t want to mislead you - it’s probably a good idea to work on something that’s been proven to be useful. If you simply try to reproduce Chiyuan Zhang (pluskid) <code>Wassertein.jl</code> layer, in the  code at the top of this thread, that would be a <strong>safe</strong> thing to do. That’s something that were reasonably sure work’s, and if you get it working in PyTorch it would be something that you could reference to, and reuse in the future.</p>
<p>Anyway, starting with the easy stuff - at the moment I can’t get the entropy regularised version in,</p>
<p><a href="https://github.com/rflamary/POT/blob/master/examples/Demo_1D_OT.ipynb" rel="nofollow noopener">https://github.com/rflamary/POT/blob/master/examples/Demo_1D_OT.ipynb</a>.</p>
<p>and the exact EMD solver used in <a href="https://github.com/wmayner/pyemd" rel="nofollow noopener">PyEMD</a> to give roughly the same number’s?  What I mean is, (I think) the functions in the POT library return the transportation plan matrix T, and to get the actual wasserstein/EMD divergence, you calculate, the inner product,</p>
<pre><code class="lang-auto">&lt;T,M&gt; = trace(T'M) 
</code></pre>
<p>Where M is the ground metric matrix. But I get different numbers from both libraries, so I’m not sure which is right? I know you need to tune the regularization parameter <code>lambda</code>, but it should be easy to do that using downhill simplex/<a href="https://docs.scipy.org/doc/scipy-0.18.1/reference/tutorial/optimize.html" rel="nofollow noopener">Nelder Mead</a></p>
<p>I want to get these test cases to reconcile - that way we’ve got something to check against when we try to code things, and do more complicated stuff in PyTorch. Otherwise, it’s too easy to make a mistake, without something solid to test against. <img src="//discuss.pytorch.org/images/emoji/apple/blush.png?v=5" title=":blush:" class="emoji" alt=":blush:"></p>

AjayTalati: <p>OK, sorry about this <a class="mention" href="/u/tom">@tom</a> <img src="//discuss.pytorch.org/images/emoji/apple/blush.png?v=5" title=":blush:" class="emoji" alt=":blush:"> , it seems the transport matrices from both PyEMD and POT do roughtly match - do want to check this?</p>
<pre><code class="lang-auto">
import numpy as np
import matplotlib.pylab as pl
import pyemd 
import ot
from ot.datasets import get_1D_gauss as gauss

#%% parameters

n=100 # nb bins

# bin positions
x=np.arange(n,dtype=np.float64)

# Gaussian distributions
a=gauss(n,m=20,s=5) # m= mean, s= std
b=gauss(n,m=60,s=10)

# loss matrix
M=ot.dist(x.reshape((n,1)),x.reshape((n,1)))
M/=M.max()

#%% EMD with PyEMD library

div, T_pyemd = pyemd.emd_with_flow(a,b,M)
T_pyemd = np.array(T_pyemd)

#%% EMD with POT library

T_pot = ot.emd(a,b,M)

#%% Sinkhorn with POT library

lambd=1e-3
T_sink = ot.sinkhorn(a,b,M,lambd)

#%% plot each of the transport matrices

ot.plot.plot1D_mat(a,b,T_pyemd,'Exact OT matrix using PyEMD')
pl.show()

ot.plot.plot1D_mat(a,b,T_pot,'Exact OT matrix using POT')
pl.show()

ot.plot.plot1D_mat(a,b,T_sink,'Sinkhorn OT matrix using POT')
pl.show()
</code></pre>

tom: <p>To compute the distance, you need to integrate:</p>
<pre data-code-wrap="python"><code class="lang-python">dist=(M*T).sum()
</code></pre>
<p>Something seems to be up with pyemd, though: If you use your example but with larger stddev (e.g. 10 and 30), the “optimal” plan pyemd emits here clearly is broken (as the monotone transport function gives the optimal plan and I see two distinct areas)…</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/1X/10fba662e226d6fee1d2143dcc9b127a6bf5ca35.png" width="375" height="264"></p>

AjayTalati: <p>Hey thanks a lot for that! I think you’ve found something!</p>
<p>I do think PyEMD is a bit finicky, but the distances seem to be roughly the same,</p>
<pre><code class="lang-auto">&gt;&gt;&gt; #%% EMD with PyEMD library
... 
&gt;&gt;&gt; div_pyemd, T_pyemd = pyemd.emd_with_flow(a,b,M)
&gt;&gt;&gt; T_pyemd = np.array(T_pyemd)
&gt;&gt;&gt; div_pyemd
0.16684170730400028
&gt;&gt;&gt; 
&gt;&gt;&gt; #%% EMD with POT library
... 
&gt;&gt;&gt; T_pot = ot.emd(a,b,M)
&gt;&gt;&gt; div_pot_emd = np.trace( np.matmul( T_pot.transpose() , M) )
&gt;&gt;&gt; div_pot_emd
0.16580581101906636
</code></pre>
<p>I guess then, we should go with the POT library, as that’s probably more reliable?</p>

AjayTalati: <p>On another note, I think their choice of <code>lambda=1e-3</code> is reasonable, I tried to tune it but it did’nt really make much difference,</p>
<pre><code class="lang-auto">
#%% Use the downhill simplex method to find lambd
 
import numpy as np
from scipy.optimize import minimize

T_pot = ot.emd(a,b,M)
div_pot_emd = np.trace( np.matmul( T_pot.transpose() , M) )

#objective functon
def tune_lambda(x):
    #%% Sinkhorn approximate EMD using POT library
    T_sink = ot.sinkhorn(a,b,M,x)
    div_pot_sink = np.trace( np.matmul( T_sink.transpose() , M) )
    error = abs(div_pot_emd - div_pot_sink)
    return error

x0 = 1e-3 # initial guess

res = minimize(tune_lambda, x0, method='nelder-mead', options={'xtol': 1e-8, 'disp': True})

print(res.x)
</code></pre>

AjayTalati: <p>Some speed tests for you <a class="mention" href="/u/tom">@tom</a>, only x10 at the moment?</p>
<p>It should be faster when i) it’s run on the gpu, ii) the histograms get bigger, iii) increase the dimensions, i.e images or higher dimension tensors, iv) tune <code>lambda</code></p>
<p>Check out figure 4, from Marco Cuturi’s original paper,</p>
<aside class="onebox pdf">
  <header class="source">
      <a href="https://papers.nips.cc/paper/4927-sinkhorn-distances-lightspeed-computation-of-optimal-transport.pdf" target="_blank" rel="nofollow noopener">papers.nips.cc</a>
  </header>
  <article class="onebox-body">
    <a href="https://papers.nips.cc/paper/4927-sinkhorn-distances-lightspeed-computation-of-optimal-transport.pdf" target="_blank" rel="nofollow noopener"><span class="pdf-onebox-logo"></span></a>
<h3><a href="https://papers.nips.cc/paper/4927-sinkhorn-distances-lightspeed-computation-of-optimal-transport.pdf" target="_blank" rel="nofollow noopener">4927-sinkhorn-distances-lightspeed-computation-of-optimal-transport.pdf</a></h3>

<p class="filesize">360.42 KB</p>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<pre><code class="lang-auto">#-----------------------------------------------------------------------------------------
#%% Speed Tests - see https://github.com/rflamary/POT/blob/master/examples/plot_OT_1D.py
#-----------------------------------------------------------------------------------------

import numpy as np
import pyemd 
import ot
import functools
from time import time
from ot.datasets import get_1D_gauss as gauss

#@functools.lru_cache (maxsize=None) # hash(tuple( yadda, yadda
def pyemd_linprog(n,a,b,M):
    start = time()
    for i in range(n): 
        div_pyemd = pyemd.emd(a,b,M)
    run_time = time() - start
    print(run_time)

#@functools.lru_cache (maxsize=None)
def pot_sinkhorn(n,a,b,M,lambd):
    start = time()
    for i in range(n): 
        T_sink = ot.sinkhorn(a,b,M, lambd)
        #div_pot_sink = np.trace( np.matmul( T_sink.transpose() , M) )
    run_time = time() - start
    print(run_time)

#%% parameters

n=128 # nb bins

# bin positions
x=np.arange(n,dtype=np.float64)

# Gaussian distributions
a=gauss(n,m=20,s=5) # m= mean, s= std
b=gauss(n,m=60,s=10)

# loss matrix
M=ot.dist(x.reshape((n,1)),x.reshape((n,1)))
M/=M.max()

##--------
runs=1000
lambd = 1

&gt;&gt;&gt; pyemd_linprog( runs, a, b, M )
21.02236819267273
&gt;&gt;&gt; pot_sinkhorn( runs, a, b, M, lambd )
2.0207273960113525

</code></pre>

tom: <p>Hi <a class="mention" href="/u/ajaytalati">@AjayTalati</a>,</p>
<p>I started an implementation here: <a href="https://github.com/t-vi/pytorch-tvmisc/blob/master/wasserstein-distance/Pytorch_Wasserstein.ipynb" rel="nofollow noopener">https://github.com/t-vi/pytorch-tvmisc/blob/master/wasserstein-distance/Pytorch_Wasserstein.ipynb</a></p>
<p>I’m not terribly impressed by the numerical stability, I’ll have to look into that.</p>
<p>Best regards</p>
<p>Thomas</p>

AjayTalati: <p>Hi <a class="mention" href="/u/tom">@tom</a></p>
<p>Wow !!!</p>
<p>That’s pretty amazing how quickly you did that !!! Very, very impressive.</p>
<p>Let me go over it and try to do some testing - I need to get my slow old brain working <img src="//discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=5" title=":slight_smile:" class="emoji" alt=":slight_smile:">  !! Be nice to do a test verses the unregularized version, i.e.</p>
<p><code>ot.lp.emd : Unregularized OT</code></p>
<p>So now that you have a Wasserstein loss that you can backprop through maybe you want to train a plain vanilla GAN with it, it would be interesting to see how it compares, here’s a quite basic version that’s new, and converges quite fast,</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/AjayTalati/generative-models/blob/master/GAN/boundary_equilibrium_gan/began_pytorch.py" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/AjayTalati/generative-models/blob/master/GAN/boundary_equilibrium_gan/began_pytorch.py" target="_blank" rel="nofollow noopener">AjayTalati/generative-models/blob/master/GAN/boundary_equilibrium_gan/began_pytorch.py</a></h4>
<pre><code class="lang-py">import torch
import torch.nn
import torch.nn.functional as nn
import torch.autograd as autograd
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import os
from torch.autograd import Variable
from tensorflow.examples.tutorials.mnist import input_data


mnist = input_data.read_data_sets('../../MNIST_data', one_hot=True)
mb_size = 32
z_dim = 10
X_dim = mnist.train.images.shape[1]
y_dim = mnist.train.labels.shape[1]
h_dim = 128
cnt = 0
</code></pre>

  This file has been truncated. <a href="https://github.com/AjayTalati/generative-models/blob/master/GAN/boundary_equilibrium_gan/began_pytorch.py" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Be interesting if you could use your loss layer to improve it?</p>
<p>Perhaps you want to get in touch with Rémi Flamary, <a href="http://remi.flamary.com/" rel="nofollow noopener">http://remi.flamary.com/</a>, I’m sure he’ll be very impressed and be bursting with ideas for possible collaboration <img src="//discuss.pytorch.org/images/emoji/apple/smile.png?v=5" title=":smile:" class="emoji" alt=":smile:"></p>
<p>Best regards,</p>
<p>Ajay</p>

AjayTalati: <p>Hi Tom,</p>
<p>Yes as you said, I’m also getting a lot of  <code>Warning: numerical errrors</code> numerical stability warnings?</p>
<p>Your implementations fine, (thanks once again for trying it), I’m guessing the problems simply the inherent instability of the different versions of the SK algorithm??</p>
<p>I’ve tried testing the linear programming <code>ot.emd</code>, against your implementation, and the numpy functions <code>ot.bregman.sinkhorn_epsilon_scaling(a,b,M,1)</code> etc. Like you showed the stabilized algorithm is much more stable than the vanilla version, although the relative rankings are still a little off?</p>
<p>So practically the task we want to address is not really can be reproduce the same values as the linear program <code>emd</code> algorithm. Rather we’re interested in ranking distributions.</p>
<p>To put it simply, if the linear program emd algorithm ranks, A and B, closer than A and C, then any approximate algorithm, (eg Sinkhorn-Knopp), should also give the same relative ranking.</p>
<p>I’m also trying it with discrete distributions, i.e. the histograms used in the original emd paper, but it works much better with Gaussians.</p>
<p>All the best,</p>
<p>Ajay</p>