diff --git a/RELEASES.md b/RELEASES.md index 063701229..96204df9c 100644 --- a/RELEASES.md +++ b/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) diff --git a/ot/stochastic.py b/ot/stochastic.py index 6df3ad771..8cbfe8f5e 100644 --- a/ot/stochastic.py +++ b/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 diff --git a/test/test_stochastic.py b/test/test_stochastic.py index 911d52b92..85e5a7701 100644 --- a/test/test_stochastic.py +++ b/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