Bayesian Context Aggregation for Neural Processes

Formulating scalable probabilistic regression models with reliable uncertainty estimates has been a long-standing challenge in machine learning research. Recently, casting probabilistic regression as a multi-task learning problem in terms of conditional latent variable (CLV) models such as the Neural Process (NP) has shown promising results. In this paper, we focus on context aggregation, a central component of such architectures, which fuses information from multiple context data points. So far, this aggregation operation has been treated separately from the inference of a latent representation of the target function in CLV models. Our key contribution is to combine these steps into one holistic mechanism by phrasing context aggregation as a Bayesian inference problem. The resulting Bayesian Aggregation (BA) mechanism allows principled handling of task ambiguity, which is key for efficiently processing context information. We demonstrate on a range of challenging experiments, including dynamics modeling and image completion, that BA consistently improves upon the performance of traditional aggregation methods while remaining computationally efficient and fully compatible with existing NP-based models.