Stabilize Gaussian mixture backward passes

pyro/distributions/diag_normal_mixture.py

@@ -202,23 +202,23 @@ def backward(ctx, grad_output):
         root_two = math.sqrt(2.0)
         shift_log_scales = log_scales[..., shift_indices]
         shift_log_scales[..., 0] = 0.0
-        sigma_products = torch.cumsum(shift_log_scales, dim=-1).exp()  # b j i
+        log_sigma_products = torch.cumsum(shift_log_scales, dim=-1)  # b j i
 
         reverse_indices = torch.tensor(
             range(dim - 1, -1, -1), dtype=torch.long, device=z.device
         )
         reverse_log_sigma_0 = sigma_0.log()[..., reverse_indices]  # b 1 i
-        sigma_0_products = torch.cumsum(reverse_log_sigma_0, dim=-1).exp()[
+        log_sigma_0_products = torch.cumsum(reverse_log_sigma_0, dim=-1)[
             ..., reverse_indices - 1
         ]  # b 1 i
-        sigma_0_products[..., -1] = 1.0
-        sigma_products *= sigma_0_products
+        log_sigma_0_products[..., -1] = 0.0
+        log_sigma_products += log_sigma_0_products
 
         logits_grad = torch.erf(z_tilde / root_two) - torch.erf(
             z_shift / root_two
         )  # l b j i
-        logits_grad *= torch.exp(-0.5 * r_sqr_ji)  # l b j i
-        logits_grad = (logits_grad * g / sigma_products).sum(-1)  # l b j
+        logits_grad *= torch.exp(-0.5 * r_sqr_ji - log_sigma_products)  # l b j i
+        logits_grad = (logits_grad * g).sum(-1)  # l b j
         logits_grad = sum_leftmost(logits_grad / q_tot, -1 - batch_dims)  # b j
         logits_grad *= 0.5 * math.pow(2.0 * math.pi, -0.5 * (dim - 1))
         logits_grad = -pis * logits_grad

pyro/distributions/diag_normal_mixture_shared_cov.py

@@ -176,25 +176,24 @@ def backward(ctx, grad_output):
         log_q_j = log_qs.sum(-1, keepdim=True)  # l b j 1
         log_q_j_max = torch.max(log_q_j, -2, keepdim=True)[0]
         q_j_prime = torch.exp(log_q_j - log_q_j_max)  # l b j 1
-        q_j = torch.exp(log_q_j)  # l b j 1
 
-        q_tot = (pis.unsqueeze(-1) * q_j).sum(-2)  # l b 1
         q_tot_prime = (pis.unsqueeze(-1) * q_j_prime).sum(-2).unsqueeze(-1)  # l b 1 1
 
         root_two = math.sqrt(2.0)
         mu_ll_ba = torch.transpose(mu_ll_ab, -1, -2)
         logits_grad = torch.erf((z_ll_ab - mu_ll_ab) / root_two) - torch.erf(
             (z_ll_ab + mu_ll_ba) / root_two
         )
-        logits_grad *= torch.exp(-0.5 * z_perp_ab_sqr)  # l b k j
+        log_q_j_max = log_q_j_max.squeeze(-1).unsqueeze(-1)  # l b 1 1
+        logits_grad *= torch.exp(-0.5 * z_perp_ab_sqr - log_q_j_max)  # l b k j
 
         #                 bi      lbi                               bkji
         mu_ab_sigma_g = ((coord_scale * g).unsqueeze(-2).unsqueeze(-2) * mu_ab).sum(
             -1
         )  # l b k j
         logits_grad *= -mu_ab_sigma_g * pis.unsqueeze(-2)  # l b k j
         logits_grad = pis * sum_leftmost(
-            logits_grad.sum(-1) / q_tot, -(1 + batch_dims)
+            logits_grad.sum(-1) / q_tot_prime.squeeze(-1), -(1 + batch_dims)
         )  # b k
         logits_grad *= math.sqrt(0.5 * math.pi)
 

pyro/distributions/gaussian_scale_mixture.py

@@ -151,21 +151,34 @@ def backward(ctx, grad_output):
         epsilons_sqr = torch.pow(epsilons, 2.0)  # l i
         r_sqr = epsilons_sqr.sum(-1, keepdim=True)  # l
         r_sqr_j = r_sqr / component_scale_sqr  # l j
-        coord_scale_product = coord_scale.prod()
-        component_scale_power = torch.pow(component_scale, float(dim))
+        log_coord_scale_product = coord_scale.log().sum()
+        log_component_scale_power = component_scale.log() * float(dim)
+
+        log_q_j = (
+            -0.5 * r_sqr_j
+            - 0.5 * math.log(2.0 * math.pi) * float(dim)
+            - log_coord_scale_product
+            - log_component_scale_power
+        )
+        log_q_tot = torch.logsumexp(pis.log() + log_q_j, dim=-1, keepdim=True)
+        q_j_over_q_tot = torch.exp(log_q_j - log_q_tot)
+
+        log_normalizer = (
+            -0.5 * math.log(2.0 * math.pi) * float(dim)
+            - log_component_scale_power
+            - log_coord_scale_product
+            - log_q_tot
+        )
 
-        q_j = torch.exp(-0.5 * r_sqr_j) / math.pow(
-            2.0 * math.pi, 0.5 * float(dim)
-        )  # l j
-        q_j /= coord_scale_product * component_scale_power  # l j
-        q_tot = (pis * q_j).sum(-1, keepdim=True)  # l
+        Phi_j = torch.exp(-0.5 * r_sqr_j + log_normalizer)
 
-        Phi_j = torch.exp(-0.5 * r_sqr_j)  # l j
-        exponents = -torch.arange(1.0, int(dim / 2) + 1.0, 1.0)
+        exponents = -torch.arange(
+            1.0, int(dim / 2) + 1.0, 1.0, device=z.device, dtype=z.dtype
+        )
         if z.dim() > 1:
-            r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, -1, int(dim / 2))  # l j d/2
+            r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, -1, int(dim / 2))
         else:
