How to add VAE latent variance variables as heteroscedastic noise into the model?

[marvinquiet]

Hello, I am trying to use GPytorch to implement SVGP-VAE in order to understand more about how to use GPytorch and I am referring to their code. I am still learning this and in my understanding, certain tweaks on the ApproximateGP could achieve the purpose of adding GP prior over the latent variable with a kernel function.

Therefore, I first constructed the ApproximateGP with the model from the Stochastic Variational GP Regression tutorial.

from gpytorch.models import ApproximateGP
from gpytorch.variational import CholeskyVariationalDistribution
from gpytorch.variational import VariationalStrategy

class GPModel(ApproximateGP):
    def __init__(self, inducing_points):
        variational_distribution = CholeskyVariationalDistribution(inducing_points.size(0))
        variational_strategy = VariationalStrategy(self, inducing_points, variational_distribution, learn_inducing_locations=True)
        super(GPModel, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

Our data has $X$ (auxiliary data showing an intrinsic correlation in the data) and $Y$ (input). We first use VAE model to get the latent variables $\mu$ and $\sigma$ denoted as qnet_mu and qnet_var below in the code. According to the paper, they separate each latent channel $l$ and take $\mu^l$ as the response variable of $Y^l$ and $\sigma^l$ as heteroscedastic noise. I, therefore, did some search and found that #653 discussed using FixedNoiseGaussianLikelihood as likelihood but it seems that the noise is fully observed. In our case, it will change over training.

I was wondering if you could provide some suggestions on this and thanks a lot in advance! I provided my example code below.

Wenjing

...

svgp = GPModel(inducing_points)
svgp.to(args.device)
likelihood = gpytorch.likelihoods.GaussianLikelihood()
likelihood.to(args.device)
svgp.train()
likelihood.train()
optimizer = optim.AdamW(list(encoder.parameters())+list(decoder.parameters())+list(dec_mean.parameters())+list(svgp.parameters()) + list(likelihood.parameters()), 
                                 lr=args.lr, weight_decay=args.weight_decay)

mll = gpytorch.mlls.VariationalELBO(likelihood, svgp, num_data=N_train) # variational loss


for epoch in range(args.maxiter):
      for batch_idx, (x_batch, y_batch) in enumerate(dataloader):
          optimizer.zero_grad()
          x_batch = x_batch.to(device) # auxillary data
          y_batch = y_batch.to(device) # input data
          qnet_mu, qnet_var = encoder(y_batch) # q(z | y), mu (batch_size, z_dim) and variance (batch_size, z_dim) from VAE
          # for each l-th element in z dimension
          p_m, p_v = [], []
          GP_loss = 0
          for l in range(z_dim):
              output = svgp(x_batch)
              ## Can I add FixedNoiseGaussianLikelihood here? 
              ## likelihood = FixedNoiseGaussianLikelihood(noise=qnet_var[:, l], learn_additional_noise=True) 
              ## likelihood.train()
              ## mll = gpytorch.mlls.VariationalELBO(likelihood, svgp, num_data=batch_size) # variational loss
              ## What should I do with optimizer because it contains likelihood parameters..
              GP_loss += -mll(output, qnet_mu[:, l])
              p_m.append(output.mean)
              p_v.append(output.variance)
          p_m = torch.stack(p_m, dim=1)
          p_v = torch.stack(p_v, dim=1)
          # GP_loss = GP_loss / z_dim