Instability with GW

[ayushjain1144]

Instability with GW

[michalk8]

Hi @ayushjain1144 ,

I noticed when running your notebook that even though you requested float64, it was still truncated to float32. At start-up, you can do:

import jax
jax.config.update("jax_enable_x64", True)

See also JAX's sharp bits for other options.

You're also using very low (1e-6) epsilon values, so on float32, so I expect there to be numerical imprecision.
Changing it to float64 in your notebook, I get mostly 1.0 accuracy, see the output below:

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

9 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 1e-06, accuracy: 1.0,  cost: 2.9970166200219775e-06
8 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 3e-06, accuracy: 1.0,  cost: 8.989837444618229e-06
8 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 5e-06, accuracy: 1.0,  cost: 1.4999405788748445e-05
8 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 7e-06, accuracy: 1.0,  cost: 2.103736864386557e-05
8 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 9e-06, accuracy: 1.0,  cost: 2.7097978185330104e-05
8 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 1e-05, accuracy: 1.0,  cost: 3.013191393879788e-05
7 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 3e-05, accuracy: 1.0,  cost: 8.99370951687826e-05
6 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 5e-05, accuracy: 1.0,  cost: 0.00014800605952555742
5 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 7e-05, accuracy: 1.0,  cost: 0.00020506084397681518
5 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 9e-05, accuracy: 1.0,  cost: 0.0002617191301163146
5 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: False
The final regularised GW cost is: 0.000
eps: 0.0001, accuracy: 1.0,  cost: 0.0002899387715525633
4 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: True
The final regularised GW cost is: 0.001
eps: 0.0005, accuracy: 0.9500000000000001,  cost: 0.001426522237415696
4 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: True
The final regularised GW cost is: 0.003
eps: 0.001, accuracy: 0.9500000000000001,  cost: 0.0027998761460317242
4 outer iterations were needed.
The last Sinkhorn iteration has converged: True
The outer loop of Gromov Wasserstein has converged: True
The final regularised GW cost is: 0.013
eps: 0.005, accuracy: 0.9500000000000001,  cost: 0.012696349927582307

[ayushjain1144]

Oh I see, my bad! Thank you so much! Closing the issue.