Repository: secondmind-labs/trieste

Issue:
SVGP update equations contain an error
When we change the inducing points before re-training SVGPs and Deep GPs, we need to update q_mu and q_sqrt such that the posterior distribution is unchanged at the new inducing points. 

The current equations only work when the mean function is zero. We need to add a correction term at those places;

(new_q_mu =new_q_mu - model mean prediction at Z)
https://github.com/secondmind-labs/trieste/blob/develop/trieste/models/gpflow/models.py#L1014

(f_mu = f_mu - model mean prediction at Z)
https://github.com/secondmind-labs/trieste/blob/a4b25f8cde453c2a19399e66c3c440d51cfbac84/trieste/models/gpflow/utils.py#L330

DeepGPs currently do not have this problem because we don't have a scheme to update inducing points.

It is unclear whereas we have a test for whiten_points, but we should definitely have one with a non-zero mean function.


Discussion:
The only test we have for this part of the code is:

https://github.com/secondmind-labs/trieste/blob/a4b25f8cde453c2a19399e66c3c440d51cfbac84/tests/unit/models/gpflow/test_models.py#L1431

The correct test here is to check that the posterior distribution at the new inducing points is unchanged after you update the model inducing points.

Resolving patch:
diff --git a/trieste/models/gpflow/models.py b/trieste/models/gpflow/models.py
--- a/trieste/models/gpflow/models.py
+++ b/trieste/models/gpflow/models.py
@@ -961,7 +961,6 @@ def update(self, dataset: Dataset) -> None:
         num_data = dataset.query_points.shape[0]
         assert self.model.num_data is not None
         self.model.num_data.assign(num_data)
-        self.update_posterior_cache()
 
         if self._inducing_point_selector is not None:
             new_inducing_points = self._inducing_point_selector.calculate_inducing_points(
@@ -974,6 +973,7 @@ def update(self, dataset: Dataset) -> None:
                 )
             ):  # only bother updating if points actually change
                 self._update_inducing_variables(new_inducing_points)
+                self.update_posterior_cache()
 
     def optimize(self, dataset: Dataset) -> None:
         """
@@ -1013,6 +1013,7 @@ def _update_inducing_variables(self, new_inducing_points: TensorType) -> None:
             new_q_mu, new_q_sqrt = _whiten_points(self, new_inducing_points)
         else:
             new_q_mu, new_f_cov = self.predict_joint(new_inducing_points)  # [N, L], [L, N, N]
+            new_q_mu -= self.model.mean_function(new_inducing_points)
             jitter_mat = DEFAULTS.JITTER * tf.eye(
                 tf.shape(new_inducing_points)[0], dtype=new_f_cov.dtype
             )
diff --git a/trieste/models/gpflow/utils.py b/trieste/models/gpflow/utils.py
--- a/trieste/models/gpflow/utils.py
+++ b/trieste/models/gpflow/utils.py
@@ -327,7 +327,9 @@ def _whiten_points(
     :para inducing_points: The new inducing point locations.
     :return: The updated q_mu and q_sqrt with shapes [N, L] and [L, N, N], respectively.
     """
+
     f_mu, f_cov = model.model.predict_f(inducing_points, full_cov=True)  # [N, L], [L, N, N]
+    f_mu -= model.model.mean_function(inducing_points)
     Knn = model.get_kernel()(inducing_points, full_cov=True)  # [N, N]
     jitter_mat = DEFAULTS.JITTER * tf.eye(tf.shape(inducing_points)[0], dtype=Knn.dtype)
     Lnn = tf.linalg.cholesky(Knn + jitter_mat)  # [N, N]


Test patch:
diff --git a/tests/unit/models/gpflow/test_models.py b/tests/unit/models/gpflow/test_models.py
--- a/tests/unit/models/gpflow/test_models.py
+++ b/tests/unit/models/gpflow/test_models.py
@@ -35,6 +35,7 @@
 import pytest
 import tensorflow as tf
 import tensorflow_probability as tfp
