Exploiting Inferential Structure in Neural Processes

TL;DR

This paper introduces a framework for Neural Processes that incorporates rich prior distributions via graphical models, enhancing their ability to model complex data distributions and improving robustness.

Contribution

It proposes a novel method to integrate complex priors into Neural Processes using graphical models and message passing, enabling better modeling of diverse data distributions.

Findings

Improved function modeling with mixture and Student-t priors.

Enhanced test-time robustness in neural process models.

Framework generality demonstrated across different prior assumptions.

Abstract

Neural Processes (NPs) are appealing due to their ability to perform fast adaptation based on a context set. This set is encoded by a latent variable, which is often assumed to follow a simple distribution. However, in real-word settings, the context set may be drawn from richer distributions having multiple modes, heavy tails, etc. In this work, we provide a framework that allows NPs' latent variable to be given a rich prior defined by a graphical model. These distributional assumptions directly translate into an appropriate aggregation strategy for the context set. Moreover, we describe a message-passing procedure that still allows for end-to-end optimization with stochastic gradients. We demonstrate the generality of our framework by using mixture and Student-t assumptions that yield improvements in function modelling and test-time robustness.