Repository: secondmind-labs/trieste

Issue:
bug with `deep_ensemble_trajectory` diversify option
Trajectory sampler at the moment with `diversify` option samples quantiles from randomly chosen networks and uses that as a tproxy for the trajectory - this is then repeated for the whole batch and idea was to use it to be able to get more trajectories than the networks we have in the ensemble.

This was not correct idea unfortunately, as we end up not drawing from distributions that contain epistemic/estimation uncertainty - this part in deep ensembles comes from combining all the networks. This particularly is problematic with noiseless data where variances of individual network vanish quickly and batch becomes not diverse enough very quickly. 

The fix is easy, we should sample trajectories as quantiles from the approximate posterior of the whole ensemble.

Discussion:

Patch:
diff --git a/trieste/models/gpflux/__init__.py b/trieste/models/gpflux/__init__.py
index 17518b51..ba8c7a51 100644
--- a/trieste/models/gpflux/__init__.py
+++ b/trieste/models/gpflux/__init__.py
@@ -21,3 +21,10 @@ from . import config
 from .builders import build_vanilla_deep_gp
 from .interface import GPfluxPredictor
 from .models import DeepGaussianProcess
+from .sampler import (
+    DeepGaussianProcessDecoupledLayer,
+    DeepGaussianProcessDecoupledTrajectorySampler,
+    DeepGaussianProcessReparamSampler,
+    ResampleableDecoupledDeepGaussianProcessFeatureFunctions,
+    dgp_feature_decomposition_trajectory,
+)
diff --git a/trieste/models/keras/__init__.py b/trieste/models/keras/__init__.py
index 312710bd..d05d3b62 100644
--- a/trieste/models/keras/__init__.py
+++ b/trieste/models/keras/__init__.py
@@ -25,5 +25,5 @@ from .architectures import GaussianNetwork, KerasEnsemble, KerasEnsembleNetwork
 from .builders import build_keras_ensemble
 from .interface import DeepEnsembleModel, KerasPredictor
 from .models import DeepEnsemble
-from .sampler import DeepEnsembleTrajectorySampler
+from .sampler import DeepEnsembleTrajectorySampler, deep_ensemble_trajectory
 from .utils import get_tensor_spec_from_data, negative_log_likelihood, sample_with_replacement
diff --git a/trieste/models/keras/interface.py b/trieste/models/keras/interface.py
index 215225f5..b7e24e9e 100644
--- a/trieste/models/keras/interface.py
+++ b/trieste/models/keras/interface.py
@@ -92,6 +92,14 @@ class DeepEnsembleModel(ProbabilisticModel, Protocol):
         """
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def num_outputs(self) -> int:
+        """
+        Returns the number of outputs trained on by each member network.
+        """
+        raise NotImplementedError
+
     @abstractmethod
     def ensemble_distributions(
         self, query_points: TensorType
diff --git a/trieste/models/keras/models.py b/trieste/models/keras/models.py
index 197877df..2ae6b385 100644
--- a/trieste/models/keras/models.py
+++ b/trieste/models/keras/models.py
@@ -232,7 +232,7 @@ class DeepEnsemble(
             + \mu^2_{\theta_m}(\mathbf{x})) - \mu^2_{*}(\mathbf{x})
 
         This method assumes that the final layer in each member of the ensemble is
-        probabilistic, an instance of :class:`¬tfp.distributions.Distribution`. In particular, given
+        probabilistic, an instance of :class:`~tfp.distributions.Distribution`. In particular, given
         the nature of the approximations stated above the final layer should be a Gaussian
         distribution with `mean` and `variance` methods.
 
@@ -318,21 +318,10 @@ class DeepEnsemble(
         Return a trajectory sampler. For :class:`DeepEnsemble`, we use an ensemble
         sampler that randomly picks a network from the ensemble and uses its predicted means
         for generating a trajectory, or optionally randomly sampled quantiles rather than means.
-
-        At the moment only models with single output are supported.
+        Only models with single output are supported with diversify option.
 
         :return: The trajectory sampler.
-        :raise NotImplementedError: If we try to use the sampler with a model that has more than
-            one output.
         """
-        if self.num_outputs > 1:
-            raise NotImplementedError(
-                f"""
-                Trajectory sampler does not currently support models with multiple outputs,
-                however received a model with {self.num_outputs} outputs.
-                """
-            )
-
         return DeepEnsembleTrajectorySampler(self, self._diversify)
 
     def update(self, dataset: Dataset) -> None:
