Efficient Queries Transformer Neural Processes
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 per query, limiting its applications.
Conversely, existing sub-quadratic NP variants perform significantly worse than that of TNPs.
Tackling this issue, we propose Efficient Queries Transformer Neural Processes (EQTNPs), a more computationally efficient NP variant.
The model encodes the context dataset into a set of vectors that is linear in the number of context datapoints.
When making predictions, the model retrieves higher-order information from the context dataset via multiple cross-attention mechanisms on the context vectors.
We empirically show that EQTNPs achieve results competitive with the state-of-the-art.
Introduction
Meta-learning aims to learn a model that can adapt quickly and computationally efficiently to new tasks.
Neural Processes (NPs) are a popular method in meta-learning that models a conditional distribution of the prediction of a target datapoint given a set of labelled (context) datapoints, providing uncertainty estimates.
NP variants (Garnelo et al., 2018a; Gordon et al., 2019; Kim et al., 2019) adapt via a conditioning step in which they compute embeddings representative of the context dataset.
NPs can be divided into two categories: (1) computationally efficient (sub-quadratic complexity) but poor performance and (2) computationally expensive (quadratic complexity) but good performance.
Early NP variants were especially computationally efficient, requiring only linear computation in the number of context datapoints but suffered from underfitting and as a result, overall poor performance.
Tackling underfitting, recent state-of-the-art methods such as TNPs have proposed to use self-attention mechanisms.
However, these state-of-the-art methods are computationally expensive in that they require quadratic computation in the number of context datapoints per query.
The quadratic computation makes the method inapplicable in settings with limited resources.
In this work, we introduce Efficient Queries Transformer Neural Processes (EQTNPs), a NP variant which has a per target datapoint query complexity that is linear in the number of context datapoints.
Our experiments show that EQTNPs achieve results competitive with the state-of-the-art methods.
Background 2.1 Meta-Learning for Predictive Uncertainty Estimation
In meta-learning, models are trained on a distribution of tasks Ω( T ) .
During each meta-training iteration, a batch of B tasks T = {T i } Bi
=1 is sampled from a task distribution Ω( T ) .
A task T i is a tuple ( X , Y , L , q ), where X is the input space, Y is the output space, L is the task-specific
[Table]
loss function, and q ( x, y ) is a distribution over data points.
During each meta-training iteration, for each T i ∈ T , we sample from q T i : a context dataset D context i = { ( x, y ) i,j } N j =1 and a target dataset D target i = { ( x, y ) i,j } Mj =1 , where N and M are the fixed number of context and target datapoints respectively. The context data is used to perform an update to the model f such that the type of update differs depending on the model. In Neural Processes, these updates refer to its conditioning step where embeddings of the dataset are computed.
Afterwards, the update is evaluated on the target data and the update rule is adjusted. In this setting, we use a neural network to model the probabilistic predictive distribution p θ ( y | x, D context ) where θ are the parameters of the NN. 2.2 Neural Processes Neural Processes (NPs) are a class of models that define an infinite family of conditional distributions where one can condition on arbitrary number of context datapoints (labelled datapoints) and make predictions for an arbitrary number of target datapoints (unlabelled datapoints), while preserving invariance in the ordering of the contexts.
Several NP models have proposed to model it as follows:
[Equation]
with r C :
=
Agg( D context ) where Agg is a deterministic function varying across the variants that aggeregates D context into a finite represention with permutation invariance in the context dataset.
These NPs aim to maximise to likelihood of the target dataset given the context dataset.
Conditional
Neural Processes (CNPs) (Garnelo et al., 2018a), proposes an aggregator using DEEPSETs (Zaheer et al., 2017) and Conditional Attentive Neural Processes (CANPs) (Kim et al., 2019)
(a deterministic variant of ANPs without the latent path) proposes an aggreagator using a single self-attention layer.
In these NP variants, for a new task, we first (in a conditioning step) compute the embedding of the context dataset r C .
Afterwards, in a separate querying step, the model makes predictions for target datapoints.
Neural Processes have several desirable properties (1) Scalability : NPs generates predictions more computationally efficiently than Gaussian Processes which scale cubically with the size of the context dataset.
In Table 1, we show the complexity of NP variants relative to the size of the context dataset.
It suffices to perform the conditioning step once; however, for each prediction, the querying step needs to be performed independently. As such, the complexity of the query step is the most important when considering the model’s efficiency. (2) Flexibility : NPs can condition on an arbitrary number of context datapoints and make predictions on an arbitrary number of target datapoints, and (3) Permutation invariance : the predictions are order invariant in the context datapoints.
2.2.1 Transformer Neural Processes
Transformer Neural Processes (TNPs) (Nguyen & Grover, 2022) comprises of a transformer-like architecture that takes both the context dataset and the target dataset as input.
However, unlike previous NP variants that model the predictive distribution of target datapoints independently, i.e., p ( y i | x i , D context ) , TNPs propose to model the predictive distribution of all target datapoints in con- junction p ( y 1 ,...,M | x 1 ,...,M , D context ) as follows:
[Equation]
with r C,T :=
Agg( x C , y C , x T ) where Agg comprises of multiple Self-Attention layers with a masking mechanism such that the context datapoints and target datapoints only attend to the context datapoints.
The output of the aggregator r C,T that corresponds to the target datapoints embeddings are passed to a predictor model for prediction.
In contrast to previous NP variants, TNPs computes its embeddings with both the context dataset and target dataset.
Thus, there is no conditioning step for TNPs.
Instead, when making queries for a batch of target datapoints, the self-attention on both context and target datapoints must be re-computed, requiring quadratic computation O (( N + M ) 2 ) in both the number of context datapoints N and number of target datapoints M .
Importantly, this is quadratic even in very small batch query scenario.
For example, if M = 1 , then the prediction is still quadratic O (( N + 1) 2 ) .
The quadratic complexity of a query regardless of the size of the target batch makes the model computationally costly for deployment.
Methodology
In this section, we introduce Efficient Queries Transformer Neural Processes (EQTNPs), an NP variant that achieves competitive results while being more computationally efficient.
3.1 Reducing the querying complexity of TNP
Due to passing both the context and query dataset through a self-attention mechanism, for TNP, each prediction requires a quadratic amount of computation in the number of context datapoints regardless of the number of target datapoints.
We propose to compute an embedding for each self-attention layer encompassing the information of the context dataset.
This way, we can perform cross-attention to retrieve information from the context dataset for predictions.
Unlike in TNP whose queries require quadratic computation O (( N + M ) 2 )) in both the number of context and target datapoints, the queries by EQTNPs can instead be made in O ( NM ) .
Importantly, this means the model has a per target datapoint query complexity that is linear in the number of context datapoints, greatly benefiting the settings with small batch queries.
For example, if M = 1 , then TNP requires O ( N 2 ) while
EQTNP requires only O ( N ) .
Figure 1 describes how the architecture of TNP is modified to give a more efficient variant that we name Efficient Queries TNPs (EQTNPs). Conditioning Phase:
In the conditioning phase, we compute the embeddings representative of the context dataset
[Equation]
Since the number of context datapoints is O ( N ) , the computational cost per SelfAttention layer is O ( N 2 ) .
Let K be the number of self-attention layers.
Then the embeddings { CEMB 1 , . . . , CEMB K } encode the model’s knowledge of the context dataset. Querying Phase: When making predictions, we make predictions by retrieving information from the context dataset, specifically from { CEMB 1 , . . . , CEMB K } , via multiple cross-attention layers.
{+QEMB 
+}
= QUERY
QEMB i
=
CrossAttention(QEMB i − 1 , CEMB i )
Output = Predictor(QEMB K )
Each CrossAttention layer is linear in CEMB .
Since there is an embedding per datapoint, the computational cost is O ( NM ) .
The multiple cross-attention layers that retrieve information from the context dataset embeddings allows the predictor to retrieve higher-order information just like in
TNPs.  Experiments We evaluate EQTNPs on the meta-regression and image completion setting.
These experiment settings have been used extensively to benchmark NP models in prior works (Garnelo et al., 2018a;
Kim et al., 2019; Lee et al., 2020; Nguyen & Grover, 2022).
In this experiments section, we focus on comparisons with the well-established vanilla versions of NP models.
Specifically, for the case of TNPs, we compare with TNP-D. Due to its masking mechanism and architecture, TNP-D models the probabilistic predictive distribution identical to that of previous NP variants.
As such, it offers a fair comparison to prior NP works.
In contrast, TNP-ND models the predictive distribution differently.
This choice of modelling by TNP-ND benefits settings where the predictive uncertainty of all target datapoints are evaluated in conjunction such as in image completion settings.
However, this does not benefit and can instead be detrimental to settings where the uncertainty of individual datapoints is evaluated independently such as in multi-armed bandits or bayesian optimization.
For a fair comparison of TNP-ND with prior works, it would require modifying the existing NP methods into ND variants (e.g., CNP-ND, CANP-ND, and BNP-ND).
However, this is outside the scope of our work.
4.1 1-D Regression
Method
RBF
Matern
[Table] [Table] All experiments are run with 5 seeds. In this experiment, the model observes a set of N context points pertaining from an unknown function f .
The model is expected to make predictions for a set of M target datapoints that comes from the same function.
In each training epoch, B =functions are sampled from a GP prior with an RBF kernel f i ∼ GP ( m, k ) where m ( x ) =[-
-]
and k ( x, x ′ )
=
σ 2 f exp (cid:16) − ( x − x ′ ) 2 2 l(cid:17) .
The hyperparameters l and σ f are randomly sampled per task.
The hyperparameters are sampled as follows: l ∼ U [0 . 6 , 1 . 0) , σ f ∼ U [0 . 1 , 1 . 0) , N ∼ U [3 , 47) , and M ∼ U [3 , 50 − N ) .
At test-time, the models are evaluated on unseen functions sampled from GPs with RBF and Matern 5 / 2 kernels.
Specifically, the methods are evaluated according to the log-likelihood of the target datapoints.
The number of context datapoints N and the number of evaluation points M are sampled according to the same distribution as in training.
Results: As shown in Table 2, EQTNPs achieve results comparable to that of TNPs.
4.2 Image Completion
In this problem, the model observes a subset of pixel values of an image and predicts the remaining pixels of the image.
As a result, the problem can be perceived as a 2-D meta-regression problem in which x is the coordinates of a pixel and y is the value of the pixel.
The image can be interpreted as a unique function in itself (Garnelo et al., 2018b).
For these experiments, the x values are rescaled to [-1, 1] and the y values are rescaled to [ − 0 .
Randomly selected subsets of pixels are selected as context datapoints and target datapoints.
For these experiments, we consider two datasets:
EMNIST (Cohen et al., 2017) and CelebA (Liu et al., 2015).
EMNIST comprises of black and white images of handwritten letters of 32 × 32 resolution.
In total,classes are used for training, N ∼ U [3 , 197) context datapoints are sampled, and
M ∼ U [3 , 200 − N ) target datapoints are sampled.
CelebA comprises of colored images of celebrity faces that are down-sampled to 32 × 32 and N ∼ U [3 , 197) and M ∼ U [3 , 200 − N ) .
Results: As shown in Table 3, EQTNPs performs comparably with TNPs.
[Table]
The first NP model was Conditional Neural Processes (CNPs) (Garnelo et al., 2018a). CNPs encode the context dataset by encoding the context datapoints via a deep sets encoder. NP (Garnelo et al., 2018b) proposes to use a global latent variable to encode functional stochasticity.
Later, other variants proposed to build in translational equivariance (Gordon et al., 2019), attention to tackle underfitting (Kim et al., 2019), transformer attention (Nguyen & Grover, 2022), bootstrapping (Lee et al., 2020), modeling predictive uncertainty (Bruinsma et al., 2020), and tackling sequence data (Singh et al., 2019; Willi et al., 2019).
From a practical perspective, NPs have also been applied to a large range of problems, including modeling sequences of stochastic processes (Singh et al., 2019), sequential decision making settings (Galashov et al., 2019), and modeling stochastic physics fields (Holderrieth et al., 2021).
Mathieu et al.
(2021) propose a contrastive learning framework to learn better representations for NPs.
Conclusion and Future Work
Prior NP methods have generally been split between (1) computationally efficient variants of Neural
Processes but poor overall performance and (2) computationally expensive attention-based variants of
Neural Processes that perform well. In this work, we propose Efficient Queries Transformer Neural Processes (EQTNPs).
EQTNPs fix the slow querying of the previous state-of-the-art method (TNPs) by proposing to apply self-attention only on the context dataset and retrieve information to the target datapoints via cross-attention mechanisms.
Each query datapoint only requires linear computation in the number of context datapoints.
In the small batch query scenario this is important since in contrast TNPs will require quadratic computation in both the number of context and target datapoints.
In our experiments, we show that EQTNPs achieve results comparable to state-of-the-art.
However, both EQTNPs and TNPs require quadratic computation to perform the conditioning step, limiting their scalability.
As such, an interesting future direction would be to develop methods that are competitive while requiring linear computation in the conditioning step.