Gaussian Processes in GPyTorch for Multidimensional Continuous Signal Reconstruction and Image Inpainting

Building models that know what they don’t know, by design.

Image source: Danu Caus

Introduction and Motivation

Today we find ourselves in the thriving era of artificial intelligence (AI), where deep learning continues to drive groundbreaking innovations. It’s the perfect moment to revisit Gaussian Processes (GPs) and explore how these classical models can be revitalized and enhanced by modern deep learning tools.

Gaussian Processes represent a universal framework with fundamental connections to major machine learning (ML) disciplines, including Bayesian statistics and deep learning. If you’re drawn to understanding the bigger picture of things, Gaussian Processes provide a powerful lens for exploring connections across multiple domains.

Figure 1: Gaussian Processes positioned at the intersection of Bayesian Modelling and Deep Learning. GPs leverage Bayesian principles for rigorous uncertainty quantification and can integrate with modern deep learning tools for scalability and efficient inference. [Image source: Author]

Why would you consider Gaussian Processes in the first place? To begin with, GPs are especially valuable when you want to reason about uncertainty in a principled manner. While neural networks excel in many practical tasks, they often remain somewhat of a “black box,” especially regarding uncertainty quantification. Gaussian Processes, by contrast, offer a mathematically rigorous framework for explicitly quantifying uncertainty.

Moreover, neural networks typically require large amounts of data to perform effectively, whereas Gaussian Processes are particularly suited to scenarios with smaller datasets. When data is limited, GPs can better capture and represent uncertainty, allowing you to reason about constraints and predictions with greater confidence. Even when large datasets are involved, efficient approximation methods exist for scaling Gaussian Processes, such as Variational Gaussian Processes with inducing points, an approach we will explore later in this article.

Intuition and Interpretation

One intuitive way to understand Gaussian Processes is to view them as an ensemble of many possible models — or more formally, a probability distribution over functions. To illustrate this idea, consider Figure 2 below. Here, the Gaussian Process is represented as a bundle of functions, each offering a legitimate explanation of the dataset. These functions are constrained to pass through (or at least near, if measurements are noisy) the observed data points, represented by the X-marks. In regions without concrete measurements, the functions are less constrained and can vary more freely, causing the bundle to spread out. This increased variability creates the shaded region, which intuitively captures the uncertainty quantified by the Gaussian Process.

Figure 2: Gaussian Processes interpreted as an ensemble of functions. Each function describes the same data in its own way. These functions vary within a confidence interval and by averaging them out, we get a mean / average model function. [Image source: Author]

To produce a similar plot as the one above, please follow -> https://github.com/ducspe/SpatioTemporalGaussianProcessInterpolation/blob/main/notebooks/GaussianProcess_demo.ipynb [1]

A second way to view GPs is through the lens of basis functions, i.e. so called kernels. This view is similar conceptually to Fourier series, where we use sines and cosines to represent more complex functions, or Taylor series, where we use polynomials. In this case we can say that Gaussian Processes are a combination of kernel basis functions, and we can use our domain knowledge and insights into the underlying task to pick the most suitable kernels. Examples of such kernels are the radial basis function / RBF, the Matern kernel, Spectral Mixture kernels and many more. Figure 3 below illustrates the general concept of centering a specific kernel to the available data points. If we gather multiple such kernels and combine them, we are able to represent a more complex model mean function, along with its corresponding uncertainty. When we train a GP model, our task is essentially to find the parameters of these kernels that allow us to better fit our model to the data (see the Mathematical Formalities, Notation and Tools section below for more insight into how these parameters look like for the RBF kernel).

Figure 3: Gaussian Processes viewed as a combination of kernel functions centered at the data measurement points. Combining multiple such simple kernels together allows us to create a complex model and judge the uncertainty at other locations where data is not available. [Image source: Author]

To create a similar plot as the one above, you can run this notebook -> https://github.com/ducspe/SpatioTemporalGaussianProcessInterpolation/blob/main/notebooks/GaussianProcess_kernel_perspective.ipynb [1]

Yet a third way to view a GP is as a neural network that consists of a single layer with an infinite number of neurons, as shown in Figure 4 below. Of course, using our discrete modern tools like PyTorch or TensorFlow we would not be able to create such a neural network in practice, but the advantage of GPs is that they offer an analytical way to build the equivalent of such infinitely wide neural nets.

Figure 4: A Gaussian Process viewed as a single-layer neural network that is infinitely wide, i.e. unlimited neurons. [Image source: Author]

For more related information building on this particular interpretation, you may check NNGPs at -> https://en.wikipedia.org/wiki/Neural_network_Gaussian_process [2]

Mathematical Formalities, Notation and Tools

For more mathematical formality and clarity regarding notation, I found the lectures and book [3] of Carl Edward Rasmussen to be tremendously helpful! So I recommend them to get more in depth understanding of the subject matter. For this brief article however, I wrote a couple of short summarizing notes in LaTeX and snipped the generated image below:

Gaussian Processes summarizing notes. [Image source: Author, using LaTeX]

