<jupyter_start><jupyter_text>AutoLowRankMultivariateNormal parametrizationsThis notebook aims to validate parametrizations of Pyro's [AutoLowRankMultivariateNormal](http://docs.pyro.ai/en/stable/infer.autoguide.htmlautolowrankmultivariatenormal) guide, in the context of https://github.com/pyro-ppl/pyro/pull/2127<jupyter_code># These are needed so the git commands below take immediate effect.
%load_ext autoreload
%autoreload 2
import torch
import pyro
import pyro.distributions as dist
from pyro.infer import SVI, Trace_ELBO
from pyro.infer.autoguide import AutoLowRankMultivariateNormal
from pyro.optim import ClippedAdam
from matplotlib import pyplot

%matplotlib inline
%config InlineBackend.figure_format = 'svg'
%config InlineBackend.rc = {'figure.facecolor': (1, 1, 1, 1)}<jupyter_output><empty_output><jupyter_text>Let's test with a generative model with a wide range of data scales (exponentially distributed), and with a low rank structure. We'll later fit this with a guide with even lower rank.<jupyter_code>class Model:
    def __init__(self, dim, rank):
        self.dim = dim
        self.rank = rank
        self.loc1 = dist.Laplace(0, 1).sample((dim,))
        self.scale1 = dist.Exponential(1).sample((dim,))
        self.loc2 = dist.Laplace(0, 1).sample((rank,))
        self.scale2 = dist.Exponential(1).sample((rank,))
        self.mat = dist.Normal(0, 1).sample((dim, rank))

    def __call__(self):
        z = pyro.sample("z",
                        dist.Normal(self.loc1, self.scale1)
                            .expand([self.dim]).to_event(1))
        pyro.sample("x",
                    dist.Normal(self.loc2, self.scale2)
                        .expand([self.rank]).to_event(1),
                    obs=z @ self.mat)
def train(dim, rank):
    pyro.clear_param_store()
    pyro.set_rng_seed(123456789)
    model = Model(dim, 2 * rank)
    guide = AutoLowRankMultivariateNormal(model, rank=rank, init_scale=0.01)
    optim = ClippedAdam({"lr": 0.01})
    elbo = Trace_ELBO()
    svi = SVI(model, guide, optim, elbo)
    losses = []
    for step in range(1000):
        loss = svi.step() / dim
        losses.append(loss)
        if step % 100 == 0:
            print("step {: >4} loss = {:0.8g}".format(step, loss))
    return losses
results = {}
!(cd ~/pyro ; git checkout dev)
results["dev-10-3"] = train(10, 3)
results["dev-100-10"] = train(100, 10)
results["dev-1000-30"] = train(1000, 30)
results["dev-10000-100"] = train(10000, 100)
!(cd ~/pyro ; git checkout auto-lowrank-mvn-reparam)
results["new-10-3"] = train(10, 3)
results["new-100-10"] = train(100, 10)
results["new-1000-30"] = train(1000, 30)
results["new-10000-100"] = train(10000, 100)
def plot(dim, rank):
    pyplot.figure(figsize=(9,3))
    for name, losses in sorted(results.items()):
        if name.endswith(f"-{dim}-{rank}"):
            pyplot.plot(losses, label=name)
    pyplot.title(f"dim={dim} rank={rank}")
    pyplot.ylabel('loss')
    pyplot.yscale('log')
    pyplot.xlabel('svi step')
    pyplot.legend()
    pyplot.tight_layout()
plot(10, 3)
plot(100, 10)
plot(1000, 30)
plot(10000, 100)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14399744572158366, 'fraction_numerical': 0.05300127713920817, 'mean_word_length': 3.4314002828854315, 'pattern_counts': {'":': 1, '<': 9, '<?xml version=': 0, '>': 10, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '106315', 'license': 'Apache-2.0', 'n_tokens_mistral': 1206, 'n_tokens_neox': 1060, 'n_words': 242, 'path': '2019-11-lowrank/experiment.ipynb', 'repo_name': 'vitkl/sandbox', 'size': '3132'}