Faster exact EMD for distributions on shared multi-dimensional cartesian grids (with cityblock or sqeuclidean ground metric)

[tvercaut]

Hi,

I got interested in computing OT between d-dimensional histograms (supported on Cartesian product grids, not scattered points).

The Ling & Okada paper [1] has shown that in this case, and with the L1 (cityblock) ground metric, we can solve a min-cost-flow problem on a neighbourhood graph instead of on the full bipartite graph. Unfortunately, the reference implementation from the authors has a restrictive license (see https://haibinling.github.io/publication-selected.htm) and there is no widely available alternative (the one in opencv contrib seems broken, see opencv/opencv_contrib#3879).

For L2 (sqeuclidean), Auricchio et al. [2], have proposed a similar construct using (d+1)-partite graphs. Unfortunately, the repo they mention in the paper is empty, see https://github.com/stegua/dpartion-nips2018.

I (or rather some coding agents with extensive supervision) took a stab at re-implementing these:
https://github.com/tvercaut/emdgrid

For L1, I found that using the graph formulation of [1] but using an off-the-shelf solver worked well and outperformed re-implementing the solver discussed in the paper. LEMON Cost Scaling has a bit of an edge over LEMON Network Simplex but the difference is rather small. In any case, computation is one or two orders of magnitude faster than using the full bipartite formulation (for example used in the general emd2_lazy). As an illustration solving EMD-L1 on a 30x30x30 grid with this approach takes ~60 ms instead of 2-5 sec. Since POT already ships with a vendored version of LEMON Network Simplex, integrating a similar version of EMD-L1 in POT should be a reasonable coding effort.

For L2, the speed up of [2] over emd2_lazy is a little bit lower than for the L1 case but still comfortably in the one or two orders of magnitude faster regime. LEMON Network Simplex is also used there so adapting it for POT's code base should still be reasonable.

Would such a feature be of interest to POT or is OT on d-dimensional Cartesian grids too niche a problem?

Best wishes,
Tom

[1] Ling H, Okada K. An efficient earth mover's distance algorithm for robust histogram comparison. IEEE transactions on pattern analysis and machine intelligence. 2007 May 31;29(5):840-53. https://doi.org/10.1109/TPAMI.2007.1058

[2] Auricchio G, Codotti M, Scarselli F, Lodi A, Yoshiyasu Y. Computing Kantorovich-Wasserstein Distances on d-dimensional histograms using (d+1)-partite graphs. In Advances in Neural Information Processing Systems (NeurIPS). 2018;31. https://arxiv.org/abs/1805.07416

[rflamary]

Hello @tvercaut , this is very interesting. I wanted to add grid based solver (beyond the convolutional sinkhorn we have) for a while. If you can proposed a, exact OT variant using the POT verision of the lemon solver we are definitely interested.

I think it woudl be nice to have a generic solver function like ot.solve_grid(A,B) that can solve the exact and Entropic version with ot.solve_grid(A,B,reg=10) and that could be a nice wrapper for your solver.

Honnestly I'm quite swamped at the moment so code reviews can take some time but we can discuss API if you have questions.

[tvercaut]

Turns out a first PR wasn't that much work (for me at least, still a fair bit of coding agent leg work):
#863
I guess follow-up discussion can happen there.

[tvercaut]

Thanks. I'll see if I can carve out a bit of time to try that. It may take me a while as well.

Re convolutional Sinkhorn based EMD compuation on grids: Sorry I can't find it in the documentation. What is the function name?

One question already: Would support for backpropagation be needed? I haven't looked at it yet. I expect we can still use the envelope theorem for the gradient of the EMD value. Getting a (sub-)gradient for the transportation plan is probably more involved.

[rflamary]

You are right there is no OT solver for grids I miced things but we do have an implementation for 2D barcenters with a similar bergman projection algo rith. If we have the exact solver on grid we should also implement the entropic solver:
https://pythonot.github.io/gen_modules/ot.bregman.html#ot.bregman.convolutional_barycenter2d

backpropagation would be nice but only on the objective wrt the weights (so that we can do iterative solvers for barycenetrs for instance). No need to for wrt the OT plan we have entropic OT for that.

I saw your PR it looks very interesting.