GPyTorch

To design scalable systems using Gaussian Processes, we naturally require efficient computational tools. While there have been numerous GP frameworks available in the past, the rise of deep learning calls for a modern solution capable of leveraging Graphics Processing Units (GPUs) effectively. A prominent example is GPyTorch [4], a framework specifically built to exploit highly optimized matrix-to-matrix operations for efficient GP training and inference [5].

GPyTorch is particularly attractive to those familiar with PyTorch, as it maintains a similar syntax and coding style, easing the transition to Gaussian Process-based systems. Moreover, its close integration with PyTorch allows users to conveniently design and implement hybrid architectures, seamlessly combining Gaussian Processes and neural networks within a unified system pipeline.

Two-Dimensional Interpolation Task

To illustrate GPyTorch in action we will use an image inpainting task / 2D continuous signal interpolation task. Concretely, suppose we have a geographic region with limited locations where sensors take measurements of certain variables, say temperature for example, see Figure 5. The task is then to learn a model that would fill in the blanks in those locations where data is not available, using the sampled measurements in other parts of the 2D map, while estimating uncertainty in parallel given the constraints.

Figure 5: Conceptually showcasing a 2D interpolation task. We assume a continuous signal with only sparse measurements available. In this case the signal is mapped on a two-dimensional grid, and can also vary in time. We can take temperature as a concrete example of such a signal, but the Gaussian Processes approach can be applied to reconstructing real world images as well, since they do generally fit the requirement of being a 2D continuous signal. [Image source: Author]

To solve this task we will use two approaches: the Exact GP approach and the Approximated GP approach via so called inducing points.

The Exact GP will take in all the training data points and build a matrix out of them in order to capture the detailed variations within the training data. We will use the gpytorch.models.ExactGP as presented here -> https://github.com/ducspe/SpatioTemporalGaussianProcessInterpolation [1] specifically in the source code from the exact_gp_related folder.

Note how in the train script we can have 2 separate kernels for the time and space dimensions and simply add them up as shown below:

# Combine temporal and spatial kernels
if self.temporal_kernel is not None:
    self.covar_module = self.spatial_kernel + self.temporal_kernel
else:
    self.covar_module = self.spatial_kernel

Since we format our data to be a list of tuples of the form:

[Image source: Author, using LaTeX]

where Time is at index 0, X is at index 1 and Y is at index 2, we use the active_dims argument of the kernel objects to specify in which dimension we want to learn the kernel hyperparameters, as follows:

# Temporal kernel
if time_kernel == "rbf":
    print("Using RBF kernel for temporal dimension.")
    self.temporal_kernel = gpytorch.kernels.ScaleKernel(
        gpytorch.kernels.RBFKernel(active_dims=[0])
        # active_dims=[0] means kernel acts only on time dimension
    )
elif time_kernel == "matern":
    # ...

And similarly, for the X and Y directions of the spatial kernel, we have:

# Spatial kernel
if space_kernel == "rbf":
    print("Using RBF kernel for spatial dimension.")
    self.spatial_kernel = gpytorch.kernels.ScaleKernel(
        gpytorch.kernels.RBFKernel(active_dims=[1, 2])
        # active_dims=[1,2] indicates X and Y spatial coordinates
    )
elif space_kernel == "matern":
    # ...

Note that it is also possible to create a linear combination of scaled kernels for each dimension, space and time, of the form:

[Image source: Author, using LaTeX]

This is essentially a mixture of scaled kernels, where each kernel can be of different types, i.e. RBF, Matern, Periodic, etc. For example if you observe that your signal contains an upwards trend with local periodic oscillations, it could be wise to design a compound kernel that consists of 2 simpler kernels, e.g. an RBF and a periodic kernel added together.

When the training dataset becomes too large to fit entirely into GPU memory, we can utilize approximations provided by the gpytorch.models.ApproximateGP object. A concrete implementation of this approach is available in the variational_gp_related folder of our repository: **https://github.com/ducspe/SpatioTemporalGaussianProcessInterpolation [1]**.

Specifically, this approximation method employs so-called inducing points, which act as a concise summary of the training data. Instead of using the entire training dataset directly in each inference step, we optimize these inducing points to efficiently represent the overall data distribution. Thus, rather than storing all training points simultaneously in memory, we rely on a much smaller set of inducing points, significantly reducing GPU memory requirements.

However, this scalability comes with a tradeoff: since the inducing points are a compressed representation, some loss of information is possible. This compression typically results in smoother predictions compared to the Exact Gaussian Process method, which directly utilizes the entire dataset.

Besides using a limited number of inducing points to summarize the dataset, another effective strategy to reduce memory usage with Approximate Gaussian Processes is the mini-batching procedure. Mini-batching, illustrated in the code block below, allows us to process large datasets efficiently in smaller chunks, making it possible to handle datasets far larger than would be feasible with an Exact Gaussian Process.

