tensorflow_probability/python/distributions/variational_gaussian_process.py
def variational_loss(self,
                       observations,
                       observation_index_points=None,
                       kl_weight=1.,
                       name='variational_loss'):
    """Variational loss for the VGP.

    Given `observations` and `observation_index_points`, compute the
    negative variational lower bound as specified in [Hensman, 2013][1].

    Args:
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below). If
        set to `None` uses `index_points` as the origin for observations.
        Default value: None.
      kl_weight: Amount by which to scale the KL divergence loss between prior
        and posterior.
        Default value: 1.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "GaussianProcess".
    Returns:
      loss: Scalar tensor representing the negative variational lower bound.
        Can be directly used in a `tf.Optimizer`.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
         https://arxiv.org/abs/1309.6835
    """

    with tf.name_scope(name or 'variational_gp_loss'):
      if observation_index_points is None:
        observation_index_points = self._index_points
      observation_index_points = tf.convert_to_tensor(
          value=observation_index_points, dtype=self._dtype,
          name='observation_index_points')
      observations = tf.convert_to_tensor(
          value=observations, dtype=self._dtype, name='observations')
      kl_weight = tf.convert_to_tensor(
          value=kl_weight, dtype=self._dtype,
          name='kl_weight')

      # The variational loss is a negative ELBO. The ELBO can be broken down
      # into three terms:
      #  1. a likelihood term
      #  2. a trace term arising from the covariance of the posterior predictive

      kzx = self.kernel.matrix(self._inducing_index_points,
                               observation_index_points)

      kzx_linop = tf.linalg.LinearOperatorFullMatrix(kzx)
      loc = (self._mean_fn(observation_index_points) +
             kzx_linop.matvec(self._kzz_inv_varloc, adjoint=True))

      likelihood = independent.Independent(
          normal.Normal(
              loc=loc,
              scale=tf.sqrt(self._observation_noise_variance + self._jitter),
              name='NormalLikelihood'),
          reinterpreted_batch_ndims=1)
      obs_ll = likelihood.log_prob(observations)

      chol_kzz_linop = tf.linalg.LinearOperatorLowerTriangular(self._chol_kzz)
      chol_kzz_inv_kzx = chol_kzz_linop.solve(kzx)
      kzz_inv_kzx = chol_kzz_linop.solve(chol_kzz_inv_kzx, adjoint=True)

      kxx_diag = tf.linalg.diag_part(
          self.kernel.matrix(
              observation_index_points, observation_index_points))
      ktilde_trace_term = (
          tf.reduce_sum(input_tensor=kxx_diag, axis=-1) -
          tf.reduce_sum(input_tensor=chol_kzz_inv_kzx ** 2, axis=[-2, -1]))

      # Tr(SB)
      # where S = A A.T, A = variational_inducing_observations_scale
      # and B = Kzz^-1 Kzx Kzx.T Kzz^-1
      #
      # Now Tr(SB) = Tr(A A.T Kzz^-1 Kzx Kzx.T Kzz^-1)
      #            = Tr(A.T Kzz^-1 Kzx Kzx.T Kzz^-1 A)
      #            = sum_ij (A.T Kzz^-1 Kzx)_{ij}^2
      other_trace_term = tf.reduce_sum(
          input_tensor=(
              self._variational_inducing_observations_posterior.scale.matmul(
                  kzz_inv_kzx) ** 2),
          axis=[-2, -1])

      trace_term = (.5 * (ktilde_trace_term + other_trace_term) /
                    self._observation_noise_variance)

      inducing_prior = gaussian_process.GaussianProcess(
          kernel=self._kernel,
          mean_fn=self._mean_fn,
          index_points=self._inducing_index_points,
          observation_noise_variance=self._observation_noise_variance)

      kl_term = kl_weight * kullback_leibler.kl_divergence(
          self._variational_inducing_observations_posterior,
          inducing_prior)

      lower_bound = (obs_ll - trace_term - kl_term)

      return -tf.reduce_mean(input_tensor=lower_bound)
