[MRG] Shift before exponentiating in coordinate_grad_semi_dual (closes #264)

RELEASES.md

@@ -15,7 +15,8 @@
 
 #### Closed issues
 
-- Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #860)
+- Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264)
+- Remove a leftover debug `print` from `ot.utils.projection_sparse_simplex` with `axis=1`, and make the `ot.datasets.make_gauss_hd` docstring a raw string so importing `ot` no longer emits a `SyntaxWarning` (PR #866)
 - Fix `ot.dist` ignoring the weights `w` for `metric="cityblock"`, which returned the unweighted distance although the weights are documented for this metric (PR #859)
 - Fix swapped arguments to `div_to_product` in `ot.gromov.fused_unbalanced_across_spaces_cost`: with `reg_type="independent"` (UCOOT) the entropic terms used the plan marginals as the reference measures and vice versa (PR #855, Issue #854)
 - Fix device placement in `ot.batch.bregman_projection_batch` so `ot.solve_batch(..., method="sinkhorn")` no longer crashes on GPU when the torch default device is CPU (PR #851)

ot/stochastic.py

@@ -67,7 +67,10 @@ def coordinate_grad_semi_dual(b, M, reg, beta, i):
         Advances in Neural Information Processing Systems (2016).
     """
     r = M[i, :] - beta
-    exp_beta = np.exp(-r / reg) * b
+    # shift before exponentiating: the factor cancels in the ratio below, but
+    # without it exp overflows once beta grows and khi becomes all-NaN
+    min_r = np.min(r)
+    exp_beta = np.exp(-(r - min_r) / reg) * b
     khi = exp_beta / (np.sum(exp_beta))
     return b - khi
 

test/test_stochastic.py

@@ -13,6 +13,8 @@
 #
 # License: MIT License
 
+import warnings
+
 import numpy as np
 import ot
 
@@ -87,6 +89,49 @@ def test_stochastic_asgd():
 # regularization term, a learning rate and a number of iteration
 
 
+def test_coordinate_grad_semi_dual_large_beta():
+    """Non-regression test for issue #264.
+
+    The semi-dual gradient exponentiates -(M[i] - beta) / reg. Once beta grows
+    that overflows and the normalised khi becomes all-NaN, so the shift by the
+    minimum is what keeps the solver usable at large scale.
+    """
+    rng = np.random.RandomState(0)
+    n_source, n_target, reg = 500, 60, 1
+    M = ot.dist(rng.randn(n_source, 2), rng.randn(n_target, 2))
+    b = ot.utils.unif(n_target)
+
+    with warnings.catch_warnings():
+        # narrow: unrelated warnings elsewhere must not make this brittle
+        warnings.simplefilter("error", RuntimeWarning)
+        grad = ot.stochastic.coordinate_grad_semi_dual(
+            b, M, reg, np.full(n_target, 800.0), 0
+        )
+
+    assert not np.any(np.isnan(grad))
+    # khi lives on the simplex, so the gradient b - khi sums to zero
+    np.testing.assert_allclose(np.sum(grad), 0.0, atol=1e-12)
+
+
+def test_coordinate_grad_semi_dual_unchanged_without_overflow():
+    """The shift cancels in the ratio, so benign inputs must be unaffected."""
+    rng = np.random.RandomState(0)
+    n_source, n_target, reg = 200, 40, 1
+    M = ot.dist(rng.randn(n_source, 2), rng.randn(n_target, 2))
+    b = ot.utils.unif(n_target)
+
+    for scale in (0.0, 0.5, 2.0):
+        beta = rng.randn(n_target) * scale
+        r = M[0, :] - beta
+        expected = b - (np.exp(-r / reg) * b) / np.sum(np.exp(-r / reg) * b)
+        np.testing.assert_allclose(
+            ot.stochastic.coordinate_grad_semi_dual(b, M, reg, beta, 0),
+            expected,
+            rtol=1e-12,
+            atol=1e-14,
+        )
+
+
 def test_sag_asgd_sinkhorn():
     # test all algorithms
     n = 10