diff --git a/trieste/models/keras/sampler.py b/trieste/models/keras/sampler.py
index ce5c58e4..013b02a5 100644
--- a/trieste/models/keras/sampler.py
+++ b/trieste/models/keras/sampler.py
@@ -19,7 +19,10 @@ of the Trieste's Keras model wrappers.
 
 from __future__ import annotations
 
+from typing import Dict
+
 import tensorflow as tf
+import tensorflow_probability as tfp
 
 from ...types import TensorType
 from ...utils import flatten_leading_dims
@@ -31,13 +34,9 @@ from .utils import sample_model_index
 class DeepEnsembleTrajectorySampler(TrajectorySampler[DeepEnsembleModel]):
     """
     This class builds functions that approximate a trajectory by randomly choosing a network from
-    the ensemble and using its predicted means as a trajectory. Currently we do not
-    support models with multiple outputs.
-
-    Option `diversify` can be used to increase the diversity in case of optimizing very large
-    batches of trajectories. In this case the sampler will use quantiles from distributions rather
-    than means as trajectories. Quantiles are sampled uniformly from a unit interval. When batch
-    size is larger than the ensemble size, multiple quantiles will be used with the same network.
+    the ensemble and using its predicted means as a trajectory. Option `diversify` can be used to
+    increase the diversity in case of optimizing very large batches of trajectories, but this is not
+    supported for multiple output models.
     """
 
     def __init__(self, model: DeepEnsembleModel, diversify: bool = False):
@@ -45,6 +44,9 @@ class DeepEnsembleTrajectorySampler(TrajectorySampler[DeepEnsembleModel]):
         :param model: The ensemble model to sample from.
         :param diversify: Whether to use quantiles from final probabilistic layer as
             trajectories (`False` by default). See class docstring for details.
+        :raise NotImplementedError: If we try to use the model that is not instance of
+            :class:`DeepEnsembleModel` or a multiple output model together with ``diversify``
+            option.
         """
         if not isinstance(model, DeepEnsembleModel):
             raise NotImplementedError(
@@ -58,6 +60,13 @@ class DeepEnsembleTrajectorySampler(TrajectorySampler[DeepEnsembleModel]):
         self._model = model
         self._diversify = diversify
 
+        if self._diversify:
+            if self._model.num_outputs > 1:
+                raise NotImplementedError(
+                    f"EnsembleTrajectorySampler with diversify=True only works single output models"
+                    f", however received a model with {self._model.num_outputs} outputs."
+                )
+
     def __repr__(self) -> str:
         """"""
         return f"{self.__class__.__name__}({self._model!r}"
@@ -104,9 +113,10 @@ class deep_ensemble_trajectory(TrajectoryFunctionClass):
     networks from the ensemble and using their predicted means as trajectories.
 
     Option `diversify` can be used to increase the diversity in case of optimizing very large
-    batches of trajectories. It is experimental feature and hence should be used with caution.
-
-    Models with multiple outputs are not supported.
+    batches of trajectories. Instead of using mean predictions of individual networks quantiles are
+    drawn from Gaussian approximation of the whole ensemble, with randomly chosen quantiles
+    approximating a trajectory. Since quantiles are difficult to compute for multivariate
+    distributions, multiple outputs are not supported with `diversify` option.
     """
 
     def __init__(self, model: DeepEnsembleModel, diversify: bool):
@@ -121,12 +131,14 @@ class deep_ensemble_trajectory(TrajectoryFunctionClass):
 
         self._initialized = tf.Variable(False, trainable=False)
         self._batch_size = tf.Variable(0, dtype=tf.int32, trainable=False)
-        self._indices = tf.Variable(tf.zeros([0], dtype=tf.int32), shape=[None], trainable=False)
 
         if self._diversify:
-            self._sample_size = tf.Variable(0, dtype=tf.int32, trainable=False)
-            self._seeds = tf.Variable(
-                tf.zeros([0, 0], dtype=tf.float32), shape=[None, None], trainable=False
+            self._quantiles = tf.Variable(
+                tf.zeros([0], dtype=tf.float32), shape=[None], trainable=False
+            )
+        else:
+            self._indices = tf.Variable(
+                tf.zeros([0], dtype=tf.int32), shape=[None], trainable=False
             )
 
     @tf.function
@@ -139,9 +151,7 @@ class deep_ensemble_trajectory(TrajectoryFunctionClass):
         """
         if not self._initialized:  # work out desired batch size from input
             self._batch_size.assign(tf.shape(x)[-2])  # B
-            if self._diversify:
-                self._sample_size.assign(-(-self._batch_size // self._ensemble_size))  # B/E
-            self.resample()  # sample network indices/seeds
+            self.resample()  # sample network indices/quantiles
             self._initialized.assign(True)
 
         tf.debugging.assert_equal(
@@ -154,20 +164,13 @@ class deep_ensemble_trajectory(TrajectoryFunctionClass):
             """,
         )
         flat_x, unflatten = flatten_leading_dims(x)  # [N*B, d]
-        ensemble_distributions = self._model.ensemble_distributions(flat_x)
 
         if self._diversify:
-            samples = tf.convert_to_tensor(
-                [
-                    tf.map_fn(dist.quantile, self._seeds[i])
-                    for i, dist in enumerate(ensemble_distributions)
-                ]
-            )  # [E, B/E, N*B, 1]
-            flattened_samples = tf.reshape(
-                samples, [self._ensemble_size * self._sample_size, *samples.shape[2:]]
-            )  # [E*B/E, N*B, 1]
-            predictions = tf.gather(flattened_samples, self._indices)  # [B, N*B, 1]
+            predicted_means, predicted_vars = self._model.predict(flat_x)  # ([N*B, 1], [N*B, 1])
+            normal = tfp.distributions.Normal(predicted_means, tf.sqrt(predicted_vars))
+            predictions = tf.map_fn(normal.quantile, self._quantiles)  # [B, N*B, 1]
         else:
+            ensemble_distributions = self._model.ensemble_distributions(flat_x)
             predicted_means = tf.convert_to_tensor([dist.mean() for dist in ensemble_distributions])
             predictions = tf.gather(predicted_means, self._indices)  # [B, N*B, 1]
 
@@ -188,15 +191,28 @@ class deep_ensemble_trajectory(TrajectoryFunctionClass):
         Efficiently resample network indices in-place, without retracing.
         """
         if self._diversify:
-            indices = tf.random.shuffle(tf.range(self._ensemble_size * self._sample_size))
-            self._indices.assign(indices[: self._batch_size])  # [B]
-            self._seeds.assign(
+            self._quantiles.assign(
                 tf.random.uniform(
-                    shape=(self._ensemble_size, self._sample_size),
+                    shape=(self._batch_size,),
                     minval=0.000001,
                     maxval=0.999999,
                     dtype=tf.float32,
                 )
-            )  # [E, B/E]
+            )  # [B]
         else:
             self._indices.assign(sample_model_index(self._ensemble_size, self._batch_size))  # [B]
+
+    def get_state(self) -> Dict[str, TensorType]:
+        """
+        Return internal state variables.
+        """
+        state = {
+            "initialized": self._initialized,
+            "batch_size": self._batch_size,
+        }
+        if self._diversify:
+            state["quantiles"] = self._quantiles
+        else:
+            state["indices"] = self._indices
+
+        return state


Test patch:
diff --git a/tests/unit/models/keras/test_models.py b/tests/unit/models/keras/test_models.py
index dd9d8238..9a78ac2e 100644
--- a/tests/unit/models/keras/test_models.py
+++ b/tests/unit/models/keras/test_models.py
@@ -125,14 +125,6 @@ def test_deep_ensemble_model_attributes() -> None:
     assert model.model is keras_ensemble.model
 
 
-def test_deep_ensemble_trajectory_sampler_raises_for_multi_output_model() -> None:
-    example_data = empty_dataset([2], [2])
-    model, _, _ = trieste_deep_ensemble_model(example_data, _ENSEMBLE_SIZE, False, False)
-
-    with pytest.raises(NotImplementedError):
-        model.trajectory_sampler()
-
-
 def test_deep_ensemble_ensemble_size_attributes(ensemble_size: int) -> None:
     example_data = empty_dataset([1], [1])
     model, _, _ = trieste_deep_ensemble_model(example_data, ensemble_size, False, False)
diff --git a/tests/unit/models/keras/test_sampler.py b/tests/unit/models/keras/test_sampler.py
index 4e773bbf..9998b773 100644
--- a/tests/unit/models/keras/test_sampler.py
+++ b/tests/unit/models/keras/test_sampler.py
@@ -23,7 +23,11 @@ import tensorflow as tf
 from tests.util.misc import empty_dataset, quadratic, random_seed
 from tests.util.models.keras.models import trieste_deep_ensemble_model
 from trieste.data import Dataset
-from trieste.models.keras import DeepEnsemble, DeepEnsembleTrajectorySampler
+from trieste.models.keras import (
+    DeepEnsemble,
+    DeepEnsembleTrajectorySampler,
+    deep_ensemble_trajectory,
+)
 from trieste.types import TensorType
 
 _ENSEMBLE_SIZE = 3
@@ -44,16 +48,27 @@ def _batch_size_fixture(request: Any) -> int:
     return request.param
 
 
+def test_ensemble_trajectory_sampler_raises_for_multi_output_model() -> None:
+    example_data = empty_dataset([2], [2])
+    model, _, _ = trieste_deep_ensemble_model(example_data, _ENSEMBLE_SIZE)
+
+    with pytest.raises(NotImplementedError):
+        DeepEnsembleTrajectorySampler(model, diversify=True)
+
+
+@pytest.mark.parametrize("diversify, num_outputs", [(True, 1), (False, 1), (False, 2)])
 def test_ensemble_trajectory_sampler_returns_trajectory_function_with_correctly_shaped_output(
     num_evals: int,
     batch_size: int,
     dim: int,
     diversify: bool,
+    num_outputs: int,
 ) -> None:
     """
