Likelihoods combining multiple latent GPs in GPflow - gpflow

Question

I am trying to implement Variational Heteroscedastic Gaussian Process Regression in GPFlow.
My idea is to use the Variational Sparse Gaussian Process model (gpflow.models.SVGP) with a custom built Likelihood which expresses the density of y given two independend GPs f, g:
p(y|f, g) = N(y | f, t(g) )
Where t(·) is some transformation to make g positive (currently using tf.nn.softplus).
To make this work, I am setting model.num_latent to 2, but will implement the Likelihood in such a way that logp, conditional_mean, conditional_variance method only output tensors with shape (N, 1). Below is my current implementation:
from gpflow.likelihoods import Likelihood
from gpflow.decors import params_as_tensors
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
class HeteroscedasticGaussian(Likelihood):
r"""
When using this class, num_latent must be 2.
It does not support multi-output (num_output will be 1)
"""
def __init__(self, transform=tf.nn.softplus, **kwargs):
super().__init__(**kwargs)
self.transform = transform
#params_as_tensors
def Y_given_F(self, F):
mu = tf.squeeze(F[:, 0])
sigma = self.transform(tf.squeeze(F[:, 1]))
Y_given_F = tfd.Normal(mu, sigma)
return Y_given_F
#params_as_tensors
def logp(self, F, Y):
return self.Y_given_F(F).log_prob(Y)
#params_as_tensors
def conditional_mean(self, F):
return self.Y_given_F(F).mean()
#params_as_tensors
def conditional_variance(self, F):
return self.Y_given_F(F).variance()
My doubt is how to make the method variational_expectations work with what would be a double integral over df dg. I intend to use Gauss-Hermite quadrature, but I could not understand how to perform this double integral with the ndiagquad.
Is it as simple as calling
ndiagquad(self.logp, self.num_gauss_hermite_points, Fmu, Fvar, Y=Y)
??
EDIT:
Some MWE, using the implementation of variational_expectations from the base class Likelihood.
import gpflow as gpf
import tensorflow as tf
import numpy as np
N = 1001
M = 100
X = np.linspace(0, 4*np.pi, N)[:, None]
F = np.sin(X)
G = np.cos(X)
E = np.logaddexp(0, G) * np.random.normal(size=(N,1))
Y = F + E
Z_idx = np.random.choice(N, M, replace=False)
kernel = gpf.kernels.SquaredExponential(input_dim=1)
likelihood = HeteroscedasticGaussian()
likelihood.num_gauss_hermite_points = 11
model = gpf.models.SVGP(
X=X, Z=X[Z_idx], Y=Y,
kern=kernel,
likelihood=likelihood,
num_latent=2
)
# This method will call
# model.likelihood.variational_expectations(...)
# internally
model.compute_log_likelihood()
I am getting the following error message:
InvalidArgumentError: Incompatible shapes: [1001,11] vs. [2002]
[[{{node SVGP-bdd79b25-24/Normal/log_prob/standardize/sub}}]]
Which I believe has something to do with f, g being stacked on top of each other (shape [2002] = 2*N, with N = 1001), and Gauss-Hermite Points (11) being generated for only one Dimension for each observation (N = 1001), otherwise we would have a shape [1001, 11, 11] or [1001, 121=11^2].
All help appreciated.

Answer 1

You were very close - we did implement multi-dimensional quadrature in ndiagquad exactly for this use-case, though it needs to be called slightly differently. It would be good to have it work such that what you wrote works out of the box. But unfortunately it's not completely trivial to find a design that works both for actual multi-output regression, single-output likelihoods with several latent GPs, and the combination of the two! ndiagquad expects tuples (or lists) for Fmu and Fvar to indicate that you want multi-dimensional integration -- this preserves backwards-compatibility for f having shape (N, L) when you want to predict multiple outputs with Y also having shape (N, L).
So you've got to write the code slightly differently. This version works with your MWE:
from gpflow.likelihoods import Likelihood
from gpflow.decors import params_as_tensors
import tensorflow as tf
import tensorflow_probability as tfp
tfd = tfp.distributions
class MultiLatentLikelihood(Likelihood):
def __init__(self, num_latent=1, **kwargs):
super().__init__(**kwargs)
self.num_latent = num_latent
def _transform(self, F):
return [F[:, i] for i in range(self.num_latent)]
def predict_mean_and_var(self, Fmu, Fvar):
return super().predict_mean_and_var(self._transform(Fmu), self._transform(Fvar))
def predict_density(self, Fmu, Fvar, Y):
return super().predict_density(self._transform(Fmu), self._transform(Fvar), Y)
def variational_expectations(self, Fmu, Fvar, Y):
return super().variational_expectations(self._transform(Fmu), self._transform(Fvar), Y)
class HeteroscedasticGaussian(MultiLatentLikelihood):
r"""
When using this class, num_latent must be 2.
It does not support multi-output (num_output will be 1)
"""
def __init__(self, transform=tf.nn.softplus, **kwargs):
super().__init__(num_latent=2, **kwargs)
self.transform = transform
#params_as_tensors
def Y_given_F(self, F, G):
mu = tf.squeeze(F)
sigma = self.transform(tf.squeeze(G))
Y_given_F = tfd.Normal(mu, sigma)
return Y_given_F
#params_as_tensors
def logp(self, F, G, Y):
return self.Y_given_F(F, G).log_prob(Y)
#params_as_tensors
def conditional_mean(self, F, G):
return self.Y_given_F(F, G).mean()
#params_as_tensors
def conditional_variance(self, F, G):
return self.Y_given_F(F, G).variance()
I separated the boiler-plate code in its own class, MultiLatentLikelihood, to make it clearer what is generic and what is specific to a heteroscedastic Gaussian.
Maybe we should put both the MultiLatentLikelihood and examples into GPflow - if you would be up for it, why not add this to GPflow and make a Pull Request on github.com/GPflow/GPflow ? Would be happy to review it.
Also, the GPflow tutorials contain a notebook demonstrating how to deal with heteroscedastic noise in case you hadn't come across it - but it doesn't allow you to learn a latent GP to model the noise variance. So again, I think it'd be great if you wanted to extend the notebook with this example ("Demo 3") and make a Pull Request for your amendments :)