+from gpflow.config import Config, as_context
 from gpflow.inducing_variables import (
     SeparateIndependentInducingVariables,
     SharedIndependentInducingVariables,
@@ -48,6 +49,7 @@
     mock_data,
     sgpr_model,
     svgp_model,
+    svgp_model_with_mean,
     two_output_svgp_model,
     vgp_matern_model,
     vgp_model,
@@ -55,9 +57,10 @@
 from tests.util.models.models import fnc_2sin_x_over_3, fnc_3x_plus_10
 from trieste.data import Dataset, add_fidelity_column
 from trieste.logging import step_number, tensorboard_writer
-from trieste.models import TrainableProbabilisticModel
+from trieste.models import ProbabilisticModelType, TrainableProbabilisticModel
 from trieste.models.gpflow import (
     GaussianProcessRegression,
+    GPflowPredictor,
     MultifidelityAutoregressive,
     MultifidelityNonlinearAutoregressive,
     SparseGaussianProcessRegression,
@@ -83,6 +86,7 @@
 from trieste.models.optimizer import BatchOptimizer, DatasetTransformer, Optimizer
 from trieste.space import Box
 from trieste.types import TensorType
+from trieste.utils import DEFAULTS
 
 
 def _3x_plus_gaussian_noise(x: tf.Tensor) -> tf.Tensor:
@@ -1600,7 +1604,6 @@ def test_sparse_variational_pairwise_covariance_for_non_whitened(
     y2 = y1 * 0.5
 
     svgp = two_output_svgp_model(x, mo_type, whiten)
-
     model = SparseVariational(svgp, BatchOptimizer(tf.optimizers.Adam(), max_iter=3, batch_size=10))
     model.model.whiten = whiten
 
@@ -1618,6 +1621,57 @@ def test_sparse_variational_pairwise_covariance_for_non_whitened(
     np.testing.assert_allclose(expected_covariance, actual_covariance, atol=1e-4)
 
 
+class DummyInducingPointSelector(InducingPointSelector[GPflowPredictor]):
+    def __init__(self, new_inducing_points: TensorType, recalc_every_model_update: bool = True):
+        super().__init__(recalc_every_model_update)
+        self._new_inducing_points = new_inducing_points
+
+    def _recalculate_inducing_points(
+        self, M: int, model: ProbabilisticModelType, dataset: Dataset
+    ) -> TensorType:
+        return self._new_inducing_points
+
+
+@random_seed
+@pytest.mark.parametrize("whiten", [False, True])
+def test_sparse_variational_inducing_updates_preserves_posterior(
+    whiten: bool,
+) -> None:
+    default_jitter = 0.0
+    with as_context(Config(jitter=default_jitter)), unittest.mock.patch.object(
+        DEFAULTS, "JITTER", default_jitter
+    ):
+        x = tf.constant(np.linspace(0.0, 1.0, 8).reshape(-1, 1), dtype=gpflow.default_float())
+        y1 = fnc_3x_plus_10(x)
+
+        num_inducing_points = 4
+        xnew = tf.constant(
+            np.linspace(0.31, 0.77, num_inducing_points).reshape(-1, 1),
+            dtype=gpflow.default_float(),
+        )
+
+        svgp = svgp_model_with_mean(x, y1, whiten, num_inducing_points)
+        inducing_point_selector = DummyInducingPointSelector(xnew)
+        model = SparseVariational(
+            svgp,
+            BatchOptimizer(tf.optimizers.Adam(), max_iter=3, batch_size=10),
+            inducing_point_selector=inducing_point_selector,
+        )
+
+        np.testing.assert_array_equal(model.model.inducing_variable.Z, x[:num_inducing_points])
+
+        old_mu, old_sqrt = model.predict_joint(xnew)  # predict old posterior
+
+        model.update(Dataset(x, y1))  # this changes inducing points to xnew
+
+        np.testing.assert_array_equal(model.model.inducing_variable.Z, xnew)
+
+        new_mu, new_sqrt = model.predict_joint(xnew)  # predict new posterior
+
+        np.testing.assert_allclose(old_mu, new_mu, atol=1e-9)
+        np.testing.assert_allclose(old_sqrt, new_sqrt, atol=1e-9)
+
+
 def multifidelity_autoregressive_nd_dataset(n_dims: int = 1) -> Dataset:
     dataset = Dataset(
         tf.Variable(
diff --git a/tests/util/models/gpflow/models.py b/tests/util/models/gpflow/models.py
--- a/tests/util/models/gpflow/models.py
+++ b/tests/util/models/gpflow/models.py
@@ -17,6 +17,7 @@
 from collections.abc import Callable, Sequence
 
 import gpflow
+import numpy as np
 import tensorflow as tf
 import tensorflow_probability as tfp
 from gpflow.models import GPR, SGPR, SVGP, VGP, GPModel
@@ -380,6 +381,31 @@ def svgp_model(x: tf.Tensor, y: tf.Tensor, num_latent_gps: int = 1) -> SVGP:
     )
 
 
+def svgp_model_with_mean(
+    x: tf.Tensor, y: tf.Tensor, whiten: bool, num_inducing_points: int, num_latent_gps: int = 1
+) -> SVGP:
+    mean_function = gpflow.mean_functions.Linear(
+        A=0.37 * np.ones((1, 1), dtype=gpflow.default_float()),
+        b=0.19 * np.ones((1,), dtype=gpflow.default_float()),
+    )
+    q_mu = np.random.randn(num_inducing_points, 1)
+    q_sqrt = np.tril(np.random.randn(1, num_inducing_points, num_inducing_points))
+    m = SVGP(
+        gpflow.kernels.Matern32(variance=0.91),
+        gpflow.likelihoods.Gaussian(variance=0.23),
+        x[:num_inducing_points],
+        num_data=len(x),
+        num_latent_gps=num_latent_gps,
+        mean_function=mean_function,
+        whiten=whiten,
+        q_mu=q_mu,
+        q_sqrt=q_sqrt,
+    )
+    gpflow.set_trainable(mean_function, False)
+    gpflow.set_trainable(m.inducing_variable, False)
+    return m
+
+
 def vgp_model(x: tf.Tensor, y: tf.Tensor, num_latent_gps: int = 1) -> VGP:
     likelihood = gpflow.likelihoods.Gaussian()
     kernel = gpflow.kernels.Matern32()