New neural OT solver ENOT (ott-jax#503)

* ExpectileNeuralDual method * nb update * nb update * dim fix * nb update * nb update * Update expectile_neural_dual * Update expectile_neural_dual * sint fix * sint fix * sint fix * sint fix * bk * Update enot nb * 123 * enot pot * nb update * 2nd rew * lint fix * merge main * Minor documentation polish
mahdi-shafiei · Jun 28, 2024 · 8aa6810 · 8aa6810
1 parent 2147cbe
commit 8aa6810
Show file tree

Hide file tree

Showing 9 changed files with 1,474 additions and 2 deletions.
diff --git a/docs/neural/methods.rst b/docs/neural/methods.rst
@@ -18,6 +18,8 @@ Neural Dual
     :toctree: _autosummary
 
     neuraldual.W2NeuralDual
+    expectile_neural_dual.ExpectileNeuralDual
+    expectile_neural_dual.ENOTPotentials
 
 ott.neural.methods.flows
 ========================

diff --git a/docs/neural/networks.rst b/docs/neural/networks.rst
@@ -12,6 +12,7 @@ Networks
     velocity_field.VelocityField
     potentials.BasePotential
     potentials.PotentialMLP
+    potentials.MLP
     potentials.PotentialTrainState
 
 

diff --git a/docs/references.bib b/docs/references.bib
@@ -916,6 +916,15 @@ @inproceedings{cho:09
   year      = {2009},
 }
 
+@misc{buzun:24,
+  title         = {ENOT: Expectile Regularization for Fast and Accurate Training of Neural Optimal Transport},
+  author        = {Nazar Buzun and Maksim Bobrin and Dmitry V. Dylov},
+  year          = {2024},
+  eprint        = {2403.03777},
+  archivePrefix = {arXiv},
+  primaryClass  = {cs.LG}
+}
+
 @misc{vesseron:24,
   author      = {Vesseron, Nina and Cuturi, Marco},
   eprint      = {2403.03071},

diff --git a/docs/spelling/misc.txt b/docs/spelling/misc.txt
@@ -3,6 +3,7 @@ Utils
 alg
 arg
 args
+coef
 cond
 dtype
 eps
@@ -16,6 +17,7 @@ gmm
 gv
 init
 iter
+iters
 jnp
 jvp
 kwargs

diff --git a/docs/spelling/technical.txt b/docs/spelling/technical.txt
@@ -7,6 +7,7 @@ DTW
 Danskin
 Datasets
 Dykstra
+Expectile
 Fenchel
 Frobenius
 Gangbo
@@ -64,6 +65,7 @@ elementwise
 embeddings
 entropic
 epigenetic
+expectile
 featurized
 grayscale
 heterogeneous

diff --git a/docs/tutorials/ENOT.ipynb b/docs/tutorials/ENOT.ipynb
diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst
@@ -58,3 +58,4 @@ Neural Optimal Transport
     icnn_inits
     MetaOT
     Monge_Gap
+    ENOT
-Original file line number
+Diff line change
@@ Expand Up / @@ -3,6 +3,7 @@ Utils @@
     alg
     arg
     args
+    coef
     cond
     dtype
     eps
@@ Expand All / @@ -16,6 +17,7 @@ gmm @@
     gv
     init
     iter
+    iters
     jnp
     jvp
     kwargs
@@ Expand Down @@