qpytorch 0.1__py3-none-any.whl

qpytorch/variational/tril_natural_variational_distribution.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+
+from typing import Tuple, Union
+
+import torch
+from linear_operator.operators import CholLinearOperator, TriangularLinearOperator
+from torch import Tensor
+from torch.autograd.function import FunctionCtx
+
+from ..distributions import Distribution, MultivariateNormal, MultivariateQExponential
+from .natural_variational_distribution import (
+    _NaturalToMuVarSqrt,
+    _NaturalVariationalDistribution,
+    _phi_for_cholesky_,
+    _triangular_inverse,
+)
+
+
+class TrilNaturalVariationalDistribution(_NaturalVariationalDistribution):
+    r"""A multivariate normal :obj:`~qpytorch.variational._VariationalDistribution`,
+    parameterized by the natural vector, and a triangular decomposition of the
+    natural matrix (which is not the Cholesky).
+
+    .. note::
+       The :obj:`~qpytorch.variational.TrilNaturalVariationalDistribution` should only
+       be used with :obj:`gpytorch.optim.NGD`, or other optimizers
+       that follow exactly the gradient direction.
+
+    .. seealso::
+        The `natural gradient descent tutorial
+        <examples/04_Variational_and_Approximate_GPs/Natural_Gradient_Descent.ipynb>`_
+        for use instructions.
+
+        The :obj:`~qpytorch.variational.NaturalVariationalDistribution`, which
+        needs less iterations to make variational regression converge, at the
+        cost of introducing numerical instability.
+
+    .. note::
+        The relationship of the parameter :math:`\mathbf \Theta_\text{tril_mat}`
+        to the natural parameter :math:`\mathbf \Theta_\text{mat}` from
+        :obj:`~qpytorch.variational.NaturalVariationalDistribution` is
+        :math:`\mathbf \Theta_\text{mat} = -1/2 {\mathbf \Theta_\text{tril_mat}}^T {\mathbf \Theta_\text{tril_mat}}`.
+        Note that this is not the form of the Cholesky decomposition of :math:`\boldsymbol \Theta_\text{mat}`.
+
+    :param int num_inducing_points: Size of the variational distribution. This implies that the variational mean
+        should be this size, and the variational covariance matrix should have this many rows and columns.
+    :param batch_shape: Specifies an optional batch size
+        for the variational parameters. This is useful for example when doing additive variational inference.
+    :type batch_shape: :obj:`torch.Size`, optional
+    :param float mean_init_std: (Default: 1e-3) Standard deviation of gaussian (q-exponential) noise to add to the mean initialization.
+    """
+
+    def __init__(self, num_inducing_points: int, batch_shape: torch.Size = torch.Size([]), mean_init_std: float = 1e-3, **kwargs):
+        super().__init__(num_inducing_points=num_inducing_points, batch_shape=batch_shape, mean_init_std=mean_init_std)
+        scaled_mean_init = torch.zeros(num_inducing_points)
+        neg_prec_init = torch.eye(num_inducing_points, num_inducing_points)
+        scaled_mean_init = scaled_mean_init.repeat(*batch_shape, 1)
+        neg_prec_init = neg_prec_init.repeat(*batch_shape, 1, 1)
+
+        # eta1 and tril_dec(eta2) parameterization of the variational distribution
+        self.register_parameter(name="natural_vec", parameter=torch.nn.Parameter(scaled_mean_init))
+        self.register_parameter(name="natural_tril_mat", parameter=torch.nn.Parameter(neg_prec_init))
+        
+        if 'power' in kwargs: self.power = kwargs.pop('power')
+
+    def forward(self) -> Distribution:
+        mean, chol_covar = _TrilNaturalToMuVarSqrt.apply(self.natural_vec, self.natural_tril_mat)
+        covar = CholLinearOperator(TriangularLinearOperator(chol_covar))
+        if not hasattr(self, 'power'):
+            return MultivariateNormal(mean, covar)
+        else:
+            return MultivariateQExponential(mean, covar, power=self.power)
+
+    def initialize_variational_distribution(self, prior_dist: Union[MultivariateNormal, MultivariateQExponential]) -> None:
+        prior_cov = prior_dist.lazy_covariance_matrix
+        chol = prior_cov.cholesky().to_dense()
+        tril_mat = _triangular_inverse(chol, upper=False)
+
+        natural_vec = prior_cov.solve(prior_dist.mean.unsqueeze(-1)).squeeze(-1)
+        noise = torch.randn_like(natural_vec).mul_(self.mean_init_std)
+
+        self.natural_vec.data.copy_(natural_vec.add_(noise))
+        self.natural_tril_mat.data.copy_(tril_mat)
+
+
+class _TrilNaturalToMuVarSqrt(torch.autograd.Function):
+    @staticmethod
+    def _forward(nat_mean: Tensor, tril_nat_covar: Tensor) -> Tuple[Tensor, Tensor]:
+        L = _triangular_inverse(tril_nat_covar, upper=False)
+        mu = L @ (L.transpose(-1, -2) @ nat_mean.unsqueeze(-1))
+        return mu.squeeze(-1), L
+        # return nat_mean, L
+
+    @staticmethod
+    def forward(ctx: FunctionCtx, nat_mean: Tensor, tril_nat_covar: Tensor) -> Tuple[Tensor, Tensor]:
+        mu, L = _TrilNaturalToMuVarSqrt._forward(nat_mean, tril_nat_covar)
+        ctx.save_for_backward(mu, L, tril_nat_covar)
+        return mu, L
+
+    @staticmethod
+    def backward(ctx: FunctionCtx, dout_dmu: Tensor, dout_dL: Tensor) -> Tuple[Tensor, Tensor]:
+        mu, L, C = ctx.saved_tensors
+        dout_dnat1, dout_dnat2 = _NaturalToMuVarSqrt._backward(dout_dmu, dout_dL, mu, L, C)
+        """
+        Now we need to do the Jacobian-Vector Product for the transformation:
+        L = inv(chol(inv(-2 theta_cov)))
+
+        C^T C = -2 theta_cov
+
+        so we need to do forward differentiation, starting with sensitivity (sensitivities marked with .dots.)
+        .theta_cov. = dout_dnat2
+
+        and ending with sensitivity .C.
+
+        if B = inv(-2 theta_cov) then:
+
+        .B.  =  d inv(-2 theta_cov)/dtheta_cov * .theta_cov.  =  -B (-2 .theta_cov.) B
+
+        if L = chol(B), B = LL^T then (https://homepages.inf.ed.ac.uk/imurray2/pub/16choldiff/choldiff.pdf):
+
+        .L. = L phi(L^{-1} .B. (L^{-1})^T) = L phi(2 L^T .theta_cov. L)
+
+        Then C = inv(L), so
+
+        .C. = -C .L. C = phi(-2 L^T .theta_cov. L)C
+        """
+        A = L.transpose(-2, -1) @ dout_dnat2 @ L
+        phi = _phi_for_cholesky_(A.mul_(-2))
+        dout_dtril = phi @ C
+        return dout_dnat1, dout_dtril

