Any model to deal with varying output noise ?[Feature Request]

[lk1983823]

Is there any model for varying output noise or heteroskedastic noise like example in GPflow?
"https://gpflow.readthedocs.io/en/develop/notebooks/advanced/varying_noise.html"

[Balandat]

There is FixedNoiseGaussianLikelihood for the first case: https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/likelihoods/gaussian_likelihood.py#L100

There currently is no pre-made likelihood for group-level estimation of noise as in the example, but that shouldn't be too hard to set up.

In BoTorch we have a model that allows you to estimate a full heteroskedastic out-of-sample noise model: https://github.com/pytorch/botorch/blob/master/botorch/models/gp_regression.py#L280

[lk1983823]

@Balandat Thanks. I don't quite understand the meaning of "estimate a full heteroskedastic out-of-sample noise model'. Does it mean that I don't have to give any test noise to a trained model during the testing period, just like the GPflow case 2 ? However, in the FixedNoiseGaussianLikelihood example, I find the test_noise has to be given. Maybe you can add a more detailed example in GPytorch as to how to use such kind of model.

Example for the model in Botorch

>>> train_X = torch.rand(20, 2)
>>> train_Y = torch.sin(train_X).sum(dim=1, keepdim=True)
>>> se = torch.norm(train_X, dim=1, keepdim=True)
>>> train_Yvar = 0.1 + se * torch.rand_like(train_Y)
>>> model = HeteroskedasticSingleTaskGP(train_X, train_Y, train_Yvar)

[Balandat]

Does it mean that I don't have to give any test noise to a trained model during the testing period, just like the GPflow case 2 ?

Yes, it means that at test time you don't need to provide test noise.

FixedNoiseGaussianLikelihood example, I find the test_noise has to be given.

The noise only needs to be given at training time. With a FixedNoiseGaussianLikelihood you can infer the posterior noise level of the latent function you are estimating without providing that noise. But technically if you want the observation noise (the posterior predictive), that only works at the training point locations, as the model does not know what observation noise to use at unobserved locations.

The HeteroskedasticSingleTaskGP estimates an out-of-sample noise model over the whole space by fitting another GP model under the hood.

[julioasotodv]

Hi,

3 years later: but I would like to leave this GPFlow example as a reference: https://gpflow.github.io/GPflow/2.9.0/notebooks/advanced/heteroskedastic.html

it basically uses two latent GPs to model mean and variance, and it introduces an heteroskedastic likelihood.

It has the advantage of not requiring to provide any noise estimates during training, as the heterogeneous noise is also inferred during training (with SVI).

I am unsure on whether this approach is similar to BoTorch's MostLikelyHeteroskedasticGP, but it does not look quite like it.

I am trying to find an easy way to reproduce that GPFlow notebook in GPyTorch, but no luck so far reproducing the notebook's likelihood.

[Balandat]

Haven't dug into the details here, but it does seem that wile it's similar in spirit to the MostLikelyHeteroskedasticGP it does not require to infer local variance estimates from the data explicitly for the log-variance model (but instead infer this jointly and implicitly via SVI).

cc @Ryan-Rhys

[julioasotodv]

Thank you @Balandat for pointing it out.

Indeed: it uses two latent, separate GPs (as defined by two independent kernels for their covariances) that get trained jointly using SVI.

I believe that currently gpytorch lacks some classes in order to reproduce a similar model. For instance:

• gpf.likelihoods.HeteroskedasticTFPConditional as a likelihood function where noise is not necessarily homoscedastic
• gpf.kernels.SeparateIndependent as a list of independent kernels (I believe this is similar to gpytorch.models.IndependentModelList but being only the kernels instead of the full models)
• A gpflow.models.svgp.SVGP that wraps one or more latent GPs into a single instance for jointly training all those included latent GPs

That's all I could decypher, as I am not used to GPFlow (nor TF for that matter).

[julioasotodv]

Ok, so I tried to use gpytorch's internal tools to create something similar to the GPFlow model mentioned previously. I was able to make it work using gpytorch's Pyro integration.

Here's a reproducible demo:

import numpy as np
import matplotlib.pyplot as plt

train_x = np.random.uniform(8, 20, 500)

train_y = (0.1 * train_x ** 2 - 3.5*train_x + np.cos(train_x*1.3) + 40 
           + np.random.normal(0, 0.01, len(train_x)) * train_x**2 * 0.3)

