compute a KL divergence for a Gaussian Mixture prior and a normal distribution posterior

[neuronphysics]

🚀 Feature Request

Hi,

I am trying to compute a KL divergence for my variational autoencoder architecture between a Gaussian Mixture prior and a normal distribution posterior. Since this KL divergence is analytically intractable and only it could be evaluated by some sort of approximation, I am wondering whether it is also possible to compute it via Monte Carlo Sampling using botorch? I will appreciate if you could suggest an implementation with your library? Here is the architecture of my encoder and prior networks

import torch
import torch.nn as nn
from torch.distributions.normal import Normal
from torch.distributions import MixtureSameFamily, Categorical
from torch.distributions.independent import Independent
class VGMM(nn.Module): 
     def __init__(self,
                  u_dim,
                  h_dim,
                  z_dim,
                  n_mixtures,
                  batch_norm=False,
                  ):
        super(VGMM, self).__init__()
        self.n_mixtures =n_mixtures
        self.u_dim= u_dim
        self.h_dim=h_dim
        self.z_dim=z_dim
        self.batch_norm=  batch_norm
        encoder_layers=[nn.Linear(self.u_dim , self.h_dim)]
        if self.batch_norm:
            encoder_layers.append(torch.nn.BatchNorm1d(self.h_dim))
        encoder_layers=encoder_layers+[
            nn.ReLU(),
            nn.Linear(self.h_dim, self.h_dim),
        ]
        if self.batch_norm:
            encoder_layers= encoder_layers+[nn.BatchNorm1d(self.h_dim)]

        encoder_layers  = encoder_layers+[nn.ReLU()]

        self.enc        = torch.nn.Sequential(*encoder_layers)

        self.enc_mean   = nn.Linear(self.h_dim, self.z_dim)

        self.enc_logvar = nn.Linear(self.h_dim, self.z_dim)
        self.dist = MixtureSameFamily
        self.comp = Normal
        self.mix = Categorical

        layers_prior = [nn.Linear(self.u_dim, self.h_dim)]
        if self.batch_norm:
            layers_prior.append(torch.nn.BatchNorm1d(self.h_dim))
        layers_prior = layers_prior + [
            nn.ReLU(),
        ]

        self.prior = torch.nn.Sequential(*layers_prior)

        self.prior_mean = nn.ModuleList(
            [nn.Linear(self.h_dim, self.z_dim) for _ in range(n_mixtures)]
        )

        self.prior_logvar = nn.ModuleList(
            [nn.Linear(self.h_dim, self.z_dim) for _ in range(n_mixtures)]
        )
        self.prior_weights = nn.Linear(self.h_dim, n_mixtures) 
     def forward(self, u):
        encoder_input = self.enc(u)
        enc_mean   = self.enc_mean(encoder_input)
        enc_logvar = self.enc_logvar(encoder_input)
        enc_logvar = nn.Softplus()(enc_logvar)
        prior_input =self.prior(u)
        prior_mean  = torch.cat([ self.prior_mean[n](prior_input).unsqueeze(1) for n in range(self.n_mixtures)],dim=1,)
        prior_logvar = torch.cat([self.prior_logvar[n](prior_input).unsqueeze(1)for n in range(self.n_mixtures)],dim=1,)
        prior_w     = self.prior_weights(prior_input)
        prior_sigma = prior_logvar.exp().sqrt()
        prior_dist = self.dist(self.mix(logits=prior_w), Independent(self.comp(prior_mean, prior_sigma), 1))
        post_dist = self.comp(enc_mean, enc_logvar.exp().sqrt())
        z_t      = self.reparametrization(enc_mean, enc_logvar)
        return prior_dist, post_dist, z_t
     def reparametrization(self, mu, log_var):
        var = torch.exp(log_var* 0.5)
        eps = torch.FloatTensor(var.size()).normal_(mean=0, std=1)
        eps = torch.autograd.Variable(eps)
        return eps.mul(var).add_(mu).add_(1e-7)

Is there any possibility to use botorch library and expliciltly Monte Carlo sampling in order to compute this KL divergence? I was wondering whether botorch.acquisition.active_learning.StatisticalDistance for KL distance_metric could be used? Thanks in advance.

[esantorella]

Hi, thanks for reaching out!

If I understand correctly, you’re asking about computing a KL divergence between two continuous distributions via Monte Carlo sampling. I would expect that it’s possible to use MCMC to do this, although I’m not sure how well it would converge.

However, BoTorch doesn’t have any functionality for estimating a KL divergence in that way, and since BoTorch is intended as a Bayesian Optimization library, we’re not planning on adding that functionality.

You can implement MCMC methods on top of BoTorch Posterior objects, which represent distributions. However, you can also do this with the underlying torch distributions that you’re already using in the VAE, so I don’t think BoTorch would add a lot of value there.

[neuronphysics]

Thank you for your reply. Is it possible to provide an example about how to combine the MCMC method with a Posterior object?

[esantorella]

So to be clear, I don't think BoTorch is going to be helpful here. It doesn't provide sufficiently generic MCMC functionality, and BoTorch's posteriors are very similar to PyTorch Distributions. So I'd recommend continuing to use PyTorch Distributions directly and, if you use MCMC, use another library's MCMC routines, like Pyro's.

Now, without my BoTorch developer hat on, I think this can solved with nothing but PyTorch and the distribution's sample method:

import torch
import torch.distributions as D

def approximate_kl_with_sampling(
    p: D.Distribution, q: D.Distribution, n_samples: int
) -> float:
    """
    Args:
        p: distribution
        q: reference distribution
        n_samples: Number of draws

    Approximate the KL divergence from q to p by sampling from p.

    The KL divergence (relative entropy) is given by integrating
    log(p(x) / q(x)) over the density of p, so it can be computed by sampling
    form p and evaluating log(p(x) / q(x)) at the samples.
    """
    samples = p.sample((n_samples,))
    p_log_probs = p.log_prob(samples)
    q_log_probs = q.log_prob(samples)
    return p_log_probs.mean() - q_log_probs.mean()

In this case p and q can be the torch.distributions.MixtureSameFamily and torch.distributions.MixtureSameNormal instances you already have.

I checked this by looking at the degenerate case where there is weight on only one component of the mixture distribution, in which case we get back the analytical solution for the KL divergence between two Gaussians. However, this seems to not converge very quickly in cases where q and p differ mainly in the tails, so I'd recommend a large number of samples.