Latent Bottlenecked Attentive Neural Processes

TL;DR

The paper introduces LBANPs, a scalable, efficient neural process model that maintains competitive performance while reducing computational complexity from quadratic to sub-quadratic, enabling larger dataset handling.

Contribution

LBANPs is a novel sub-quadratic neural process variant that encodes context into latent vectors, improving scalability and efficiency without sacrificing accuracy.

Findings

Achieves competitive results on meta-regression, image completion, and bandits.

Scales beyond existing attention-based NPs to larger datasets.

Allows trading off computational cost and performance.

Abstract

Neural Processes (NPs) are popular methods in meta-learning that can estimate predictive uncertainty on target datapoints by conditioning on a context dataset. Previous state-of-the-art method Transformer Neural Processes (TNPs) achieve strong performance but require quadratic computation with respect to the number of context datapoints, significantly limiting its scalability. Conversely, existing sub-quadratic NP variants perform significantly worse than that of TNPs. Tackling this issue, we propose Latent Bottlenecked Attentive Neural Processes (LBANPs), a new computationally efficient sub-quadratic NP variant, that has a querying computational complexity independent of the number of context datapoints. The model encodes the context dataset into a constant number of latent vectors on which self-attention is performed. When making predictions, the model retrieves higher-order information from the context dataset via multiple cross-attention mechanisms on the latent vectors. We empirically show that LBANPs achieve results competitive with the state-of-the-art on meta-regression, image completion, and contextual multi-armed bandits. We demonstrate that LBANPs can trade-off the computational cost and performance according to the number of latent vectors. Finally, we show LBANPs can scale beyond existing attention-based NP variants to larger dataset settings.