-            r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, int(dim / 2))  # l j d/2
+            r_j_poly = r_sqr_j.unsqueeze(-1).expand(-1, int(dim / 2))
         r_j_poly = coeffs * torch.pow(r_j_poly, exponents)
         Phi_j *= r_j_poly.sum(-1)
         if dim % 2 == 1:
@@ -174,20 +187,19 @@ def backward(ctx, grad_output):
                 coeffs[-1]
                 * math.sqrt(0.5 * math.pi)
                 * (1.0 - torch.erf(r_sqr_j.sqrt() / root_two))
-            )  # l j
-            Phi_j += extra_term * torch.pow(r_sqr_j, -0.5 * float(dim))
-
-        logits_grad = (z.unsqueeze(-2) * Phi_j.unsqueeze(-1) * g).sum(-1)  # l j
-        logits_grad /= q_tot
-        logits_grad = sum_leftmost(logits_grad, -1) * math.pow(
-            2.0 * math.pi, -0.5 * float(dim)
-        )
-        logits_grad = pis * logits_grad / (component_scale_power * coord_scale_product)
+            )
+            Phi_j += (
+                extra_term
+                * torch.pow(r_sqr_j, -0.5 * float(dim))
+                * torch.exp(log_normalizer)
+            )
+
+        logits_grad = (z.unsqueeze(-2) * Phi_j.unsqueeze(-1) * g).sum(-1)
+        logits_grad = sum_leftmost(logits_grad, -1)
+        logits_grad = pis * logits_grad
         logits_grad = logits_grad - logits_grad.sum() * pis
 
-        prefactor = (
-            pis.unsqueeze(-1) * q_j.unsqueeze(-1) * g / q_tot.unsqueeze(-1)
-        )  # l j i
+        prefactor = pis.unsqueeze(-1) * q_j_over_q_tot.unsqueeze(-1) * g  # l j i
         coord_scale_grad = sum_leftmost(prefactor * epsilons.unsqueeze(-2), -1)
         component_scale_grad = sum_leftmost(
             (prefactor * z.unsqueeze(-2)).sum(-1) / component_scale, -1

tests/distributions/test_gaussian_mixtures.py

@@ -211,6 +211,81 @@ def test_gsm_log_prob():
     )
 
 
+def test_mix_of_diag_normals_backward_small_scales_finite_grads():
+    torch.manual_seed(0)
+
+    K = 5
+    D = 50
+
+    locs = torch.randn(K, D, requires_grad=True)
+    coord_scale = torch.full((K, D), 0.1, dtype=torch.float32, requires_grad=True)
+    component_logits = torch.zeros(K, requires_grad=True)
+
+    dist = MixtureOfDiagNormals(locs, coord_scale, component_logits)
+    z = dist.rsample(sample_shape=(32,))
+
+    assert torch.isfinite(z).all()
+
+    loss = z.pow(2).sum()
+    assert torch.isfinite(loss)
+
+    loss.backward()
+
+    assert torch.isfinite(locs.grad).all()
+    assert torch.isfinite(coord_scale.grad).all()
+    assert torch.isfinite(component_logits.grad).all()
+
+
+def test_gsm_backward_small_coord_scale_finite_grads():
+    torch.manual_seed(0)
+
+    K = 5
+    D = 50
+
+    coord_scale = torch.full((D,), 0.1, requires_grad=True)
+    component_logits = torch.zeros(K, requires_grad=True)
+    component_scale = torch.ones(K, dtype=torch.float32, requires_grad=True)
+
+    dist = GaussianScaleMixture(coord_scale, component_logits, component_scale)
+    z = dist.rsample(sample_shape=(32,))
+
+    assert torch.isfinite(z).all()
+
+    loss = z.pow(2).sum()
+    assert torch.isfinite(loss)
+
+    loss.backward()
+
+    assert torch.isfinite(coord_scale.grad).all()
+    assert torch.isfinite(component_logits.grad).all()
+    assert torch.isfinite(component_scale.grad).all()
+
+
+def test_mix_of_diag_normals_shared_cov_backward_high_dim_finite_grads():
+    torch.manual_seed(0)
+
+    K = 3
+    D = 300
+
+    locs = torch.randn(K, D, requires_grad=True)
+    coord_scale = torch.ones(D, dtype=torch.float32, requires_grad=True)
+    component_logits = torch.zeros(K, requires_grad=True)
+
+    dist = MixtureOfDiagNormalsSharedCovariance(locs, coord_scale, component_logits)
+    z = dist.rsample(sample_shape=(4,))
+
+    assert torch.isfinite(z).all()
+
+    loss = z.pow(2).sum()
+    assert torch.isfinite(loss)
+
+    loss.backward()
+
+    assert torch.isfinite(locs.grad).all()
+    assert torch.isfinite(coord_scale.grad).all()
+    assert torch.isfinite(component_logits.grad).all()
+
+
 @pytest.mark.parametrize("batch_size", [1, 3])
 def test_mix_of_diag_normals_log_prob(batch_size):
     sigmas = torch.tensor([[2.0, 1.5], [1.5, 2.0]])