# Mini-batching approach to avoid GPU memory issues on large datasets
for iter_idx in range(MAX_ITERS):
    # Set model and likelihood into training mode
    model.train()
    likelihood.train()
    total_train_loss = 0.0

    # Loop over mini-batches from DataLoader
    for x_batch, y_batch in train_loader:
        optimizer.zero_grad()                   # Reset gradients before backward pass
        output = model(x_batch)                 # Forward pass to get GP predictions
        loss = -mll(output, y_batch)            # Compute the negative marginal log-likelihood loss
        loss.backward()                         # Compute gradients via backpropagation
        optimizer.step()                        # Update parameters using gradients
        total_train_loss += loss.item() * x_batch.size(0)  # Accumulate loss (scaled by batch size)

    # Compute average loss per training data point for current iteration
    total_train_loss /= len(train_dataset)

Note that in our repository source code we use a combination of these tricks, i.e. a mini-batched version of the Inducing Points Variational GP.

Qualitative Results

Below are some qualitative examples of a test frame reconstruction using various kernels. We can observe how the Matern kernel exhibits slightly lower errors and lower uncertainties than the RBF Kernel, and how the periodic kernel results in relatively worse reconstructions, hinting that it’s not the most adequate kernel selection for this particular data subset.

Figure 6: Example of a test frame reconstruction using the RBF kernel model. GT frame is the ground truth. The second from left is the Mean frame, i.e. the reconstructed map after 90% of the pixels were removed from the ground truth. “Reconstructed Var” map reflects the certainty of the model, and “Error Map” is the difference between the GT and predicted maps. [Image source: Author]

Figure 7: Same test frame reconstruction, but using the Matern kernel model this time. GT frame is the ground truth. The second from left is the Mean frame, i.e. the reconstructed map after 90% of the pixels were removed from the ground truth. “Reconstructed Var” map reflects the certainty of the model, and “Error Map” is the difference between the GT and predicted maps. [Image source: Author]

Figure 8: Example of what happens when we use a slightly unsuitable kernel for our specific dataset / task. In this case the mean and variance maps have checkerboard artifacts and even if the model prediction certainty is high, the model makes larger errors as compared to other models that were built with more appropriate kernels for this particular task, namely RBF and Matern. [Image source: Author]

Quantitative Results

To show also some quantitative results we run several experiments with different kernels, i.e. RBF, Matern and Periodic; as well as different masking probabilities: 50% grid values dropped off, then 80%, and finally 90% masked values as more of an extreme scenario.

Figure 9: Quantitative results for several experiments, with different kernels and data sparsity. MAE is the Mean Absolute Error, and RMSE is the Root Mean Squared Error. Note that prob=0.5 means 50% of values were masked, and 0.8/0.9 means 80% and 90% respectively were masked, i.e. for the extrem cases only 10% of the points were available to reconstruct from. Note that Gaussian Processes offer similar results across masking degrees, taking advantage of the signal continuity as very important prior knowledge, i.e. a priori assumption. [Image source: Author]

We can observe in the plot above that if we use appropriate kernels for our task, the error between the reconstructed output and the ground truth is rather small, whereas if we use a kernel not very adequate for our particular task or dataset, in this case the periodic kernel, the errors get naturally larger. Also note that because the signal is continuous, we can reconstruct it very well even if we mask 90% of the original frame.

Of course, GPs are not perfect; they can still be overconfident and make mistakes. Nevertheless, especially when constrained by lack of data, they are one of the best tools available to quantify uncertainty in a more principled manner than neural networks.

And just to emphasize an already mentioned note for this section, you may still wonder given the qualitative and quantitative results we’ve seen, how is it possible to interpolate rather well, even when we mask 90% of the grid values, i.e. using only 10% of pixels and filling in the rest. In fact, in these extreme cases we take advantage of the idea that our signal to predict is continuous in nature. In other words, knowing beforehand that our signal is continuous constitutes a very strong a priori knowledge in the Bayesian sense, i.e. a very important assumption that we leverage in such cases.

Conclusions

In conclusion, Gaussian Processes provide a valuable approach for predicting average trends and quantifying uncertainties. We’ve explored how to leverage GPyTorch to reconstruct multi-dimensional data, specifically 2D continuous signals evolving over time. With modern deep learning libraries like GPyTorch, we can achieve faster and more efficient GPU training and significantly scale Gaussian Processes to handle larger and higher dimensional datasets than previously possible. Although Gaussian Processes may sometimes not match neural networks in raw predictive power, they offer the crucial advantage of explicitly choosing kernels or combinations of kernels guided by domain expertise and intuition. This flexibility enables more principled modeling of the underlying processes relevant to specific tasks.

References

[1] This Medium article is accompanied by a GitHub repository -> **https://github.com/ducspe/SpatioTemporalGaussianProcessInterpolation**

[2] Neural Nets vs Gaussian Processes -> **https://en.wikipedia.org/wiki/Neural_network_Gaussian_process**

[3] Carl Edward Rasmussen’s book -> **https://gaussianprocess.org/gpml/**

[4] GPyTorch framework documentation -> **https://gpytorch.ai**

[5] Paper by Gardner et al. on GPyTorch conceptual background and internals -> GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration