SVI for multivariate distribution

Hello all, I want to predict the mean and the covariance matrix from a multivariate normal distribution. If I just want to get the point estimation of mean and the covariance matrix, it will be easy enough using MAP like in this code:

import matplotlib.pyplot as plt
import torch
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import time
import math
import pyro
import pyro.distributions as dist
from torch.distributions import constraints
from pyro.infer import SVI, Trace_ELBO
from pyro.optim import Adam

loc = np.array([0.,-6.])
scale = np.array([[1.,-0.7],[-0.7,1.]])
truth = np.random.multivariate_normal(loc, scale,5000)

def model(data):
    loc = pyro.param("loc",torch.zeros(2))
    scale_tril = pyro.param("scale",torch.eye(2,2))

    with pyro.iarange("count_arange", len(data)):
        return pyro.sample("count", dist.MultivariateNormal(loc, scale_tril), obs=data)


def guide(data):
    pass


def main(data, num_iter=2000):
    optim = Adam({"lr": 0.005})
    svi = SVI(model, guide, optim, loss=Trace_ELBO())
    pyro.clear_param_store()
    start = time.time()
    losses = np.zeros(num_iter)
    for i in range(num_iter):
        losses[i] = svi.step(data)
        if i % (num_iter // 10) == 0:
            print("[iteration %04d] loss: %.4f" % (i + 1, losses[i]))
            elapsed_time = time.time()
            print("elapsed time: %.2f" % (elapsed_time - start))
    end = time.time()
    print("Loop take time %.2f" % (end - start))
    plt.plot(losses)
    plt.show()


main(torch.tensor(truth).float(), 5000)

loc_prediction = pyro.param("loc").data.numpy()
scale_tril = pyro.param("scale").data.numpy()
scale_prediction = scale_tril @ scale_tril.T
prediction = np.random.multivariate_normal(loc_prediction, scale_prediction,5000)

plt.scatter(truth[:,0],truth[:,1],c="r",label ="truth",alpha = 0.1)
plt.scatter(prediction[:,0],prediction[:,1],c="b",label ="prediction",alpha = 0.1)
plt.legend(loc = "lower left")

However what I need is the distribution of mean and the covariance matrix. So I made this model:

loc = np.array([0.,-6.])
scale = np.array([[1.,-0.7],[-0.7,1.]])
truth = np.random.multivariate_normal(loc, scale,5000)

def model(data):
    loc_of_loc = torch.zeros(2)
    scale_of_loc = torch.ones(2)
    
    loc = pyro.sample("loc", dist.Normal(loc_of_loc,scale_of_loc))
    scale = pyro.param("scale",torch.eye(2,2))
    
    with pyro.iarange("count_arange", len(data)):
        pyro.sample("count", dist.MultivariateNormal(loc, scale), obs=data)


def guide(data):
    loc_of_loc_q = pyro.param("loc_of_loc_q", torch.zeros(2))
    scale_of_loc_q = pyro.param("scale_of_loc_q", torch.ones(2))
    pyro.sample("loc", dist.Normal(loc_of_loc_q,scale_of_loc_q))



def main(data, num_iter=2000):
    optim = Adam({"lr": 0.005})
    svi = SVI(model, guide, optim, loss=Trace_ELBO())
    pyro.clear_param_store()
    start = time.time()
    losses = np.zeros(num_iter)
    for i in range(num_iter):
        losses[i] = svi.step(data)
        if i % (num_iter // 10) == 0:
            print("[iteration %04d] loss: %.4f" % (i + 1, losses[i]))
            elapsed_time = time.time()
            print("elapsed time: %.2f" % (elapsed_time - start))
    end = time.time()
    print("Loop take time %.2f" % (end - start))
    plt.plot(losses)
    plt.show()


main(torch.tensor(truth).float(), 5000)

loc_prediction = pyro.param("loc_of_loc_q").data.numpy()
scale_tril = pyro.param("scale").data.numpy()
scale_prediction = scale_tril @ scale_tril.T
prediction = np.random.multivariate_normal(loc_prediction, scale_prediction,5000)

plt.scatter(truth[:,0],truth[:,1],c="r",label ="truth",alpha = 0.1)
plt.scatter(prediction[:,0],prediction[:,1],c="b",label ="prediction",alpha = 0.1)
plt.legend(loc = "lower left")

In the code above I am using MAP for evaluating covariance matrix but I want the distribution of the covariance matrix. In STAN there is LKJ prior for prior distribution of the covariance matrix, in pyro how to get the distribution of covariance matrix?

Hi @yusri_dh, we haven’t implemented the LKJ or inverse Wishart distributions for covariance matrices, but you could try the AutoDiagonalNormal or AutoMultivariateNormal guides from pyro.contrib.autoguide:

def model(data):
    ...
    loc = pyro.sample("loc", ...)
    scale = pyro.sample("scale", ...)
    ...

guide = AutoMultivariateNormal(model)

then after training you can call the guide to generate samples:

sample = guide(data)
print(sample['loc'])
print(sample['scale'])

Note that the learned distribution uses a weird transformed multivariate normal distribution rather than the more standard LKJ prior, and I’m not sure how well it will work. Let us know either way :slight_smile: