Heteroskedastic Gaussian Likelihood

[gustavocmv]

PR type: New Feature

Related issue(s)/PRs: #1559

Summary

Proposed changes
Analytical implementation of Heteroskedastic Gaussian Likelihood.

Tests are not implemented yet: I'm open to ideas about test cases and best file to place them. I thought about adding it to tests/gpflow/likelihoods/test_heteroskedastic_constant_variance.py, but couldn't figure out how to on my first try.

What alternatives have you considered?
This provides analytical expressions for the defaults of MultiLatentTFPConditional class as a separate HeteroskedasticGaussian class.
The other option would be to add multiple if 's to use analytical expressions when they are available, but I think it would cause more confusion than help.

Minimal working example

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
import gpflow as gpf

N = 1001

np.random.seed(0)
tf.random.set_seed(0)

X = np.linspace(0, 4 * np.pi, N)[:, None]
f1 = np.sin
f2 = np.cos
transform = np.exp
loc = f1(X)
scale = transform(f2(X))
Y = np.random.normal(loc, scale)


### This is the analytical implementation
likelihood = gpf.likelihoods.HeteroskedasticGaussian()

## This is the non-analytical implementation, uncomment to check equivalence
# likelihood = gpf.likelihoods.HeteroskedasticTFPConditional(
#     distribution_class=tfp.distributions.Normal,
#     scale_transform=tfp.bijectors.Exp(),
# )


kernel = gpf.kernels.SeparateIndependent(
    [
        gpf.kernels.SquaredExponential(),
        gpf.kernels.SquaredExponential(),
    ]
)

M = 20
Z = np.linspace(X.min(), X.max(), M)[:, None]
inducing_variable = gpf.inducing_variables.SeparateIndependentInducingVariables(
    [
        gpf.inducing_variables.InducingPoints(Z),
        gpf.inducing_variables.InducingPoints(Z),
    ]
)
model = gpf.models.SVGP(
    kernel=kernel,
    likelihood=likelihood,
    inducing_variable=inducing_variable,
    num_latent_gps=likelihood.latent_dim,
)

data = (X, Y)
loss_fn = model.training_loss_closure(data)

gpf.utilities.set_trainable(model.q_mu, False)
gpf.utilities.set_trainable(model.q_sqrt, False)

variational_vars = [(model.q_mu, model.q_sqrt)]
natgrad_opt = gpf.optimizers.NaturalGradient(gamma=0.1)

adam_vars = model.trainable_variables
adam_opt = tf.optimizers.Adam(0.01)

@tf.function
def optimisation_step():
    natgrad_opt.minimize(loss_fn, variational_vars)
    adam_opt.minimize(loss_fn, adam_vars)

epochs = 100
log_freq = 20

for epoch in range(1, epochs + 1):
    optimisation_step()
    if epoch % log_freq == 0 and epoch > 0:
        print(f"Epoch {epoch} - Loss: {loss_fn().numpy() : .4f}")

PR checklist

• New features: code is well-documented
  • detailed docstrings (API documentation)
  • notebook examples (usage demonstration)
• The bug case / new feature is covered by unit tests
• Code has type annotations
• I ran the black+isort formatter (make format)
• I locally tested that the tests pass (make check-all)

Release notes

Fully backwards compatible: yes