Repository: pyro-ppl/numpyro

Issue:
Massive memory consumption by AutoLowRankMultivariateNormal quantiles
Hi Numpyro team!

I am trying to `AutoLowRankMultivariateNorma` autoguide for a model with a high number of parameters. There are no issues during the training procedure; however, when I use the `quantiles` function, the code crashes due to the failed memory allocation (321.46GiB of GPU memory).

That is probably because it uses a dense covariance matrix.

```python
.../miniconda/envs/cellpymc/lib/python3.7/site-packages/numpyro/infer/autoguide.py in quantiles(self, params, quantiles)
    603 
    604     def quantiles(self, params, quantiles):
--> 605         transform = self.get_transform(params)
    606         quantiles = jnp.array(quantiles)[..., None]
    607         latent = dist.Normal(transform.loc, jnp.diagonal(transform.scale_tril)).icdf(quantiles)

.../miniconda/envs/cellpymc/lib/python3.7/site-packages/numpyro/infer/autoguide.py in get_transform(self, params)
    585     def get_transform(self, params):
    586         posterior = self.get_posterior(params)
--> 587         return LowerCholeskyAffine(posterior.loc, posterior.scale_tril)
    588 
    589     def get_posterior(self, params):

.../miniconda/envs/cellpymc/lib/python3.7/site-packages/numpyro/distributions/util.py in __get__(self, instance, obj_type)
    522         if instance is None:
    523             return self
--> 524         value = self.wrapped(instance)
    525         setattr(instance, self.wrapped.__name__, value)
    526         return value

...miniconda/envs/cellpymc/lib/python3.7/site-packages/numpyro/distributions/continuous.py in scale_tril(self)
    846         cov_diag_sqrt_unsqueeze = jnp.expand_dims(jnp.sqrt(self.cov_diag), axis=-1)
    847         Dinvsqrt_W = self.cov_factor / cov_diag_sqrt_unsqueeze
--> 848         K = jnp.matmul(Dinvsqrt_W, jnp.swapaxes(Dinvsqrt_W, -1, -2))
    849         K = jnp.add(K, jnp.identity(K.shape[-1]))
    850         scale_tril = cov_diag_sqrt_unsqueeze * jnp.linalg.cholesky(K)
```

Is there a way to avoid it? 
@fritzo 
CC @fehiepsi 


Discussion:
@fehiepsi I think we can simply port the low-memory version from Pyro, something like this:
```py
class AutoContinuous(...):
    def quantiles(self, params, quantiles):
        quantiles = jnp.array(quantiles)[..., None]
        loc, scale = self._loc_scale(params)
        latent = dist.Normal(loc, scale).icdf(quantiles)
        return self._unpack_and_constrain(latent, params)

class AutoLowRankMultivariateNormal(AutoContinuous):
    def _loc_scale(self, params):
        cov_factor = params[f"{self.prefix}_cov_factor"]
        scale = params[f"{self.prefix}_scale"] * ((cov_factor ** 2).sum(-1) + 1).sqrt()
        return self.loc, scale
```
I even think we should delete the method `AutoLowRankMultivariateNormal.get_transform()` method because it effectively cannot be called on large models.
It might be better to fix `quantiles` method directly. The `transform` is still useful in NeuTra transport.
```python
def quantiles(self, params, quantiles):
    posterior = self.get_posterior(params)
    quantiles = jnp.array(quantiles)[..., None]
    latent = dist.Normal(posterior.loc, jnp.sqrt(posterior.variance)).icdf(quantiles)
    return self._unpack_and_constrain(latent, params)
```

Resolving patch:
diff --git a/numpyro/infer/autoguide.py b/numpyro/infer/autoguide.py
--- a/numpyro/infer/autoguide.py
+++ b/numpyro/infer/autoguide.py
@@ -691,9 +691,12 @@ def median(self, params):
         return self._unpack_and_constrain(loc, params)
 
     def quantiles(self, params, quantiles):
-        transform = self.get_transform(params)
+        loc = params[f'{self.prefix}_loc']
+        cov_factor = params[f'{self.prefix}_cov_factor']
+        scale = params[f'{self.prefix}_scale']
+        scale = scale * jnp.sqrt(jnp.square(cov_factor).sum(-1) + 1)
         quantiles = jnp.array(quantiles)[..., None]
-        latent = dist.Normal(transform.loc, jnp.diagonal(transform.scale_tril)).icdf(quantiles)
+        latent = dist.Normal(loc, scale).icdf(quantiles)
         return self._unpack_and_constrain(latent, params)
 
 


Test patch: