In-Context In-Context Learning with Transformer Neural Processes

TL;DR

This paper introduces ICICL-TNP, a novel transformer-based neural process model that enables in-context learning from multiple datasets, improving meta-learning capabilities by conditioning on both data points and datasets.

Contribution

The paper develops ICICL-TNP, the first neural process model capable of in-context in-context learning by conditioning on multiple datasets, extending the functionality of existing NPs.

Findings

ICICL-TNP effectively performs in-context in-context learning.

It improves prediction accuracy by leveraging multiple datasets.

The model scales efficiently with pseudo-token transformer architecture.

Abstract

Neural processes (NPs) are a powerful family of meta-learning models that seek to approximate the posterior predictive map of the ground-truth stochastic process from which each dataset in a meta-dataset is sampled. There are many cases in which practitioners, besides having access to the dataset of interest, may also have access to other datasets that share similarities with it. In this case, integrating these datasets into the NP can improve predictions. We equip NPs with this functionality and describe this paradigm as in-context in-context learning. Standard NP architectures, such as the convolutional conditional NP (ConvCNP) or the family of transformer neural processes (TNPs), are not capable of in-context in-context learning, as they are only able to condition on a single dataset. We address this shortcoming by developing the in-context in-context learning pseudo-token TNP (ICICL-TNP). The ICICL-TNP builds on the family of PT-TNPs, which utilise pseudo-token-based transformer architectures to sidestep the quadratic computational complexity associated with regular transformer architectures. Importantly, the ICICL-TNP is capable of conditioning on both sets of datapoints and sets of datasets, enabling it to perform in-context in-context learning. We demonstrate the importance of in-context in-context learning and the effectiveness of the ICICL-TNP in a number of experiments.