Preserve SSIM gradients through zero covariance

dm_pix/_src/metrics.py

@@ -278,9 +278,8 @@ def filter_fn_x(z):
   epsilon = jnp.finfo(jnp.float32).eps ** 2
   sigma00 = jnp.maximum(epsilon, sigma00)
   sigma11 = jnp.maximum(epsilon, sigma11)
-  sigma01 = jnp.sign(sigma01) * jnp.minimum(
-      jnp.sqrt(sigma00 * sigma11), jnp.abs(sigma01)
-  )
+  covariance_bound = jnp.sqrt(sigma00 * sigma11)
+  sigma01 = jnp.clip(sigma01, -covariance_bound, covariance_bound)
 
   c1 = (k1 * max_val) ** 2
   c2 = (k2 * max_val) ** 2

dm_pix/_src/metrics_test.py

@@ -15,9 +15,11 @@
 import os
 
 from absl.testing import absltest
+from absl.testing import parameterized
 import chex
 from dm_pix._src import metrics
 import jax
+import jax.numpy as jnp
 import numpy as np
 import tensorflow as tf
 
@@ -206,6 +208,112 @@ def test_ssim_ignore_nans(self):
     assert not np.any(np.isnan(ssim))
 
 
+def _mean_filter(x):
+  return jnp.mean(x, axis=(-3, -2), keepdims=True)
+
+
+def _smooth_ssim(a, b):
+  ma, mb = _mean_filter(a), _mean_filter(b)
+  va, vb = _mean_filter((a - ma) ** 2), _mean_filter((b - mb) ** 2)
+  covariance = _mean_filter((a - ma) * (b - mb))
+  # These fixtures have valid covariance, so no clipping is required.
+  result = (
+      (2 * ma * mb + 0.01**2)
+      * (2 * covariance + 0.03**2)
+      / ((ma**2 + mb**2 + 0.01**2) * (va + vb + 0.03**2))
+  )
+  return result.mean(axis=(-3, -2, -1))
+
+
+def _ssim_gradient_images(batch=False, channels=1):
+  a = jnp.array([[0.75, 0.75], [0.25, 0.25]])[..., None]
+  b = jnp.array([[0.75, 0.25], [0.75, 0.25]])[..., None]
+  a, b = jnp.repeat(a, channels, -1), jnp.repeat(b, channels, -1)
+  if batch:
+    a, b = jnp.stack([a, a]), jnp.stack([b, b])
+  return a, b
+
+
+class SsimCovarianceGradientTest(parameterized.TestCase):
+
+  @parameterized.product(
+      batch=[False, True], channels=[1, 3], custom_filter=[False, True]
+  )
+  def test_zero_covariance_gradient_matches_independent_formula(
+      self, batch, channels, custom_filter
+  ):
+    a, b = _ssim_gradient_images(batch, channels)
+    kwargs = (
+        {"filter_fn": _mean_filter}
+        if custom_filter
+        else {"filter_size": 2, "filter_sigma": 1.0}
+    )
+    objective = lambda x: metrics.ssim(a, x, **kwargs).sum()
+    reference = lambda x: _smooth_ssim(a, x).sum()
+    expected = jax.value_and_grad(reference)(b)
+    for fn in (
+        jax.value_and_grad(objective),
+        jax.jit(jax.value_and_grad(objective)),
+    ):
+      actual = fn(b)
+      for value, target in zip(actual, expected):
+        np.testing.assert_allclose(value, target, rtol=3e-6, atol=1e-7)
+    _, tangent = jax.jvp(objective, (b,), (a - 0.5,))
+    step = 0.01
+    finite_diff = (
+        objective(b + step * (a - 0.5)) - objective(b - step * (a - 0.5))
+    ) / (2 * step)
+    np.testing.assert_allclose(tangent, finite_diff, rtol=5e-4, atol=1e-6)
+
+  def test_default_gaussian_passes_gradients_at_a_zero_second_image(self):
+    a, b = _ssim_gradient_images()
+    b = jnp.zeros_like(b)
+    actual = jax.jit(
+        jax.value_and_grad(lambda x: metrics.ssim(a, x, filter_size=2))
+    )(b)
+    expected = jax.value_and_grad(lambda x: _smooth_ssim(a, x))(b)
+    for value, target in zip(actual, expected):
+      np.testing.assert_allclose(value, target, rtol=3e-6, atol=1e-8)
+
+  @parameterized.parameters(-0.5, 0.5)
+  def test_nonzero_covariance_values_and_gradients_are_preserved(self, weight):
+    a, b = _ssim_gradient_images()
+    b = b + weight * (a - 0.5)
+    actual = jax.value_and_grad(
+        lambda x: metrics.ssim(a, x, filter_fn=_mean_filter)
+    )(b)
+    expected = jax.value_and_grad(lambda x: _smooth_ssim(a, x))(b)
+    for value, target in zip(actual, expected):
+      np.testing.assert_allclose(value, target, rtol=3e-6, atol=1e-7)
+
+  def test_correlation_parameter_can_learn_from_zero(self):
+    a, b = _ssim_gradient_images()
+    score = lambda parameter: metrics.ssim(
+        a, b + parameter * (a - 0.5), filter_size=2, filter_sigma=1.0
+    )
+    parameter = jnp.array(0.0)
+    derivative = jax.jit(jax.grad(score))(parameter)
+    self.assertGreater(float(derivative), 0.9)
+    self.assertGreater(
+        float(score(parameter + 0.1 * derivative)), float(score(parameter))
+    )
+
+  @parameterized.parameters(False, True)
+  def test_map_and_identity_scores_keep_existing_values(self, negative):
+    a, _ = _ssim_gradient_images(batch=True)
+    b = 1 - a if negative else a
+    actual = metrics.ssim(a, b, filter_size=2, return_map=True)
+    expected = _smooth_ssim(a, b)
+    np.testing.assert_allclose(
+        actual.mean(axis=(-3, -2, -1)), expected, rtol=3e-6, atol=1e-7
+    )
+    self.assertEqual(actual.shape, (2, 1, 1, 1))
+
+
+if __name__ == "__main__":
+  absltest.main()
+
+
 if __name__ == "__main__":
   os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
   absltest.main()