Gpytorch: [Bug] Possible bug in variational strategy

[ptonner]

🐛 Bug

I'm reviewing the Variational Strategy code and there may be a
bug in how the predictive mean is calculated. Alternatively
(and more likely) I am not understanding the logic.

Conern starts here:

    L = self._cholesky_factor(induc_induc_covar)
    if L.shape != induc_induc_covar.shape:
        # Aggressive caching can cause nasty shape incompatibilies when evaluating with different batch shapes
        # TODO: Use a hook fo this
        pop_from_cache(self, "cholesky_factor")
        L = self._cholesky_factor(induc_induc_covar)
    interp_term = L.inv_matmul(induc_data_covar.double()).to(full_inputs.dtype)

    # Compute the mean of q(f)
    # k_XZ K_ZZ^{-1/2} (m - K_ZZ^{-1/2} \mu_Z) + \mu_X
    predictive_mean = (
        torch.matmul(
            interp_term.transpose(-1, -2), (inducing_values - self.prior_distribution.mean).unsqueeze(-1)
        ).squeeze(-1)
        + test_mean
    )

Trying to maintain notation with the documentation provided,
I believe we want to calculate: K_{XZ}K_{ZZ}^{-1}(u - \mu_u) =
K_{XZ}L_{ZZ}^{-T}L_{ZZ}^{-1}(u - \mu_u)

From my understanding, the interp_term is giving us
K_{XZ}L_{ZZ}^{-T} and I think we're missing a second
L_{ZZ}^{-1} term.

Apologies if I'm just missing something or got my equations
wrong!

[jacobrgardner]

Hi @ptonner --

Thanks for the careful look! However, this is actually not a bug, as VariationalStrategy implements the whitened parameterization of variational inference (although the docstrings of VS could be more explicit about this). In this parameterization, rather than learning m and S, we learn whitened parameters m' = Kuu^{-1/2}m and S' = Kuu^{-1/2}SKuu^{-1/2}. Doing this makes optimization performance significantly better as it helps decouple the GP hyperparameters from the variational parameters. Furthermore, the KL divergence term is dramatically simplified under this parameterization. Plugging these parameterizations in to the standard variational inference equations you are familiar with results in the equations above.

[ptonner]

great thanks, figured I was missing something!

[ZhiliangWu]

Hi @jacobrgardner, thanks for the great package and documentation. I am also reading this forward() method of the variational_strategy. I think the formula you are following is the equation 18 in Hensman et al. 2015. I am a bit confused about the computation of the mean, where you also involved ...\mu_Z) + \mu_X. Could you explain the reason behind these terms or offer some pointers of the papers to clarify? Thanks for your reply in advance.

[jacobrgardner]

VariationalStrategy does not follow equation 18 from that paper. UnwhitenedVariationalStrategy does.

VariationalStrategy applies the vartiaional parameters in a transformed space that is generally easier to work with (a technique commonly known as whitening). See, for example:

M. Kuss and C. E. Rasmussen. Assessing approximate inference for binary Gaussian process classification.
Journal of Machine Learning Research, 6(Oct):1679–1704, 2005.

A. G. d. G. Matthews. Scalable Gaussian process inference using variational methods. PhD thesis,
University of Cambridge, 2017.

The computation involving \mu_Z and \mu_X account for the fact that equation 18 directly assumes a 0 mean Gaussian process prior, while in code you may have any sort of prior mean function.

[ZhiliangWu]

Hi @jacobrgardner, thanks for the instant reply and references. I think you are right. But I am also a bit confused by the last point referring to zero mean prior. As in variational inference, we assume the q(u) comes from a normal distribution with (trainable) mean m and variance S and get q(f) by integrating u out. I didn't see any other possible zero mean prior. Did I miss anything?

[jacobrgardner]

The variational mean is the (variational) posterior mean that we learn for u. There is still a prior mean for u (and f).