Sinkhorn algorithm questions

[edgarriba]

hi @fwilliams nice lib ! :)

I have a couple of questions about the sinkhorn algorithm you have implemented:

1- Do you have any paper reference about the original algorithm?
2- The a, b dirac weights shouldn't be 1/num_samples ? Instead in your examples you initialize to ones, and in practice I have found that with other values the algorithm doesn't converge properly.

Thanks,

[fwilliams]

Hey @edgarriba thanks for posting this issue.

1. The original paper for the algorithm can be found here: https://arxiv.org/abs/1306.0895
2. Yes in principle a and b should be probability distributions though I've found that there aren't issues if they are both scaled by a constant factor. I've mostly just used this for distributions in low dimensional spaces (not sure if that matters here)

Could you post a simplified example of an input that fails to converge so I can look into it?

[fwilliams]

Any updates on this? If there's a bug I'd love to fix it but need an example to reproduce it.

[edgarriba]

hey @fwilliams , assuming this small toy example:

N: int = 10                                                 
x1 = torch.rand(1, N, 128) * 1e1                            
x2 = torch.rand(1, N, 128) * 1e1                            
x2[:, :5] = x1[:, :5]                                       
                                                            
M = torch.cdist(x1, x2, p=2)                                
A = torch.ones(1, N) / N                                    
B = torch.ones(1, N) / N                                    
                                                            
P = sinkhorn(A, B, M, 1e-1, max_iters=100, stop_thresh=1e-5)

I expected always to get back probabilities between (0,1)

[fwilliams]

Hey @edgarriba,

So I ran your example and the output is indeed correct (See output below).

In your case, the tensor P should have shape (1, 10, 10) and P[0] will be a Joint PMF i.e. P[0].sum() == 1.

If you want P[0] as a doubly stochastic matrix (i.e. each row and column is a probability distribution), then in your case you need to multiply P by N (or pass in A and B as torch.ones(1, N) like in the example code).

In [1]: import torch                                                                                                                                                                                                                                           

In [2]: from fml.functional import sinkhorn, pairwise_distances                                                                                                                                                                                                

In [3]: N = 10                                                                                                                                                                                                                                                 

In [4]: x1 = torch.rand(1, N, 128) * 1e1                                                                                                                                                                                                                       

In [5]: x2 = torch.rand(1, N, 128) * 1e1                                                                                                                                                                                                                       

In [6]: x2[:, :5] = x1[:, :5]                                                                                                                                                                                                                                  

In [7]: M = pairwise_distances(x1, x2)                                                                                                                                                                                                                         

In [8]: A = torch.ones(1, N) / N                                                                                                                                                                                                                               

In [9]: B = torch.ones(1, N) / N                                                                                                                                                                                                                               

In [10]: P = sinkhorn(A, B, M, 1e-1, max_iters=100, stop_thresh=1e-5)                                                                                                                                                                                          

In [11]: P.sum()                                                                                                                                                                                                                                               
Out[11]: tensor(1.0000)

In [12]: P                                                                                                                                                                                                                                                     
Out[12]: 
tensor([[[0.1000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0000],
         [0.0000, 0.1000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0000],
         [0.0000, 0.0000, 0.1000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.1000, 0.0000, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0500, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0500],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0500],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1000,
          0.1000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.1000, 0.1000, 0.0000,
          0.0000, 0.0000],
         [0.0000, 0.0000, 0.0000, 0.0000, 0.0500, 0.0000, 0.0000, 0.0000,
          0.0000, 0.0000]]])

[fwilliams]

Does this resolve your issue? Did I miss anything?

[fwilliams]

I'm closing this issue due to lack of activity and because I think it's resolved. If that's not the case @edgarriba please feel free to re-open it.