#!/usr/bin/env python3
from typing import Any, Optional
import torch
from linear_operator import to_dense
from linear_operator.operators import DiagLinearOperator, LinearOperator, TriangularLinearOperator
from linear_operator.utils.cholesky import psd_safe_cholesky
from torch import LongTensor, Tensor
from ..distributions import MultivariateNormal
from ..models import ApproximateGP
from ..utils.errors import CachingError
from ..utils.memoize import add_to_cache, cached, pop_from_cache
from ..utils.nearest_neighbors import NNUtil
from ._variational_distribution import _VariationalDistribution
from .mean_field_variational_distribution import MeanFieldVariationalDistribution
from .unwhitened_variational_strategy import UnwhitenedVariationalStrategy
[docs]class NNVariationalStrategy(UnwhitenedVariationalStrategy):
r"""
This strategy sets all inducing point locations to observed inputs,
and employs a :math:`k`-nearest-neighbor approximation. It was introduced as the
`Variational Nearest Neighbor Gaussian Processes (VNNGP)` in `Wu et al (2022)`_.
See the `VNNGP tutorial`_ for an example.
VNNGP assumes a k-nearest-neighbor generative process for inducing points :math:`\mathbf u`,
:math:`\mathbf q(\mathbf u) = \prod_{j=1}^M q(u_j | \mathbf u_{n(j)})`
where :math:`n(j)` denotes the indices of :math:`k` nearest neighbors for :math:`u_j` among
:math:`u_1, \cdots, u_{j-1}`. For any test observation :math:`\mathbf f`,
VNNGP makes predictive inference conditioned on its :math:`k` nearest inducing points
:math:`\mathbf u_{n(f)}`, i.e. :math:`p(f|\mathbf u_{n(f)})`.
VNNGP's objective factorizes over inducing points and observations, making stochastic optimization over both
immediately available. After a one-time cost of computing the :math:`k`-nearest neighbor structure,
the training and inference complexity is :math:`O(k^3)`.
Since VNNGP uses observations as inducing points, it is a user choice to either (1)
use the same mini-batch of inducing points and observations (recommended),
or (2) use different mini-batches of inducing points and observations. See the `VNNGP tutorial`_ for
implementation and comparison.
.. note::
The current implementation only supports :obj:`~gpytorch.variational.MeanFieldVariationalDistribution`.
We recommend installing the `faiss`_ library (requiring separate package installment)
for nearest neighbor search, which is significantly faster than the `scikit-learn` nearest neighbor search.
GPyTorch will automatically use `faiss` if it is installed, but will revert to `scikit-learn` otherwise.
Different inducing point orderings will produce in different nearest neighbor approximations.
:param ~gpytorch.models.ApproximateGP model: Model this strategy is applied to.
Typically passed in when the VariationalStrategy is created in the
__init__ method of the user defined model.
:param inducing_points: Tensor containing a set of inducing
points to use for variational inference.
:param variational_distribution: A
VariationalDistribution object that represents the form of the variational distribution :math:`q(\mathbf u)`
:param k: Number of nearest neighbors.
:param training_batch_size: The number of data points that will be in the training batch size.
:param jitter_val: Amount of diagonal jitter to add for Cholesky factorization numerical stability
.. _Wu et al (2022):
https://arxiv.org/pdf/2202.01694.pdf
.. _VNNGP tutorial:
examples/04_Variational_and_Approximate_GPs/VNNGP.html
.. _faiss:
https://github.com/facebookresearch/faiss
"""
def __init__(
self,
model: ApproximateGP,
inducing_points: Tensor,
variational_distribution: _VariationalDistribution,
k: int,
training_batch_size: int,
jitter_val: Optional[float] = None,
):
assert isinstance(
variational_distribution, MeanFieldVariationalDistribution
), "Currently, NNVariationalStrategy only supports MeanFieldVariationalDistribution."
super().__init__(
model, inducing_points, variational_distribution, learn_inducing_locations=False, jitter_val=jitter_val
)
# Make sure we don't try to initialize variational parameters - because of minibatching
self.variational_params_initialized.fill_(1)
# Model
object.__setattr__(self, "model", model)
self.inducing_points = inducing_points
self.M: int = inducing_points.shape[-2]
self.D: int = inducing_points.shape[-1]
self.k = k
assert self.k <= self.M, (
f"Number of nearest neighbors k must be smaller than or equal to number of inducing points, "
f"but got k = {k}, M = {self.M}."
)
self._inducing_batch_shape: torch.Size = inducing_points.shape[:-2]
self._model_batch_shape: torch.Size = self._variational_distribution.variational_mean.shape[:-1]
self._batch_shape: torch.Size = torch.broadcast_shapes(self._inducing_batch_shape, self._model_batch_shape)
self.nn_util: NNUtil = NNUtil(
k, dim=self.D, batch_shape=self._inducing_batch_shape, device=inducing_points.device
)
self._compute_nn()
self.training_batch_size = training_batch_size
self._set_training_iterator()
@property
@cached(name="prior_distribution_memo")
def prior_distribution(self) -> MultivariateNormal:
out = self.model.forward(self.inducing_points)
res = MultivariateNormal(out.mean, out.lazy_covariance_matrix.add_jitter(self.jitter_val))
return res
def _cholesky_factor(self, induc_induc_covar: LinearOperator) -> TriangularLinearOperator:
# Uncached version
L = psd_safe_cholesky(to_dense(induc_induc_covar))
return TriangularLinearOperator(L)
def __call__(self, x: Tensor, prior: bool = False, **kwargs: Any) -> MultivariateNormal:
# If we're in prior mode, then we're done!
if prior:
return self.model.forward(x, **kwargs)
if x is not None:
assert self.inducing_points.shape[:-2] == x.shape[:-2], (
f"x batch shape must matches inducing points batch shape, "
f"but got train data batch shape = {x.shape[:-2]}, "
f"inducing points batch shape = {self.inducing_points.shape[:-2]}."
)
# Delete previously cached items from the training distribution
if self.training:
self._clear_cache()
return self.forward(
x, self.inducing_points, inducing_values=None, variational_inducing_covar=None, **kwargs
)
else:
# Ensure inducing_points and x are the same size
inducing_points = self.inducing_points
return self.forward(x, inducing_points, inducing_values=None, variational_inducing_covar=None, **kwargs)
def forward(
self,
x: Tensor,
inducing_points: Tensor,
inducing_values: Optional[Tensor] = None,
variational_inducing_covar: Optional[LinearOperator] = None,
**kwargs: Any,
) -> MultivariateNormal:
if self.training:
# In training mode, note that the full inducing points set = full training dataset
# Users have the option to choose input None or a tensor of training data for x
# If x is None, will sample training data from inducing points
# Otherwise, will find the indices of inducing points that are equal to x
if x is None:
x_indices = self._get_training_indices()
kl_indices = x_indices
predictive_mean = self._variational_distribution.variational_mean[..., x_indices]
predictive_var = self._variational_distribution._variational_stddev[..., x_indices] ** 2
else:
# find the indices of inducing points that correspond to x
x_indices = self.nn_util.find_nn_idx(x.float(), k=1).squeeze(-1) # (*inducing_batch_shape, batch_size)
expanded_x_indices = x_indices.expand(*self._batch_shape, x_indices.shape[-1])
expanded_variational_mean = self._variational_distribution.variational_mean.expand(
*self._batch_shape, self.M
)
expanded_variational_var = (
self._variational_distribution._variational_stddev.expand(*self._batch_shape, self.M) ** 2
)
predictive_mean = expanded_variational_mean.gather(-1, expanded_x_indices)
predictive_var = expanded_variational_var.gather(-1, expanded_x_indices)
# sample a different indices for stochastic estimation of kl
kl_indices = self._get_training_indices()
kl = self._kl_divergence(kl_indices)
add_to_cache(self, "kl_divergence_memo", kl)
return MultivariateNormal(predictive_mean, DiagLinearOperator(predictive_var))
else:
nn_indices = self.nn_util.find_nn_idx(x.float())
x_batch_shape = x.shape[:-2]
x_bsz = x.shape[-2]
assert nn_indices.shape == (*x_batch_shape, x_bsz, self.k), nn_indices.shape
expanded_nn_indices = nn_indices.unsqueeze(-1).expand(*x_batch_shape, x_bsz, self.k, self.D)
expanded_inducing_points = inducing_points.unsqueeze(-2).expand(*x_batch_shape, self.M, self.k, self.D)
inducing_points = expanded_inducing_points.gather(-3, expanded_nn_indices)
assert inducing_points.shape == (*x_batch_shape, x_bsz, self.k, self.D)
# get variational mean and covar for nearest neighbors
batch_shape = torch.broadcast_shapes(self._model_batch_shape, x_batch_shape)
inducing_values = self._variational_distribution.variational_mean
expanded_inducing_values = inducing_values.unsqueeze(-1).expand(*batch_shape, self.M, self.k)
expanded_nn_indices = nn_indices.expand(*batch_shape, x_bsz, self.k)
inducing_values = expanded_inducing_values.gather(-2, expanded_nn_indices)
assert inducing_values.shape == (*batch_shape, x_bsz, self.k)
variational_stddev = self._variational_distribution._variational_stddev
assert variational_stddev.shape == (*self._model_batch_shape, self.M)
expanded_variational_stddev = variational_stddev.unsqueeze(-1).expand(*batch_shape, self.M, self.k)
variational_inducing_covar = expanded_variational_stddev.gather(-2, expanded_nn_indices) ** 2
assert variational_inducing_covar.shape == (*batch_shape, x_bsz, self.k)
variational_inducing_covar = DiagLinearOperator(variational_inducing_covar)
assert variational_inducing_covar.shape == (*batch_shape, x_bsz, self.k, self.k)
# Make everything batch mode
x = x.unsqueeze(-2)
assert x.shape == (*x_batch_shape, x_bsz, 1, self.D)
# Compute forward mode in the standard way
dist = super().forward(x, inducing_points, inducing_values, variational_inducing_covar, **kwargs)
predictive_mean = dist.mean # (*batch_shape, x_bsz, 1)
predictive_covar = dist.covariance_matrix # (*batch_shape, x_bsz, 1, 1)
# Undo batch mode
predictive_mean = predictive_mean.squeeze(-1)
predictive_var = predictive_covar.squeeze(-2).squeeze(-1)
assert predictive_var.shape == predictive_covar.shape[:-2]
assert predictive_mean.shape == predictive_covar.shape[:-2]
# Return the distribution
return MultivariateNormal(predictive_mean, DiagLinearOperator(predictive_var))
def _set_training_iterator(self) -> None:
self._training_indices_iter = 0
training_indices = torch.randperm(self.M - self.k, device=self.inducing_points.device) + self.k
self._training_indices_iterator = (torch.arange(self.k),) + training_indices.split(self.training_batch_size)
self._total_training_batches = len(self._training_indices_iterator)
def _get_training_indices(self) -> LongTensor:
self.current_training_indices = self._training_indices_iterator[self._training_indices_iter]
self._training_indices_iter += 1
if self._training_indices_iter == self._total_training_batches:
self._set_training_iterator()
return self.current_training_indices
def _firstk_kl_helper(self) -> Tensor:
# Compute the KL divergence for first k inducing points
train_x_firstk = self.inducing_points[..., : self.k, :]
full_output = self.model.forward(train_x_firstk)
induc_mean, induc_induc_covar = full_output.mean, full_output.lazy_covariance_matrix
induc_induc_covar = induc_induc_covar.add_jitter(self.jitter_val)
prior_dist = MultivariateNormal(induc_mean, induc_induc_covar)
inducing_values = self._variational_distribution.variational_mean[..., : self.k]
variational_covar_fisrtk = self._variational_distribution._variational_stddev[..., : self.k] ** 2
variational_inducing_covar = DiagLinearOperator(variational_covar_fisrtk)
variational_distribution = MultivariateNormal(inducing_values, variational_inducing_covar)
kl = torch.distributions.kl.kl_divergence(variational_distribution, prior_dist) # model_batch_shape
return kl
def _stochastic_kl_helper(self, kl_indices: Tensor) -> Tensor:
# Compute the KL divergence for a mini batch of the rest M-1 inducing points
# See paper appendix for kl breakdown
kl_bs = len(kl_indices)
variational_mean = self._variational_distribution.variational_mean
variational_stddev = self._variational_distribution._variational_stddev
# compute logdet_q
inducing_point_log_variational_covar = (variational_stddev[..., kl_indices] ** 2).log()
logdet_q = torch.sum(inducing_point_log_variational_covar, dim=-1)
# Select a mini-batch of inducing points according to kl_indices, and their k-nearest neighbors
inducing_points = self.inducing_points[..., kl_indices, :]
nearest_neighbor_indices = self.nn_xinduce_idx[..., kl_indices - self.k, :].to(inducing_points.device)
expanded_inducing_points_all = self.inducing_points.unsqueeze(-2).expand(
*self._inducing_batch_shape, self.M, self.k, self.D
)
expanded_nearest_neighbor_indices = nearest_neighbor_indices.unsqueeze(-1).expand(
*self._inducing_batch_shape, kl_bs, self.k, self.D
)
nearest_neighbors = expanded_inducing_points_all.gather(-3, expanded_nearest_neighbor_indices)
# compute interp_term
cov = self.model.covar_module.forward(nearest_neighbors, nearest_neighbors)
cross_cov = self.model.covar_module.forward(nearest_neighbors, inducing_points.unsqueeze(-2))
interp_term = torch.linalg.solve(
cov + self.jitter_val * torch.eye(self.k, device=self.inducing_points.device), cross_cov
).squeeze(-1)
# compte logdet_p
invquad_term_for_F = torch.sum(interp_term * cross_cov.squeeze(-1), dim=-1)
cov_inducing_points = self.model.covar_module.forward(inducing_points, inducing_points, diag=True)
F = cov_inducing_points - invquad_term_for_F
F = F + self.jitter_val
logdet_p = F.log().sum(dim=-1)
# compute trace_term
expanded_variational_stddev = variational_stddev.unsqueeze(-1).expand(*self._batch_shape, self.M, self.k)
expanded_variational_mean = variational_mean.unsqueeze(-1).expand(*self._batch_shape, self.M, self.k)
expanded_nearest_neighbor_indices = nearest_neighbor_indices.expand(*self._batch_shape, kl_bs, self.k)
nearest_neighbor_variational_covar = (
expanded_variational_stddev.gather(-2, expanded_nearest_neighbor_indices) ** 2
)
bjsquared_s = torch.sum(interp_term**2 * nearest_neighbor_variational_covar, dim=-1)
inducing_point_covar = variational_stddev[..., kl_indices] ** 2
trace_term = (1.0 / F * (bjsquared_s + inducing_point_covar)).sum(dim=-1)
# compute invquad_term
nearest_neighbor_variational_mean = expanded_variational_mean.gather(-2, expanded_nearest_neighbor_indices)
Bj_m = torch.sum(interp_term * nearest_neighbor_variational_mean, dim=-1)
inducing_point_variational_mean = variational_mean[..., kl_indices]
invquad_term = torch.sum((inducing_point_variational_mean - Bj_m) ** 2 / F, dim=-1)
kl = (logdet_p - logdet_q - kl_bs + trace_term + invquad_term) * (1.0 / 2)
assert kl.shape == self._batch_shape, kl.shape
kl = kl.mean()
return kl
def _kl_divergence(
self, kl_indices: Optional[LongTensor] = None, compute_full: bool = False, batch_size: Optional[int] = None
) -> Tensor:
if compute_full:
if batch_size is None:
batch_size = self.training_batch_size
kl = self._firstk_kl_helper()
for kl_indices in torch.split(torch.arange(self.k, self.M), batch_size):
kl += self._stochastic_kl_helper(kl_indices)
else:
assert kl_indices is not None
if (self._training_indices_iter == 1) or (self.M == self.k):
assert len(kl_indices) == self.k, (
f"kl_indices sould be the first batch data of length k, "
f"but got len(kl_indices) = {len(kl_indices)} and k = {self.k}."
)
kl = self._firstk_kl_helper() * self.M / self.k
else:
kl = self._stochastic_kl_helper(kl_indices) * self.M / len(kl_indices)
return kl
def kl_divergence(self) -> Tensor:
try:
return pop_from_cache(self, "kl_divergence_memo")
except CachingError:
raise RuntimeError("KL Divergence of variational strategy was called before nearest neighbors were set.")
def _compute_nn(self) -> "NNVariationalStrategy":
with torch.no_grad():
inducing_points_fl = self.inducing_points.data.float()
self.nn_util.set_nn_idx(inducing_points_fl)
self.nn_xinduce_idx = self.nn_util.build_sequential_nn_idx(inducing_points_fl)
return self