# Create some "gap" in the data
# to test epistemic uncertainty
# estimation:
gap_mask = (train_x<12) | (train_x>16)
train_x = train_x[gap_mask]
train_y = train_y[gap_mask]

# Scaling:
train_x = (train_x - train_x.min()) / (train_x.max() - train_x.min())
train_y = (train_y - train_y.mean()) / train_y.std()

# Plot training data:
plt.plot(train_x, train_y, 'k*')

import torch
import pyro
import gpytorch

class HeteroskedasticGP(gpytorch.models.ApproximateGP):
    
    def __init__(self, num_inducing=64, name_prefix="heteroskedastic_gp"):
        
        self.name_prefix = name_prefix

        # Define all the variational stuff
        inducing_points_task_one = torch.linspace(-1, 1, num_inducing)
        inducing_points_task_two = torch.linspace(-1, 1, num_inducing)
        
        inducing_points = torch.stack([inducing_points_task_one.unsqueeze(1), 
                                       inducing_points_task_two.unsqueeze(1)])
        
        # We have to mark the CholeskyVariationalDistribution as batch
        # so that we learn a variational distribution for each task (we have 2):
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            num_inducing_points=inducing_points.size(-2), 
            batch_shape=torch.Size([2]))
        
        single_variational_strategy = gpytorch.variational.VariationalStrategy(
                self, inducing_points, variational_distribution, learn_inducing_locations=True)
        
        # Wrap the single variational strategy into a
        # Linear Model of Coregionalization one, so the two
        # tasks (and therefore latent GPs) are assumed to be
        # somehow related (and we learn such relationship):
        variational_strategy = gpytorch.variational.LMCVariationalStrategy(single_variational_strategy,
                                                                           num_tasks=2,
                                                                           num_latents=2)
        # Standard initializtation
        super().__init__(variational_strategy)
        
        # The mean and covariance modules should be marked as batch
        # so we learn a different set of hyperparameters:
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([2]))
        
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(batch_shape=torch.Size([2])),
            batch_shape=torch.Size([2])
        )

    def forward(self, x):
        # The forward function should be written as if we were dealing with each output
        # dimension in batch
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

    def guide(self, x, y):
        
        # Get q(f) - variational (guide) distribution of latent function
        function_dist = self.pyro_guide(x)

        # Use a plate here to mark conditional independencies.
        # Our samples are of shape (nobs, 2), therefore the dim=-2
        # in the plate:
        with pyro.plate(self.name_prefix + ".data_plate", dim=-2):
            
            # Sample from latent function distribution
            function_samples = pyro.sample(self.name_prefix + ".f(x)", function_dist)

    def model(self, x, y):
        
        pyro.module(self.name_prefix + ".gp", self)

        # Get p(f) - prior distribution of latent function
        function_dist = self.pyro_model(x)

        # Use a plate here to mark conditional independencies.
        # Our samples are of shape (nobs, 2), therefore the dim=-2
        # in the plate:
        with pyro.plate(self.name_prefix + ".data_plate", dim=-2):
            
            # Sample from latent function distribution
            function_samples = pyro.sample(self.name_prefix + ".f(x)", function_dist)
            
            mean_samples = function_samples[...,0]
            std_samples = function_samples[...,1]
            
            # Exp to force always nonnegative stddevs:
            transformed_std_samples = torch.exp(std_samples)

            # Sample from observed distribution
            return pyro.sample(self.name_prefix + ".y",
                               pyro.distributions.Normal(mean_samples, transformed_std_samples),
                               obs=y)

# Instantiate model:
pyro.clear_param_store()  # Good practice
model = HeteroskedasticGP()

num_iter = 1000
num_particles = 256

# training routine:
def train():
    optimizer = pyro.optim.Adam({"lr": 0.01})
    elbo = pyro.infer.Trace_ELBO(num_particles=num_particles, vectorize_particles=True, retain_graph=True)
    svi = pyro.infer.SVI(model.model, model.guide, optimizer, elbo)

    model.train()
    iterator = range(num_iter)
    for i in iterator:
        model.zero_grad()
        loss = svi.step(torch.from_numpy(train_x).float().unsqueeze(1), torch.from_numpy(train_y).float())
        print(loss)

# Train the GP:
train()