def variational_loss(self,
                       observations,
                       observation_index_points=None,
                       kl_weight=1.,
                       name='variational_loss'):
    """Variational loss for the VGP.

    Given `observations` and `observation_index_points`, compute the
    negative variational lower bound as specified in [Hensman, 2013][1].

    Args:
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below). If
        set to `None` uses `index_points` as the origin for observations.
        Default value: None.
      kl_weight: Amount by which to scale the KL divergence loss between prior
        and posterior.
        Default value: 1.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "GaussianProcess".
    Returns:
      loss: Scalar tensor representing the negative variational lower bound.
        Can be directly used in a `tf.Optimizer`.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
         https://arxiv.org/abs/1309.6835
    """

    with tf.name_scope(name or 'variational_gp_loss'):
      if observation_index_points is None:
        observation_index_points = self._index_points
      observation_index_points = tf.convert_to_tensor(
          value=observation_index_points, dtype=self._dtype,
          name='observation_index_points')
      observations = tf.convert_to_tensor(
          value=observations, dtype=self._dtype, name='observations')
      kl_weight = tf.convert_to_tensor(
          value=kl_weight, dtype=self._dtype,
          name='kl_weight')

      # The variational loss is a negative ELBO. The ELBO can be broken down
      # into three terms:
      #  1. a likelihood term
      #  2. a trace term arising from the covariance of the posterior predictive

      kzx = self.kernel.matrix(self._inducing_index_points,
                               observation_index_points)

      kzx_linop = tf.linalg.LinearOperatorFullMatrix(kzx)
      loc = (self._mean_fn(observation_index_points) +
             kzx_linop.matvec(self._kzz_inv_varloc, adjoint=True))

      likelihood = independent.Independent(
          normal.Normal(
              loc=loc,
              scale=tf.sqrt(self._observation_noise_variance + self._jitter),
              name='NormalLikelihood'),
          reinterpreted_batch_ndims=1)
      obs_ll = likelihood.log_prob(observations)

      chol_kzz_linop = tf.linalg.LinearOperatorLowerTriangular(self._chol_kzz)
      chol_kzz_inv_kzx = chol_kzz_linop.solve(kzx)
      kzz_inv_kzx = chol_kzz_linop.solve(chol_kzz_inv_kzx, adjoint=True)

      kxx_diag = tf.linalg.diag_part(
          self.kernel.matrix(
              observation_index_points, observation_index_points))
      ktilde_trace_term = (
          tf.reduce_sum(input_tensor=kxx_diag, axis=-1) -
          tf.reduce_sum(input_tensor=chol_kzz_inv_kzx ** 2, axis=[-2, -1]))

      # Tr(SB)
      # where S = A A.T, A = variational_inducing_observations_scale
      # and B = Kzz^-1 Kzx Kzx.T Kzz^-1
      #
      # Now Tr(SB) = Tr(A A.T Kzz^-1 Kzx Kzx.T Kzz^-1)
      #            = Tr(A.T Kzz^-1 Kzx Kzx.T Kzz^-1 A)
      #            = sum_ij (A.T Kzz^-1 Kzx)_{ij}^2
      other_trace_term = tf.reduce_sum(
          input_tensor=(
              self._variational_inducing_observations_posterior.scale.matmul(
                  kzz_inv_kzx) ** 2),
          axis=[-2, -1])

      trace_term = (.5 * (ktilde_trace_term + other_trace_term) /
                    self._observation_noise_variance)

      inducing_prior = gaussian_process.GaussianProcess(
          kernel=self._kernel,
          mean_fn=self._mean_fn,
          index_points=self._inducing_index_points,
          observation_noise_variance=self._observation_noise_variance)

      kl_term = kl_weight * kullback_leibler.kl_divergence(
          self._variational_inducing_observations_posterior,
          inducing_prior)

      lower_bound = (obs_ll - trace_term - kl_term)

      return -tf.reduce_mean(input_tensor=lower_bound)
Variational loss for the VGP.

    Given `observations` and `observation_index_points`, compute the
    negative variational lower bound as specified in [Hensman, 2013][1].

    Args:
      observations: `float` `Tensor` representing collection, or batch of
        collections, of observations corresponding to
        `observation_index_points`. Shape has the form `[b1, ..., bB, e]`, which
        must be brodcastable with the batch and example shapes of
        `observation_index_points`. The batch shape `[b1, ..., bB]` must be
        broadcastable with the shapes of all other batched parameters
        (`kernel.batch_shape`, `observation_index_points`, etc.).
      observation_index_points: `float` `Tensor` representing finite (batch of)
        vector(s) of points where observations are defined. Shape has the
        form `[b1, ..., bB, e1, f1, ..., fF]` where `F` is the number of feature
        dimensions and must equal `kernel.feature_ndims` and `e1` is the number
        (size) of index points in each batch (we denote it `e1` to distinguish
        it from the numer of inducing index points, denoted `e2` below). If
        set to `None` uses `index_points` as the origin for observations.
        Default value: None.
      kl_weight: Amount by which to scale the KL divergence loss between prior
        and posterior.
        Default value: 1.
      name: Python `str` name prefixed to Ops created by this class.
        Default value: "GaussianProcess".
    Returns:
      loss: Scalar tensor representing the negative variational lower bound.
        Can be directly used in a `tf.Optimizer`.
    Raises:
      ValueError: if `mean_fn` is not `None` and is not callable.

    #### References

    [1]: Hensman, J., Lawrence, N. "Gaussian Processes for Big Data", 2013
         https://arxiv.org/abs/1309.6835