-    Inputs should be [N,B,d] while output should be [N,B]
+    Inputs should be [N,B,d] while output should be [N,B,M]. Note that for diversify
+    option only single output models are allowed.
     """
-    example_data = empty_dataset([dim], [1])
+    example_data = empty_dataset([dim], [num_outputs])
     test_data = tf.random.uniform([num_evals, batch_size, dim])  # [N, B, d]
 
     model, _, _ = trieste_deep_ensemble_model(example_data, _ENSEMBLE_SIZE)
@@ -61,7 +76,7 @@ def test_ensemble_trajectory_sampler_returns_trajectory_function_with_correctly_
     sampler = DeepEnsembleTrajectorySampler(model, diversify=diversify)
     trajectory = sampler.get_trajectory()
 
-    assert trajectory(test_data).shape == (num_evals, batch_size, 1)
+    assert trajectory(test_data).shape == (num_evals, batch_size, num_outputs)
 
 
 def test_ensemble_trajectory_sampler_returns_deterministic_trajectory(
@@ -340,3 +355,38 @@ def test_ensemble_trajectory_sampler_trajectory_is_continuous(diversify: bool) -
     trajectory = trajectory_sampler.get_trajectory()
 
     npt.assert_array_less(tf.abs(trajectory(test_data + 1e-20) - trajectory(test_data)), 1e-20)
+
+
+def test_ensemble_trajectory_sampler_returns_state(batch_size: int, diversify: bool) -> None:
+    dim = 3
+    num_evals = 10
+
+    example_data = empty_dataset([dim], [1])
+    test_data = tf.random.uniform([num_evals, batch_size, dim])  # [N, B, d]
+
+    model, _, _ = trieste_deep_ensemble_model(example_data, _ENSEMBLE_SIZE)
+
+    sampler = DeepEnsembleTrajectorySampler(model, diversify=diversify)
+    trajectory = cast(deep_ensemble_trajectory, sampler.get_trajectory())
+
+    if diversify:
+        dtype = tf.float32
+        rnd_state_name = "quantiles"
+    else:
+        dtype = tf.int32
+        rnd_state_name = "indices"
+
+    # before calling the trajectory internal state should not be initialized
+    state_pre_call = trajectory.get_state()
+    assert not state_pre_call["initialized"]
+    assert state_pre_call["batch_size"] == 0
+    assert tf.equal(tf.size(state_pre_call[rnd_state_name]), 0)
+    assert state_pre_call[rnd_state_name].dtype == dtype
+
+    # after calling the trajectory internal state should be initialized
+    _ = trajectory(test_data)
+    state_post_call = trajectory.get_state()
+    assert state_post_call["initialized"]
+    assert state_post_call["batch_size"] == batch_size
+    assert tf.equal(tf.size(state_post_call[rnd_state_name]), batch_size)
+    assert state_post_call[rnd_state_name].dtype == dtype