Batched independent 1D variational GPs as NN activation

[schneider-source]

Extending from the batch regression example, I would like to implement a 1D -> 1D variational GP activation function for later use in neural networks. So for each node in a neural network layer an independent variational GP activation can be learned using VI. How to correctly implement that, as below minimal (executable) example does not work as expected. Thanks for help!

import torch
import gpytorch


class Batched1DApproximateGP(gpytorch.models.deep_gps.DeepGPLayer):
    def __init__(self, B, L, Z, I):
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            num_inducing_points=Z,
            batch_shape=torch.Size([B, L])
        )
        variational_strategy = gpytorch.variational.IndependentMultitaskVariationalStrategy(
            gpytorch.variational.VariationalStrategy(
                self,
                inducing_points=torch.randn(L, Z, I),
                variational_distribution=variational_distribution,
                learn_inducing_locations=True
            ),
            num_tasks=L,
            task_dim=-1
        )
        super().__init__(variational_strategy, input_dims=I, output_dims=I)
        self.mean_module = gpytorch.means.ConstantMean(batch_shape=torch.Size([B, L]))
        self.covar_module = gpytorch.kernels.ScaleKernel(
            base_kernel=gpytorch.kernels.RBFKernel(batch_shape=torch.Size([B, L])),
            batch_shape=torch.Size([B, L])
        )

    def forward(self, x: torch.Tensor):
        # shape of x here: (B, L, L + Z, 1) ???
        mean_x = self.mean_module(x)  # (B, L, L + Z, I)
        covar_x = self.covar_module(x)  # (B, L, L + Z, L + Z)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)


print(gpytorch.__version__)

B = 19  # batch size
L = 5  # layer size
Z = 3  # num inducing points
I = 1  # input/output size

x = torch.randn((B, L, I))
model = Batched1DApproximateGP(B, L, Z, I)
y = model(x)
print(f'obtained shape: {y.mean.shape}')
print(f'expected shape: {torch.Size([B, L])}')

Output:

1.11
obtained shape: torch.Size([10, 25, 19])
expected shape: torch.Size([19, 5])

[gpleiss]

Don't use a deep GP - just use an ApproximateGP. Your model should look something like what we have in the SVI regression example: https://docs.gpytorch.ai/en/stable/examples/04_Variational_and_Approximate_GPs/SVGP_Regression_CUDA.html

If you want to learn multiple independent activation functions, see this example: https://docs.gpytorch.ai/en/stable/examples/04_Variational_and_Approximate_GPs/SVGP_Multitask_GP_Regression.html