PseudoObs in Stheno

[patel-zeel]

The PseudoObs implementation in Stheno is inspired by any paper or is it a natural way of doing sparse approximation? In any case, can you please guide me to the theoretical resource/reference behind this implementation?

[wesselb]

Hey @patel-zeel! PseudoObs does nothing too fancy: it uses the pseudo-point approximation by Titsias (2009). In a little more technical detail, it implements the optimal variational approximation over the specified inducing points. The particular implementation of PseudoObs, however, is not directly taken from the paper; rather, it is a custom formulation of Titsias' solution.

[wesselb]

To gain more clarity, does your answer mean that If I implement Titsias 2009 from scratch, my results should match with PseodoObs version (ignoring numerical glitches)?

Definitely! If you were to implement the formulas from the paper, the results should exactly line up with the results from PseudoObs. (Just out of curiosity, have you used PseudoObs? If so, how has it been working for you?)

Bonus question: (...)

Generally, the variational approximation by Titsias is regarded as the state-of-the-art approximation (the one best approximating the true posterior) if no other properties of the kernel can be exploited. One fundamental reason which suggests why this may be is that Titsias' pseudo-point approximation computes the optimal variational posterior in the correct model, whereas other approximations, like the one by Snelson & Ghahramani (also called FITC), instead compute an exact posterior in a modified model.

Would you be interested in also trying other approximations, like FITC and DTC? (Both those should be fairly easy to add to Stheno with a similar interface as PseudoObs.)

[patel-zeel]

Thanks for the detailed and reasonable answer @wesselb. I am just starting with PseudoObs application so haven't tried yet for my own data. I am planning to implement FITC in Stheno soon to learn it by doing it :) Shall I open a PR if I am able to do it?

[wesselb]

@patel-zeel A PR would be fantastic! The way I was thinking to go about it is to make a PseudoObsFITC class (or whatever its name should be) which mimics the structure of PseudoObs.

[patel-zeel]

@wesselb Sure, I will try to do it.

[patel-zeel]

Hi @wesselb,

I have attempted FITC implementation in the same fashion as the VFE. I created this notebook as a kind of initial testing. Also, I am attaching the main class below wondering if you can give feedback on unnecessary or wrong computations at some places.

I am thinking of renaming each as PseudoObsVFE and PseudoObsFITC, what do you think?

class PseudoObservationsFITC(AbstractObservations):
    """Observations through inducing points (FITC).

    Further takes arguments according to the constructor of
    :class:`.measure.AbstractObservations`. Can also take in tuples of inducing points.

    Args:
        u (:class:`.fdd.FDD`): Inducing points
    """

    @_dispatch
    def __init__(self, u: FDD, *args):
        AbstractObservations.__init__(self, *args)
        self.u = u
        self._K_z_store = {}
        self._elbo_store = {}
        self._mu_store = {}
        self._Q_store = {}

    @_dispatch
    def __init__(self, us: tuple, *args):
        PseudoObservationsFITC.__init__(self, combine(*us), *args)

    def K_z(self, measure):
        """Kernel matrix of the data.

        Args:
            measure (:class:`.measure.Measure`): Measure.

        Returns:
            matrix: Kernel matrix.
        """
        try:
            return self._K_z_store[id(measure)]
        except KeyError:
            self._compute(measure)
            return self._K_z_store[id(measure)]

    def elbo(self, measure):
        """ELBO.

        Args:
            measure (:class:`.measure.Measure`): Measure.

        Returns:
            scalar: ELBO.
        """
        try:
            return self._elbo_store[id(measure)]
        except KeyError:
            self._compute(measure)
            return self._elbo_store[id(measure)]

    def mu(self, measure):
        """Mean of optimal approximating distribution.

        Args:
            measure (:class:`.measure.Measure`): Measure.

        Returns:
            matrix: Mean.
        """
        try:
            return self._mu_store[id(measure)]
        except KeyError:
            self._compute(measure)
            return self._mu_store[id(measure)]

    def Q(self, measure):
        """Parameter of the corrective variance of the kernel of the optimal
        approximating distribution.

        Args:
            measure (:class:`.measure.Measure`): Measure.

        Returns:
            matrix: Corrective variance.
        """
        try:
            return self._Q_store[id(measure)]
        except KeyError:
            self._compute(measure)
            return self._Q_store[id(measure)]

    def _compute(self, measure):
        # Extract processes and inputs.
        p_x, x, noise_x = self.fdd.p, self.fdd.x, self.fdd.noise
        p_z, z, noise_z = self.u.p, self.u.x, self.u.noise

        # Construct the necessary kernel matrices.
        K_zx = measure.kernels[p_z, p_x](z, x)
        K_z = B.add(measure.kernels[p_z](z), noise_z)
        self._K_z_store[id(measure)] = K_z

        # Noise kernel matrix:
        K_n = noise_x 

        # The approximation can only handle diagonal noise matrices.
        if not isinstance(K_n, Diagonal):
            raise RuntimeError(
                f"Kernel matrix of observation noise must be diagonal, "
                f'not "{type(K_n).__name__}".'
            )

        # And construct the components for the inducing point approximation.
        diag_part = K_n + Diagonal(measure.kernels[p_x].elwise(x)[:, 0]) - Diagonal(B.iqf_diag(K_z, K_zx))
        L_z = B.cholesky(K_z)        
        A = B.add(B.eye(K_z), B.iqf(diag_part, B.transpose(B.solve(L_z, K_zx))))
        Q = K_z + B.iqf(diag_part, B.transpose(K_zx))
        self._Q_store[id(measure)] = Q
        
        y_bar = B.subtract(B.uprank(self.y), measure.means[p_x](x))
        prod_y_bar = B.solve(L_z, B.iqf(diag_part, B.transpose(K_zx), y_bar))

        # Compute the optimal mean.
        mu = B.add(
            measure.means[p_z](z),
            B.iqf(diag_part, B.transpose(K_zx), y_bar),
        )
        self._mu_store[id(measure)] = mu

        # Compute the ELBO.

        dtype = B.dtype_float(K_n)
        det_part = B.logdet(B.multiply(B.cast(dtype, 2 * B.pi), K_n)) + B.logdet(A)
        iqf_part = B.iqf(K_n, y_bar)[0, 0] - B.iqf(A, prod_y_bar)[0, 0]
        self._elbo_store[id(measure)] = -0.5 * (det_part + iqf_part)

    def posterior_kernel(self, measure, p_i, p_j):
        return PosteriorKernel(
            measure.kernels[p_i, p_j],
            measure.kernels[self.u.p, p_i],
            measure.kernels[self.u.p, p_j],
            self.u.x,
            self.K_z(measure),
        ) + SubspaceKernel(
            measure.kernels[self.u.p, p_i],
            measure.kernels[self.u.p, p_j],
            self.u.x,
            self.Q(measure),
        )

    def posterior_mean(self, measure, p):
        return PosteriorMean(
            measure.means[p],
            measure.means[self.u.p],
            measure.kernels[self.u.p, p],
            self.u.x,
            self.Q(measure),
            self.mu(measure),
        )