# Function to plot predictions:
def plot_preds(ax, feature, mean_samples, var_samples, bootstrap=True, n_boots=100,
               show_epistemic=False, epistemic_mean=None, epistemic_var=None):
    """Plots the overall mean and variance of the aggregate system.
    Inherited from https://www.pymc.io/projects/examples/en/latest/gaussian_processes/GP-Heteroskedastic.html.

    We can represent the overall uncertainty via explicitly sampling the underlying normal
    distributrions (with `bootstrap=True`) or as the mean +/- the standard deviation from
    the Law of Total Variance. 
    
    For systems with many observations, there will likely be
    little difference, but in cases with few observations and informative priors, plotting
    the percentiles will likely give a more accurate representation.
    """
    if bootstrap:
        means = np.expand_dims(mean_samples.T, axis=2)
        stds = np.sqrt(np.expand_dims(var_samples.T, axis=2))
        
        samples_shape = mean_samples.T.shape + (n_boots,)
        samples = np.random.normal(means, stds, samples_shape)
        
        reshaped_samples = samples.reshape(mean_samples.shape[1], -1).T
        
        l, m, u = [np.percentile(reshaped_samples, p, axis=0) for p in [2.5, 50, 97.5]]
        ax.plot(feature, m, label="Median", color="b")
    
    else:
        m = mean_samples.mean(axis=0)
        sd = np.sqrt(mean_samples.var(axis=0) + var_samples.mean(axis=0))
        l, u = m - 1.96 * sd, m + 1.96 * sd
        ax.plot(feature, m, label="Mean", color="b")
    
    if show_epistemic:
        ax.fill_between(feature,
                        l,
                        epistemic_mean-1.96*np.sqrt(epistemic_var), 
                        alpha=0.2,
                        color="#88b4d2",
                        label="Total Uncertainty (95%)")
        ax.fill_between(feature, 
                        u, 
                        epistemic_mean+1.96*np.sqrt(epistemic_var), 
                        alpha=0.2,
                        color="#88b4d2")
        ax.fill_between(feature, 
                        epistemic_mean-1.96*np.sqrt(epistemic_var), 
                        epistemic_mean+1.96*np.sqrt(epistemic_var),
                        alpha=0.4,
                        color="#88b4d2",
                        label="Epistemic Uncertainty (95%)")
    else:
        ax.fill_between(feature, 
                        l, 
                        u, 
                        alpha=0.2,
                        color="#88b4d2",
                        label="Total Uncertainty (95%)")

# Test data to predict:
x_padding = 0.1

test_x = torch.linspace(train_x.min() - (train_x.max() - train_x.min()) * x_padding, 
                        train_x.max() + (train_x.max() - train_x.min()) * x_padding, 
                        100).float()

# Predict:
model.eval()
with torch.no_grad():
    output_dist = model(test_x)

# Extract predictions:
output_samples = output_dist.sample(torch.Size([1000]))
mu_samples = output_samples[...,0]
sigma_samples = torch.exp(output_samples[...,1])

# Plot predictions:
plt.plot(train_x, train_y, 'k*', label="Observed Data")
ax = plt.gca()
plot_preds(ax, 
           test_x.numpy(),
           mu_samples.numpy(), 
           sigma_samples.numpy()**2,
           bootstrap=False, 
           n_boots=100,
           show_epistemic=True,
           epistemic_mean=output_dist.mean[:,0].detach().numpy(),
           epistemic_var=output_dist.stddev[:,0].detach().numpy()**2)
ax.legend();

Which yields the following model:

Not bad! My biggest surprise is how coherent gpytorch's API is, as everything was in place to create this type of model (even though AFAIK it was never tested). Both epistemic and aleatoric uncertainties are jointly learned correctly.

I still have to figure out a way for the model to be stable when noise is low (as we don't have here explicit marginal likelihood to which we can "request" a minimum noise level). Any suggestions are more than welcome!

Thanks all!

[Balandat]

Neat! Thanks for sharing!

[julioasotodv]

@Balandat you're welcome! Another way to deal with heteroskedastic noise could be using the predictive log likelihood objective function, as shown in https://docs.gpytorch.ai/en/stable/examples/04_Variational_and_Approximate_GPs/Approximate_GP_Objective_Functions.html

I guess that this issue can be closed now.

[Balandat]

converted this into a discussion for future reference