ott icon indicating copy to clipboard operation
ott copied to clipboard

New neural OT solver ENOT

Open nazarblch opened this issue 1 year ago • 5 comments
trafficstars

This pull request includes a new method (solver) for Neural OT training, called ENOT. It solves the dual OT problem in d-dimensional Euclidean spaces with the use of specific expectile regularization on the Kantorovich potentials.

The method is similar to W2NeuralDual (already implemented in the repository) but extends its functionality with the ability to use different cost functions, and eliminates the need for computationally intensive fine-tuning with conjugate optimizers.

Our paper arXiv:2403.03777 provides a detailed description of the method and its experimental evaluation. ENOT outperforms previous state-of-the-art approaches on the Wasserstein-2 benchmark tasks by a large margin (up to a 3-fold improvement in quality and up to a 10-fold improvement in runtime).

We have also included in this pull request a detailed tutorial, demonstrating how our method works. Please find the source code with ENOT implementation and the tutorial with applications on some 2D datasets by the following links: https://github.com/nazarblch/ott/blob/dev/src/ott/neural/solvers/expectile_neural_dual.py https://github.com/nazarblch/ott/blob/dev/docs/tutorials/ENOT.ipynb

nazarblch avatar Mar 19 '24 08:03 nazarblch