Repository: pyro-ppl/pyro

Issue:
O(obs_dim) algorithms for GaussianHMM methods
@martinjankowiak is trying to forecast with a `GaussianHMM` likelihood with large `obs_dim` and small `hidden_dim` however the time complexity is awful for forecasting (it's fine for training). We could solve this by implementing `AffineNormal.event_permute()`, `.condition()` and `.rsample()` as  used in `GaussianHMM.rsample()` [code](https://github.com/pyro-ppl/pyro/blob/fca6c9546e57b31e5d327ecb7fd9f6340ac3dc66/pyro/distributions/hmm.py#L514).

When I implemented those `AffineNormal` method I took a correct but expensive shortcut and simply converted to a `Gaussian` and now we need to pay off that technical debt e.g. [here](https://github.com/pyro-ppl/pyro/blob/fca6c9546e57b31e5d327ecb7fd9f6340ac3dc66/pyro/ops/gaussian.py#L302).

We should test the solution with a local test (or a `@pytest.mark.skip`'d test) with say `obs_dim=5000`, `hidden_dim=1` to make sure all relevant ops are cheap. It's tough to determine just by reading code.


Discussion:

Resolving patch:
diff --git a/pyro/distributions/hmm.py b/pyro/distributions/hmm.py
--- a/pyro/distributions/hmm.py
+++ b/pyro/distributions/hmm.py
@@ -504,14 +504,10 @@ def log_prob(self, value):
     def rsample(self, sample_shape=torch.Size()):
         assert self.duration is not None
         sample_shape = torch.Size(sample_shape)
-        obs_dim = self.obs_dim
-        hidden_dim = self.hidden_dim
         trans = self._trans + self._obs.marginalize(right=self.obs_dim).event_pad(left=self.hidden_dim)
         trans = trans.expand(trans.batch_shape[:-1] + (self.duration,))
         z = _sequential_gaussian_filter_sample(self._init, trans, sample_shape)
-        perm = torch.cat([torch.arange(hidden_dim, hidden_dim + obs_dim, device=z.device),
-                          torch.arange(hidden_dim, device=z.device)])
-        x = self._obs.event_permute(perm).condition(z).rsample()
+        x = self._obs.left_condition(z).rsample()
         return x
 
     def rsample_posterior(self, value, sample_shape=torch.Size()):
diff --git a/pyro/ops/gaussian.py b/pyro/ops/gaussian.py
--- a/pyro/ops/gaussian.py
+++ b/pyro/ops/gaussian.py
@@ -161,9 +161,11 @@ def condition(self, value):
             g.log_density(x) == g.condition(right).log_density(left)
         """
         assert isinstance(value, torch.Tensor)
-        assert value.size(-1) <= self.info_vec.size(-1)
+        right = value.size(-1)
+        dim = self.dim()
+        assert right <= dim
 
-        n = self.dim() - value.size(-1)
+        n = dim - right
         info_a = self.info_vec[..., :n]
         info_b = self.info_vec[..., n:]
         P_aa = self.precision[..., :n, :n]
@@ -178,6 +180,30 @@ def condition(self, value):
                           b.mul(info_b).sum(-1))
         return Gaussian(log_normalizer, info_vec, precision)
 
+    def left_condition(self, value):
+        """
+        Condition this Gaussian on a leading subset of its state.
+        This should satisfy::
+
+            g.condition(y).dim() == g.dim() - y.size(-1)
+
+        Note that since this is a non-normalized Gaussian, we include the
+        density of ``y`` in the result. Thus :meth:`condition` is similar to a
+        ``functools.partial`` binding of arguments::
+
+            left = x[..., :n]
+            right = x[..., n:]
+            g.log_density(x) == g.left_condition(left).log_density(right)
+        """
+        assert isinstance(value, torch.Tensor)
+        left = value.size(-1)
+        dim = self.dim()
+        assert left <= dim
+
+        perm = torch.cat([torch.arange(left, dim, device=value.device),
+                          torch.arange(left, device=value.device)])
+        return self.event_permute(perm).condition(value)
+
     def marginalize(self, left=0, right=0):
         """
         Marginalizing out variables on either side of the event dimension::
@@ -251,33 +277,60 @@ class AffineNormal:
     """
     def __init__(self, matrix, loc, scale):
         assert loc.shape == scale.shape
+        assert matrix.shape[:-2] == loc.shape[:-1]
+        assert matrix.size(-1) == loc.size(-1)
         self.matrix = matrix
         self.loc = loc
         self.scale = scale
+        self._gaussian = None
 
-    @property
+    @lazy_property
     def batch_shape(self):
         return self.matrix.shape[:-2]
 
     def condition(self, value):
+        if value.size(-1) == self.loc.size(-1):
+            prec_sqrt = self.matrix / self.scale.unsqueeze(-2)
+            precision = matmul(prec_sqrt, prec_sqrt.transpose(-1, -2))
+            delta = (value - self.loc) / self.scale
+            info_vec = matvecmul(prec_sqrt, delta)
+            log_normalizer = (-0.5 * self.loc.size(-1) * math.log(2 * math.pi)
+                              - 0.5 * delta.pow(2).sum(-1) - self.scale.log().sum(-1))
+            return Gaussian(log_normalizer, info_vec, precision)
+        else:
+            return self.to_gaussian().condition(value)
+
+    def left_condition(self, value):
+        """
+        If ``value.size(-1) == x_dim``, this returns a Normal distribution with
+        ``event_dim=1``. After applying this method, the cost to draw a sample is
+        ``O(y_dim)`` instead of ``O(y_dim ** 3)``.
         """
-        Condition on a ``Y`` value.
+        if value.size(-1) == self.matrix.size(-2):
+            loc = matvecmul(self.matrix.transpose(-1, -2), value) + self.loc
+            matrix = value.new_zeros(loc.shape[:-1] + (0, loc.size(-1)))
+            scale = self.scale.expand(loc.shape)
+            return AffineNormal(matrix, loc, scale)
+        else:
+            return self.to_gaussian().left_condition(value)
 
-        :param torch.Tensor value: A value of ``Y``.
-        :return Gaussian: A gaussian likelihood over ``X``.
+    def rsample(self, sample_shape=torch.Size()):
         """
-        assert value.size(-1) == self.loc.size(-1)
-        prec_sqrt = self.matrix / self.scale.unsqueeze(-2)
-        precision = matmul(prec_sqrt, prec_sqrt.transpose(-1, -2))
-        delta = (value - self.loc) / self.scale
-        info_vec = matvecmul(prec_sqrt, delta)
-        log_normalizer = (-0.5 * self.loc.size(-1) * math.log(2 * math.pi)
-                          - 0.5 * delta.pow(2).sum(-1) - self.scale.log().sum(-1))
-        return Gaussian(log_normalizer, info_vec, precision)
+        Reparameterized sampler.
+        """
+        if self.matrix.size(-2) > 0:
+            raise NotImplementedError
+        shape = sample_shape + self.batch_shape + self.loc.shape[-1:]
+        noise = torch.randn(shape, dtype=self.loc.dtype, device=self.loc.device)
+        return self.loc + noise * self.scale
 
     def to_gaussian(self):
-        mvn = torch.distributions.MultivariateNormal(self.loc, scale_tril=self.scale.diag_embed())
-        return matrix_and_mvn_to_gaussian(self.matrix, mvn)
+        if self._gaussian is None:
+            mvn = torch.distributions.Independent(
+                torch.distributions.Normal(self.loc, scale=self.scale), 1)
+            y_gaussian = mvn_to_gaussian(mvn)
+            self._gaussian = _matrix_and_gaussian_to_gaussian(self.matrix, y_gaussian)
+        return self._gaussian
 
     def expand(self, batch_shape):
         matrix = self.matrix.expand(batch_shape + self.matrix.shape[-2:])
@@ -305,7 +358,14 @@ def __add__(self, other):
         return self.to_gaussian() + other
 
     def marginalize(self, left=0, right=0):
-        return self.to_gaussian().marginalize(left, right)
+        if left == 0 and right == self.loc.size(-1):
+            n = self.matrix.size(-2)
+            precision = self.scale.new_zeros(self.batch_shape + (n, n))
+            info_vec = self.scale.new_zeros(self.batch_shape + (n,))
+            log_normalizer = self.scale.new_zeros(self.batch_shape)
+            return Gaussian(log_normalizer, info_vec, precision)
+        else:
+            return self.to_gaussian().marginalize(left, right)
 
 
 def mvn_to_gaussian(mvn):
@@ -337,6 +397,23 @@ def mvn_to_gaussian(mvn):
     return Gaussian(log_normalizer, info_vec, precision)
 
 
+def _matrix_and_gaussian_to_gaussian(matrix, y_gaussian):
+    P_yy = y_gaussian.precision
+    neg_P_xy = matmul(matrix, P_yy)
+    P_xy = -neg_P_xy
+    P_yx = P_xy.transpose(-1, -2)
+    P_xx = matmul(neg_P_xy, matrix.transpose(-1, -2))
+    precision = torch.cat([torch.cat([P_xx, P_xy], -1),
+                           torch.cat([P_yx, P_yy], -1)], -2)
+    info_y = y_gaussian.info_vec
+    info_x = -matvecmul(matrix, info_y)
+    info_vec = torch.cat([info_x, info_y], -1)
+    log_normalizer = y_gaussian.log_normalizer
+
+    result = Gaussian(log_normalizer, info_vec, precision)
+    return result
+
+
 def matrix_and_mvn_to_gaussian(matrix, mvn):
     """
     Convert a noisy affine function to a Gaussian. The noisy affine function is defined as::
@@ -363,19 +440,7 @@ def matrix_and_mvn_to_gaussian(matrix, mvn):
         return AffineNormal(matrix, mvn.base_dist.loc, mvn.base_dist.scale)
 
     y_gaussian = mvn_to_gaussian(mvn)
-    P_yy = y_gaussian.precision
-    neg_P_xy = matmul(matrix, P_yy)
-    P_xy = -neg_P_xy
-    P_yx = P_xy.transpose(-1, -2)
-    P_xx = matmul(neg_P_xy, matrix.transpose(-1, -2))
-    precision = torch.cat([torch.cat([P_xx, P_xy], -1),
-                           torch.cat([P_yx, P_yy], -1)], -2)
-    info_y = y_gaussian.info_vec
-    info_x = -matvecmul(matrix, info_y)
-    info_vec = torch.cat([info_x, info_y], -1)
-    log_normalizer = y_gaussian.log_normalizer
-
-    result = Gaussian(log_normalizer, info_vec, precision)
+    result = _matrix_and_gaussian_to_gaussian(matrix, y_gaussian)
     assert result.batch_shape == batch_shape
     assert result.dim() == x_dim + y_dim
     return result


Test patch:
diff --git a/tests/distributions/test_hmm.py b/tests/distributions/test_hmm.py
--- a/tests/distributions/test_hmm.py
+++ b/tests/distributions/test_hmm.py
@@ -334,6 +334,24 @@ def test_gaussian_hmm_shape(diag, init_shape, trans_mat_shape, trans_mvn_shape,
         assert d2.event_shape == (f, obs_dim)
 
 
+def test_gaussian_hmm_high_obs_dim():
+    hidden_dim = 1
+    obs_dim = 1000
+    duration = 10
+    sample_shape = (100,)
+    init_dist = random_mvn((), hidden_dim)
+    trans_mat = torch.randn((duration,) + (hidden_dim, hidden_dim))
+    trans_dist = random_mvn((duration,), hidden_dim)
+    obs_mat = torch.randn((duration,) + (hidden_dim, obs_dim))
+    loc = torch.randn((duration, obs_dim))
+    scale = torch.randn((duration, obs_dim)).exp()
+    obs_dist = dist.Normal(loc, scale).to_event(1)
+    d = dist.GaussianHMM(init_dist, trans_mat, trans_dist, obs_mat, obs_dist,
+                         duration=duration)
+    x = d.rsample(sample_shape)
+    assert x.shape == sample_shape + (duration, obs_dim)
+
+
 @pytest.mark.parametrize('sample_shape', [(), (5,)], ids=str)
 @pytest.mark.parametrize('batch_shape', [(), (4,), (3, 2)], ids=str)
 @pytest.mark.parametrize('obs_dim', [1, 2])
diff --git a/tests/ops/test_gaussian.py b/tests/ops/test_gaussian.py
--- a/tests/ops/test_gaussian.py
+++ b/tests/ops/test_gaussian.py
@@ -193,6 +193,14 @@ def test_condition(sample_shape, batch_shape, left, right):
     expected = gaussian.log_density(value)
     assert_close(actual, expected)
 
+    # test left_condition
+    permute_conditioned = gaussian.left_condition(left_value)
+    assert permute_conditioned.batch_shape == sample_shape + gaussian.batch_shape
+    assert permute_conditioned.dim() == right
+
+    permute_actual = permute_conditioned.log_density(right_value)
+    assert_close(permute_actual, expected)
+
 
 @pytest.mark.parametrize("batch_shape", [(), (4,), (3, 2)], ids=str)
 @pytest.mark.parametrize("dim", [1, 2, 3])
@@ -234,6 +242,18 @@ def test_affine_normal(batch_shape, x_dim, y_dim):
     assert_close(actual_like.info_vec, expected_like.info_vec)
     assert_close(actual_like.precision, expected_like.precision)
 
+    x = torch.randn(batch_shape + (x_dim,))
+    permute_actual = actual.left_condition(x)
+    assert isinstance(permute_actual, AffineNormal)
+    permute_actual = permute_actual.to_gaussian()
+
+    permute_expected = expected.left_condition(y)
+    assert isinstance(permute_expected, Gaussian)
+
+    assert_close(permute_actual.log_normalizer, permute_actual.log_normalizer)
+    assert_close(permute_actual.info_vec, permute_actual.info_vec)
+    assert_close(permute_actual.precision, permute_actual.precision)
+
 
 @pytest.mark.parametrize("sample_shape", [(), (7,), (6, 5)], ids=str)
 @pytest.mark.parametrize("batch_shape", [(), (4,), (3, 2)], ids=str)