Cambridge, Google & Secondmind’s Neural Diffusion Processes Challenge Gaussian Processes for Describing Rich Distributions Over Functions

While researchers have traditionally employed Gaussian processes (GP) for specifying prior and posterior distributions over functions, this approach becomes computationally expensive when scaled, is limited by the expressivity of its covariance function, and struggles with adapting a point estimation for the hyperparameters.

A research team from the University of Cambridge, Secondmind, and Google Research addresses these issues in the new paper Neural Diffusion Processes, proposing Neural Diffusion Processes (NDPs). The novel framework learns to sample from rich distributions over functions at a lower computational cost and capture distributions that are close to the true Bayesian posterior of a conventional Gaussian process.

The paper’s lead author, Vincent Dutordoir, explains, “Bayesian inference for regression is great, but it is often very costly and requires making a priori modelling assumptions. What if we can train a big neural net to sample plausible posterior samples over functions? This is the premise of our Neural Diffusion Processes.”

The team summarizes their main contributions as:

1. We propose a novel model, the Neural Diffusion Process (NDP), which extends the use case of diffusion models to stochastic processes and is able to describe a rich distribution over functions.
2. We take particular care to enforce known symmetries and properties of stochastic processes, including exchangeability, into the model, facilitating the training process.
3. We showcase the abilities and versatility of NDPs by applying them to a range of Bayesian inference tasks including prior and conditional sampling, regression, hyperparameter marginalization, and Bayesian optimization.
4. We also present a novel global optimization method using NDPs.

The proposed NDP is a denoising diffusion model-based approach for learning probabilities from a function and producing prior and conditional samples of functions. It allows full marginalization over the GP hyperparameters while reducing the computational burden compared to GPs.

The team first examined existing state-of-the-art neural network-based generative models in terms of sample quality. Based on their findings, they designed NDP to generalize diffusion models to infinite-dimensional function spaces by enabling the indexing of random variables onto which the model diffuses.

The researchers also adopted a novel bi-dimensional attention block to guarantee equivariance over the input dimensionality and sequence and enable the model to draw samples from a stochastic process. As such, NDP can leverage the benefits of stochastic processes, such as exchangeability.

In their empirical study, the team evaluated the proposed NDP’s ability to produce high-quality conditional samples and marginalize over kernel hyperparameters; and on its input dimensionality invariance.

The results show that NDP is able to capture functional distributions that are close to the true Bayesian posterior while reducing computational burdens.

The researchers note that while NDP sample quality improves with the number of diffusion steps, this also results in slower inference times. They suggest inference acceleration or sample parameterizing techniques could be explored in future studies to address this issue.

The paper Neural Diffusion Processes is on arXiv.

---

Author: Hecate He | Editor: Michael Sarazen, Chain Zhang

---

We know you don’t want to miss any news or research breakthroughs. Subscribe to our popular newsletter Synced Global AI Weekly to get weekly AI updates.