qpytorch/variational/uncorrelated_multitask_variational_strategy.py

@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+
+import warnings
+
+import torch
+from linear_operator.operators import RootLinearOperator
+
+from ..distributions import MultitaskMultivariateQExponential, MultivariateQExponential 
+from ..module import Module
+from ._variational_strategy import _VariationalStrategy
+
+
+class UncorrelatedMultitaskVariationalStrategy(_VariationalStrategy):
+    """
+    UncorrelatedMultitaskVariationalStrategy wraps an existing
+    :obj:`~qpytorch.variational.VariationalStrategy` to produce vector-valued (multi-task)
+    output distributions. Each task will be uncorrelated to one another.
+
+    The output will either be a :obj:`~qpytorch.distributions.MultitaskMultivariateQExponential` distribution
+    (if we wish to evaluate all tasks for each input) or a :obj:`~qpytorch.distributions.MultivariateQExponential`
+    (if we wish to evaluate a single task for each input).
+
+    The base variational strategy is assumed to operate on a batch of QEPs. One of the batch
+    dimensions corresponds to the multiple tasks.
+
+    :param ~qpytorch.variational.VariationalStrategy base_variational_strategy: Base variational strategy
+    :param int num_tasks: Number of tasks. Should correspond to the batch size of task_dim.
+    :param int task_dim: (Default: -1) Which batch dimension is the task dimension
+    """
+
+    def __init__(self, base_variational_strategy, num_tasks, task_dim=-1):
+        Module.__init__(self)
+        self.base_variational_strategy = base_variational_strategy
+        self.task_dim = task_dim
+        self.num_tasks = num_tasks
+
+    @property
+    def prior_distribution(self):
+        return self.base_variational_strategy.prior_distribution
+
+    @property
+    def variational_distribution(self):
+        return self.base_variational_strategy.variational_distribution
+
+    @property
+    def variational_params_initialized(self):
+        return self.base_variational_strategy.variational_params_initialized
+
+    def kl_divergence(self):
+        return super().kl_divergence().sum(dim=-1)
+
+    def __call__(self, x, task_indices=None, prior=False, **kwargs):
+        r"""
+        See :class:`LMCVariationalStrategy`.
+        """
+        function_dist = self.base_variational_strategy(x, prior=prior, **kwargs)
+
+        if task_indices is None:
+            # Every data point will get an output for each task
+            if (
+                self.task_dim > 0
+                and self.task_dim > len(function_dist.batch_shape)
+                or self.task_dim < 0
+                and self.task_dim + len(function_dist.batch_shape) < 0
+            ):
+                return MultitaskMultivariateQExponential.from_repeated_qep(function_dist, num_tasks=self.num_tasks)
+            else:
+                function_dist = MultitaskMultivariateQExponential.from_batch_qep(function_dist, task_dim=self.task_dim)
+                assert function_dist.event_shape[-1] == self.num_tasks
+                return function_dist
+
+        else:
+            # Each data point will get a single output corresponding to a single task
+
+            if self.task_dim > 0:
+                raise RuntimeError(f"task_dim must be a negative indexed batch dimension: got {self.task_dim}.")
+            num_batch = len(function_dist.batch_shape)
+            task_dim = num_batch + self.task_dim
+
+            # Create a mask to choose specific task assignment
+            shape = list(function_dist.batch_shape + function_dist.event_shape)
+            shape[task_dim] = 1
+            task_indices = task_indices.expand(shape).squeeze(task_dim)
+
+            # Create a mask to choose specific task assignment
+            task_mask = torch.nn.functional.one_hot(task_indices, num_classes=self.num_tasks)
+            task_mask = task_mask.permute(*range(0, task_dim), *range(task_dim + 1, num_batch + 1), task_dim)
+
+            mean = (function_dist.mean * task_mask).sum(task_dim)
+            covar = (function_dist.lazy_covariance_matrix * RootLinearOperator(task_mask[..., None])).sum(task_dim)
+            return MultivariateQExponential(mean, covar, power=function_dist.power)
+
+
+class MultitaskVariationalStrategy(UncorrelatedMultitaskVariationalStrategy):
+    """
+    UncorrelatedMultitaskVariationalStrategy wraps an existing
+    :obj:`~qpytorch.variational.VariationalStrategy`
+    to produce a :obj:`~qpytorch.variational.MultitaskMultivariateQExponential` distribution.
+    All outputs will be uncorrelated to one another.
+
+    The base variational strategy is assumed to operate on a batch of QEPs. One of the batch
+    dimensions corresponds to the multiple tasks.
+
+    :param ~qpytorch.variational.VariationalStrategy base_variational_strategy: Base variational strategy
+    :param int num_tasks: Number of tasks. Should correspond to the batch size of task_dim.
+    :param int task_dim: (Default: -1) Which batch dimension is the task dimension
+    """
+
+    def __init__(self, base_variational_strategy, num_tasks, task_dim=-1):
+        warnings.warn(
+            "MultitaskVariationalStrategy has been renamed to UncorrelatedMultitaskVariationalStrategy",
+            DeprecationWarning,
+        )
+        super().__init__(base_variational_strategy, num_tasks, task_dim=-1)

qpytorch/variational/unwhitened_variational_strategy.py

@@ -0,0 +1,225 @@
+#!/usr/bin/env python3
+
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+from linear_operator import to_dense
+from linear_operator.operators import (
+    CholLinearOperator,
+    DiagLinearOperator,
+    LinearOperator,
+    PsdSumLinearOperator,
+    RootLinearOperator,
+    TriangularLinearOperator,
+    ZeroLinearOperator,
+)
+from linear_operator.utils.cholesky import psd_safe_cholesky
+from linear_operator.utils.errors import NotPSDError
+from torch import Tensor
+
+from .. import settings
+from ..distributions import MultivariateNormal, MultivariateQExponential
+from gpytorch.utils.memoize import add_to_cache, cached
+from ._variational_strategy import _VariationalStrategy
+from .cholesky_variational_distribution import CholeskyVariationalDistribution
+
+
+class UnwhitenedVariationalStrategy(_VariationalStrategy):
+    r"""
+    Similar to :obj:`~qpytorch.variational.VariationalStrategy`, but does not perform the
+    whitening operation. In almost all cases :obj:`~qpytorch.variational.VariationalStrategy`
+    is preferable, with a few exceptions:
+
+    - When the inducing points are exactly equal to the training points (i.e. :math:`\mathbf Z = \mathbf X`).
+      Unwhitened models are faster in this case.
+
+    - When the number of inducing points is very large (e.g. >2000). Unwhitened models can use CG for faster
+      computation.
+
+    :param ~model: Model this strategy is applied to.
+        Typically passed in when the VariationalStrategy is created in the
+        __init__ method of the user defined model.
+        It should contain power if Q-Exponential distribution is involved in.
+    :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 learn_inducing_locations: (default True): Whether or not
+        the inducing point locations :math:`\mathbf Z` should be learned (i.e. are they
+        parameters of the model).
+    :param jitter_val: Amount of diagonal jitter to add for Cholesky factorization numerical stability
+    """
+    has_fantasy_strategy = True
+
+    @cached(name="cholesky_factor", ignore_args=True)
+    def _cholesky_factor(self, induc_induc_covar: LinearOperator) -> TriangularLinearOperator:
+        # Maybe used - if we're not using CG
+        L = psd_safe_cholesky(to_dense(induc_induc_covar))
+        return TriangularLinearOperator(L)
+
+    @property
+    @cached(name="prior_distribution_memo")
+    def prior_distribution(self) -> Union[MultivariateNormal, MultivariateQExponential]:
+        out = self.model.forward(self.inducing_points)
+        if hasattr(self.model, 'power'):
+            res = MultivariateQExponential(out.mean, out.lazy_covariance_matrix.add_jitter(), power=self.model.power)
+        else:
+            res = MultivariateNormal(out.mean, out.lazy_covariance_matrix.add_jitter())
+        return res
+
+    @property
+    @cached(name="pseudo_points_memo")
+    def pseudo_points(self) -> Tuple[Tensor, Tensor]:
+        # TODO: implement for other distributions
+        # retrieve the variational mean, m and covariance matrix, S.
+        if not isinstance(self._variational_distribution, CholeskyVariationalDistribution):
+            raise NotImplementedError(
+                "Only CholeskyVariationalDistribution has pseudo-point support currently, ",
+                "but your _variational_distribution is a ",
+                self._variational_distribution.__name__,
+            )
+
+        # retrieve the variational mean, m and covariance matrix, S.
+        var_cov_root = TriangularLinearOperator(self._variational_distribution.chol_variational_covar)
+        var_cov = CholLinearOperator(var_cov_root)
+        var_mean = self.variational_distribution.mean  # .unsqueeze(-1)
+        if var_mean.shape[-1] != 1:
+            var_mean = var_mean.unsqueeze(-1)
+
+        # R = K - S
+        Kmm = self.model.covar_module(self.inducing_points)
+        res = Kmm - var_cov
+
+        cov_diff = res
+
+        # D_a = (S^{-1} - K^{-1})^{-1} = S + S R^{-1} S
+        # note that in the whitened case R = I - S, unwhitened R = K - S
+        # we compute (R R^{T})^{-1} R^T S for stability reasons as R is probably not PSD.
+        eval_lhs = var_cov.to_dense()
+        eval_rhs = cov_diff.transpose(-1, -2).matmul(eval_lhs)
+        inner_term = cov_diff.matmul(cov_diff.transpose(-1, -2))
+        # TODO: flag the jitter here
+        inner_solve = inner_term.add_jitter(self.jitter_val).solve(eval_rhs, eval_lhs.transpose(-1, -2))
+        inducing_covar = var_cov + inner_solve
+
+        # mean term: D_a S^{-1} m
+        # unwhitened: (S - S R^{-1} S) S^{-1} m = (I - S R^{-1}) m
+        rhs = cov_diff.transpose(-1, -2).matmul(var_mean)
+        inner_rhs_mean_solve = inner_term.add_jitter(self.jitter_val).solve(rhs)
+        pseudo_target_mean = var_mean + var_cov.matmul(inner_rhs_mean_solve)
+
+        # ensure inducing covar is psd
+        try:
+            pseudo_target_covar = CholLinearOperator(inducing_covar.add_jitter(self.jitter_val).cholesky()).to_dense()
+        except NotPSDError:
+            from linear_operator.operators import DiagLinearOperator
+
+            evals, evecs = torch.linalg.eigh(inducing_covar)
+            pseudo_target_covar = (
+                evecs.matmul(DiagLinearOperator(evals + self.jitter_val)).matmul(evecs.transpose(-1, -2)).to_dense()
+            )
+
+        return pseudo_target_covar, pseudo_target_mean
+
+    def forward(
+        self,
+        x: Tensor,
+        inducing_points: Tensor,
+        inducing_values: Tensor,
+        variational_inducing_covar: Optional[LinearOperator] = None,
+        **kwargs,
+    ) -> Union[MultivariateNormal, MultivariateQExponential]:
+        # If our points equal the inducing points, we're done
+        if torch.equal(x, inducing_points):
+            if variational_inducing_covar is None:
+                raise RuntimeError
+            else:
+                if hasattr(self.model, 'power'):
+                    return MultivariateQExponential(inducing_values, variational_inducing_covar, power=self.model.power)
+                else:
+                    return MultivariateNormal(inducing_values, variational_inducing_covar)
+
+        # Otherwise, we have to marginalize
+        num_induc = inducing_points.size(-2)
+        full_inputs = torch.cat([inducing_points, x], dim=-2)
+        full_output = self.model.forward(full_inputs)
+        full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix
+
+        # Mean terms
+        test_mean = full_mean[..., num_induc:]
+        induc_mean = full_mean[..., :num_induc]
+        mean_diff = (inducing_values - induc_mean).unsqueeze(-1)
+
+        # Covariance terms
+        induc_induc_covar = full_covar[..., :num_induc, :num_induc].add_jitter(self.jitter_val)
+        induc_data_covar = full_covar[..., :num_induc, num_induc:].to_dense()
+        data_data_covar = full_covar[..., num_induc:, num_induc:]
+
+        # Compute Cholesky factorization of inducing covariance matrix
+        if settings.fast_computations.log_prob.off() or (num_induc <= settings.max_cholesky_size.value()):
+            induc_induc_covar = CholLinearOperator(self._cholesky_factor(induc_induc_covar))
+
+        # If we are making predictions and don't need variances, we can do things very quickly.
+        if not self.training and settings.skip_posterior_variances.on():
+            self._mean_cache = induc_induc_covar.solve(mean_diff).detach()
+            predictive_mean = torch.add(
+                test_mean, induc_data_covar.transpose(-2, -1).matmul(self._mean_cache).squeeze(-1)
+            )
+            predictive_covar = ZeroLinearOperator(test_mean.size(-1), test_mean.size(-1))
+            if hasattr(self.model, 'power'):
+                return MultivariateQExponential(predictive_mean, predictive_covar, power=self.model.power)
+            else:
+                return MultivariateNormal(predictive_mean, predictive_covar)
+
+        # Expand everything to the right size
+        shapes = [mean_diff.shape[:-1], induc_data_covar.shape[:-1], induc_induc_covar.shape[:-1]]
+        root_variational_covar = None
+        if variational_inducing_covar is not None:
+            root_variational_covar = variational_inducing_covar.root_decomposition().root.to_dense()
+            shapes.append(root_variational_covar.shape[:-1])
+        shape = torch.broadcast_shapes(*shapes)
+        mean_diff = mean_diff.expand(*shape, mean_diff.size(-1))
+        induc_data_covar = induc_data_covar.expand(*shape, induc_data_covar.size(-1))
+        induc_induc_covar = induc_induc_covar.expand(*shape, induc_induc_covar.size(-1))
+        if variational_inducing_covar is not None:
+            root_variational_covar = root_variational_covar.expand(*shape, root_variational_covar.size(-1))
+
+        # Cache the kernel matrix with the cached CG calls
+        if self.training:
+            if hasattr(self.model, 'power'):
+                prior_dist = MultivariateQExponential(induc_mean, induc_induc_covar, power=self.model.power)
+            else:
+                prior_dist = MultivariateNormal(induc_mean, induc_induc_covar)
+            add_to_cache(self, "prior_distribution_memo", prior_dist)
+
+        # Compute predictive mean
+        if variational_inducing_covar is None:
+            left_tensors = mean_diff
+        else:
+            left_tensors = torch.cat([mean_diff, root_variational_covar], -1)
+        inv_products = induc_induc_covar.solve(induc_data_covar, left_tensors.transpose(-1, -2))
+        predictive_mean = torch.add(test_mean, inv_products[..., 0, :])
+
+        # Compute covariance
+        if self.training:
+            interp_data_data_var, _ = induc_induc_covar.inv_quad_logdet(
+                induc_data_covar, logdet=False, reduce_inv_quad=False
+            )
+            data_covariance = DiagLinearOperator(
+                (data_data_covar.diagonal(dim1=-1, dim2=-2) - interp_data_data_var).clamp(0, math.inf)
+            )
+        else:
+            neg_induc_data_data_covar = torch.matmul(
+                induc_data_covar.transpose(-1, -2).mul(-1), induc_induc_covar.solve(induc_data_covar)
+            )
+            data_covariance = data_data_covar + neg_induc_data_data_covar
+        predictive_covar = PsdSumLinearOperator(
+            RootLinearOperator(inv_products[..., 1:, :].transpose(-1, -2)), data_covariance
+        )
+
+        # Done!
+        if hasattr(self.model, 'power'):
+            return MultivariateQExponential(predictive_mean, predictive_covar, power=self.model.power)
+        else:
+            return MultivariateNormal(predictive_mean, predictive_covar)