Latent Bottlenecked Attentive Neural Processes
Authors: Leo Feng
Hossein Hajimirsadeghi
Yoshua Bengio
Mohamed Osama Ahmed
Venue: ICLR 2023 poster (2023)
Keywords: Neural Processes
Meta-Learning
Uncertainty Estimation

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.

Published as a conference paper at ICLR 2023
LATENT BOTTLENECKED ATTENTIVE NEURAL PRO- CESSES
Leo Feng
Hossein Hajimirsadeghi
Mila - Université de Montréal
Borcalis AI
& Borealis AI
hossein ha fimi raadeghi @borea isai ccom
leo. feng@mi la. quebec
Yoshua Bengio
Mohamed Osama Ahmed
Mila Université de Montréal
Borealis AI
yoshua bengio@mi 1a. quebec
mohamed. o stmeed@orreliisa..coo
ABSTRACT
Neural Processes (NPs) are popular methods in meta-learning that can estimate predictive uncertainty target datapoints by conditioning Previous state-of-the-ar method Transformer Neural Processes (TNPs) achieve strong performance but require quadratic computation with respect 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), new computationally efficient sub-quadratic NP variant, that has 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 trade-off the computational cost and performance according to the number latent vectors. Finally, show LBANPs scale beyond existing attention- based NP variants to larger dataset settings.
context dataset.
INTRODUCTION
Meta-learning aims to learn a model that can adapt quickly and computationally efficiently to new tasks. Neural Processes (NPs) are popular method in meta-learning that models a conditional distribution of the prediction of target datapoint given set of labelled (context) datapoints, providing uncertainty estimates. NP variants (Garnelo et al.) 2018a Gordon et al., 2019 Kim et al. 2019) adapt via 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. In contrast, recent state-of-the-an methods have proposed to use self-attention mechanisms such as transformers. However, these wate-of-the of-the-art methods are computationally expensive in that they require quadratic computation in the number of context datapoints. The quadratic computation makes the method inapplicable settings with large number of datapoints and in low-resource settings. ConvCNPs (Gordon et al. 2019) partly address this problem by proposing to use convolutional neural self-attention mechanism, but this (1) requires
networks to encode the context dataset instead the data satisfy grid-like structure, limiting the method to oow-dinensional settings, and (2) the recent attention based method Transformer Neural Processes (TNPs) have been shown to greatly outperform ConyCNPs.

---

Published as a conference paper at ICLR 2023
Computational Complexity (Big-O)
Method
Training
Evaluation
Step
Condition
Query
CNP (Garnelo et al.
2018a
N M
N
M
CANP (Kmm et al.
2019)
N2 + NM
N
NM
NP Garnelo et al.
2018b
N M
N
M
ANP Kim al.
2019
N2 NM
2
NM
BNP Lee et al.
2020
(N + M)K
KN
KM
BANP (Lee al.
2020)
(N2 NM)K
KN2
KNM
TNP-D (Nguyen & Grover 2022)
(N + M)²
(N + M)²
EQTNP (Ours)
NNM N + NM
N
NM
LBANP (Ours)
(N +
NL
ML
Table 1: Computational Complexity in Big-O notation of the model with respect to the number of context datapoints (N) and the number of target datapoints per batch (M). L and K are prespecified hyperparameters. is the number of latent vectors. K the number of bootstrapping samples for BNP and BANP. It is important that the cost of performing the query step is low.
Inspired by recent developments efficient attention mechanisms, (1) we propose Latent Bottle- necked Attentive Neural Processes (LBANP), computationally efficient NP variant that has querying computational complexity independent of the number of context datapoints. Furthermore, the model requires only linear computation overhead for conditioning the context dataset. The model encodes the context dataset into a fixed number of latent vectors on which self-attention performed. When making predictions, the model retrieves higher-order information from the context dataset via multiple cross-attention mechanisms on the latent vectors. (2) We empirically show that LBANPs achieve results competitive with the ofate-ofthhe-eee on meta- restarregrssion, image completion, and contextual multi-armed bandits. (3) We demonstrate that LBANPs can trade-off the computational cost and performance according to the number of latent vectors. (4) In addition, we show that LBANPs can scale to larger dataset settings where existing frenttoo-bbsed NP variants fail to run because their expensive computational requirements. (5) Lastly, we show that similarly to TNPs, we can propose different variants of LBANPs for different settings.
2 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))) B is sampled from task distribution ((T).. A task T a tuple (X, ), L. q). where X' is the input space, y is the output space, L is the task-specific loss function, and q(x, y) is distribution over data points. During each meta-training iteration, for each T, E T, we sample from dT; a context dataset Deontext
{22, y) } and a target dataset Diurger = {xx * ". 1 where N and M 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 model the probabilistic predictive distribution Po (y) Deceteett where 0 are the parameters of the NN.
2.2 NEURAL PROCESSES
Neural Processes (NPs) are class of models that define infinite family of conditional distributions where one can condition on arbitrary number 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:
p(y|2r, pr, Deumten) plylxrrc))
(1)
2

---

Published as a conference paper at ICLR 2023
with ra == Agg(D contexi) where Agg is a deterministic function varying across the variants that aggeregates Deomext into a finite represention with permutation invariance in the context dataset. These NPs aim to maximise likelihood of the target dataset given the context dataset. Conditional Neural Processes (CNPs) (Garnelo et al. 2018a), proposes an aggregator using DEEPSETs (Zaheer 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 new task, we first (in a conditioning step) compute the embedding of the context dataset rc. Afterwards, in separate querying step, the model makes predictions for target datapoints. Note that for given context dataset and independent of the number of batches of predictions, the conditioning step needs only be performed once.
Neural Processes have several desirable properties (1) Scalability: NPs generate predictions more computationally efficiently than Gaussian Processes which scale cubically with the size of the context dataset. In Table [ step once but 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 Is efficiency. (2) Flexibility: NPs can condition an arbitrary number of target datapoints, and (3) Permutation invariance: the predictions are order invariant in the context datapoints.
we show the complexity of NP variants. It suffices to perform the conditioning
arbitrary number of context datapoints and make predictions on
2.2.1 TRANSFORMER NEURAL PROCESSES
Transformer Neural Processes (TNPs) (Nguyen & Grover, 2022) comprises of transformer-lik 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... p(yi x4, Do TNPs propose to model the predictive distribution of all target datapoints in con- junction p(y MM D Dootess) as follows:
M
(2) with TC,T' := Agg (C, 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 rc,r that corresponds to the target datapoints embeddings passed to predictor model for prediction.
p [I....MM D contex) ==
Three variants of TNPs have been proposed: TNP-D, TNP-ND, and TNP-A However, TNP-A, as the authors highlighted has low tractability due to being an autoregressive model. As such, TNP-D and TNP-ND are proposed as the practical methods for comparison. Specifically, TNP-D proposes a predictor model comprising of vanilla MLP which outputs the mean and diagonal covariance matrix for the target datapoints. In contrast, TNP-ND outputs a full covariance matrix for all target datapoints. TNP-ND's s predictor architecture comprises of self-attention layers applied to the embeddings of target datapoints and a network to output a low-rank matrix to compute full covariance matrix either via cholesky decomposition or low-rank decomposition. We highlight that although the ND predictor architecture was only introduced for TNP, in fact, it can be easily extended to other NP variants such as CNP, CANP, and BNP. which case, would result in models such as CNP-ND, CANP-ND, and BNP-ND.
In contrast to previous NP variants, TNPs computes its embeddings with both the context dataset and target dataset. Thus, there is 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(NN + M)² in both the number of context datapoints N and number of target datapoints M. The quadratic complexity makes deployment of the model (1) computationally costly for deployment and (2) restricts the model from scaling to settings with larger amounts of context or target datapoints.
3 METHODOLOGY
In this section, we introduce Latent Bottlenecked Attentive Neural Processes (LBANPs). an NP variant that achieves competitive results while being computationally efficient. Firstly, show how to fix the quadratic computation query issue in TNPs. Secondly, we propose, LBANP, a model that utilises a latent bottlenecked attention mechanism reduce the querying step's complexity to
3

---

Published as a conference paper at ICLR 2023
constant computation per target datapoint and conditioning step's overhead to linear computation. Lastly, we show how LBANP can be extended to other variants similar to TNP depending on the problem setting.
3.1 REDUCING QUERYING COMPLEXITY
Due to passing both the context and query dataset through self-attention mechanism, for TNP, each prediction requires a quadratic amount of computation both N and M. Unlike previous attentive-based NP variants (Kim et al. 2019), the multiple self-attention layers in EQTNP allow for higher order interactions between the target and context datapoints.
This way,
perform eross ansss-attntiion retrieve information from the context dataset for predic- tions. Unlike in TNP whose queries require quadratic computation O((N + M)²) in both the number of context and target datapoints, the queries instead be made in O(N M). Figure describes how the architecture of TNP is modified to give more efficient variant that we name Efficient Queries TNPs (EQTNPs). Although EQTNPs are more efficient than TNPs, they unfortunately still require quadratic computation to compute the embeddings of the context dataset, limiting its applications.
Used when Querying
Dataset Dataset
Context Query
Dataset
Context
xxD MxD
Self Attention
Self- Attention
Self Attentionn
x Dr:
SA
Predictor
Predictor
Output
M x
Dataset
Query
Output
Figure 1: Model architecture for TNPs (left) and EQTNPs (right). The blue highlights the parts of the architecture used when querying. SA refers to Self-Attention and CA refers to CrossAttention.
3.2 LATENT BOTTLENECKED ATTENTIVE NEURAL PROCESSES
x
Dataset
Context
Used when Querying
Vectors
Latent
xx
52
Predictor
Query Dataset M x
Output
CA
Figure 2: Architecture for Latent Bottlenecked Attentive Neural Processes (LBANPs). The blue highlights the parts of the architecture used when querying.
Inspired by the recent successes in efficient attention mechanisms Lee et al. 2019 Jaegle et al. 2021b), we propose Latent Bottlenecked Attentive Neural Processes (LBANPs), an NP variant that reduces the quadratic conditioning complexity to linear and further reduces the query complexity to linear in the number of target datapoints, while retaining similar performance. See Figure 2 for the architecture.
Conditioning Phase: Instead of performing Self-Attention on the embeddings of all datapoints, we propose to reduce the dimensionality of the context dataset embeddings to set of L x D,, latent vectors LEMB where LEMB, is meta-learned during training. The size of the latent vectors (L DL) represents the size of the latent bottleneck and is set according to prespecified hyperparameters L and D,.. By interleaving retrieving information from the context dataset (CONTEXT) via a cross- attention and performing self-attention, we can compute embeddings of the context dataset that
4

---

Published as a conference paper at ICLR 2023
comprise of higher-order interactions between datapoints that is typically achieved by multilayer self-attention mechanism. Note that LBANPs also satisfy the property of being invariant to the ordering of the context and target dataset. By scaling the number of latent vectors, we can increase the size of the bottleneck, allowing for more expressive models.
LEMB, SelfAttention (Cross (LEM;: CONTEXT))
Since the cross-attention is between the latent vectors and context dataset, thus the cost of this computation is O(NL). Furthermore, since the self-attention is only performed set of L latent vectors, thus it only requires ((1)) computation. Therefore, the amount of computation required to compute these latent vectors is linear in the number of context datapoints i.e.. O + L ) = (N) since L is prespecified constant. The outputs of the conditioning phase are the latent embeddings LEMB, representative of the context dataset. In contrast with EQTNP, the number of (latent) vectors representing the context dataset is constant and is thus independent of the number of context datapoints. This crucial difference contributes to a significant improvement in the overall computational complexity over both EQTNP and TNP.
Querying Phase:
The querying phase similar to that of EQTNP that multiple cross-attention layers are used to retrieve information from the context dataset. The multiple ross-attention layers allow for the model to compute higher-order information between a target datapoint and the context datapoints to use in its prediction. Specifically, this is performed follows:
QEMB, = QUERY
QEMB, = ErosssAttentionn(EMMM. LEMB,)
Output = Predictor(QEMByr)
Notice that computing QEMB (short for Query Emheddings) is dependent the number of latent vectors (a constant L) and is independent of the number of context datapoints. As a result, making the prediction of single datapoint is overall constant in computational complexity, i.e., O1). Naturally, for batch of M datapoints, the amount of computation thus would be ((M). Furthermore, note that LEMB only requires constant amount of memory to store regardless of the number of datapoints. As such, unlike TNP, only constant memory is required to make a query per datapoint, allowing the method to be more applicable in settings with imited-resources.
3.2.1 LBANP VARIANTS
Similar TNP-Not Diagonal (TNP-ND) (Nguyen & Grover 2022) which predicts covariance matrix for batch of targets, we can propose a similar variant of LBANP for settings where the predictions on target datapoints ought to be conditionally dependent. When making predictions using the target datapoint embeddings, TNP-ND first applies self-attention mechanism solely on the target datapoint embeddings before passing into an MLP that outputs a low-rank matrix that is used to construct approximations of the full covariance matrix e.g., via Cholesky decomposition. Similarly, we could follow the same framework, resulting in a method called LBANP-ND.
EXPERIMENTS
We evaluate Latent Bottlenecked Attentive Neural Processes (LBANPs) on several tasks: meta regression, image completion, and contextual bandits. These experiment settings have been used extensively to benchmark NP models in prior works (Garnelo al., 2018a| Kim et al. 2019 Lee
et al. 2020 Nguyen& & Grover (022). We compare LBANPs with the following members of the NP family: Conditional Neural Processes (CNPs) (Garneoo et al. (0188). Neural Processes (NPs)
(Garnelo et al. 20188) 8b), Bootstrapping Neural Processes (BNPs) (Lee et al. 2020), and Transformer Neural Processes (TNPs) (Nguyen & Grover 2022). In addition, we compare with their attentive variants (Kim et al. 2019) (CANPs, ANPs, and BANPs). Implementation details of the baselines and M For the sake of consistency, we evaluated LBANPs with a fixed number of latent vector (L= 8 or 7. 128).
LBANPs are included in the Appendix
In this experiments section, we focus on comparisons with the well-established vanilla versions of NP models. Specifically, for the case of TNPs, we directly compare with TNP-D. Due to its masking
The code is available at http://tthh...ooommaaaaaa//ltltt
5

---

Published as a conference paper at ICLR 2023
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 and TNP-A model the predictive distribution differently. TNP-A in particular has low tractability due to its autoregressive nature. For a fair comparison of TNP-ND with prior works, 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. Nonetheless, for the sake of completeness, we (1) include results for TNP-ND and TNP-A in the tables and (2) detailed how to construct the ND variants previously in Section 3.2. and include comparisons with TNP-ND in the Appendix.
The aim of the experimental evaluation is to answer the following questions: (1) What is the performance of LBANP compared to the existing NP methods? (2) Does LBANP scale to problems with larger number of context datapoints? (3) How does the number of latent vectors affect the performance of LBANPs?
4.1 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 2-D meta-segression problem in which x is the coordinates of a pixel and y is the value of the pixel. The image can be interpreted as 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.5,, 0.5]. 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 x 32 resolution. In total, 10 classes are used for training. NN ~ =2[3, 197) context datapoints are sampled, and M ~ 200-200 NN) target datapoints are sampled. CelebA comprises of colored images of celebrity faces. We evaluate the models on settings of various resolutions. For CelebA32, the images are down-sampled to 32 x 32 and N ~ U[3 197) and M U(3, 200 N). For CelebA64, the images down-sampled to 64 x 64 and N ~ [3, 797) and M ~ [3, 800 NN.. For CelebA the images are down-sampled 128 128 and N [3, 1597) and M [33, 1600 - N).
Results: As shown in Table 3. using only 8 latent vectors, LBANP with only 8 latent vectors significantly outperforms all baselines except for TNP-D. In addition, we see in both CelebA and EMNIST experiments that going from 8 to 128 latent vectors improves the performance. For CelebA (32 x 32), using 128 latent vectors, L.ANP outperforms TNP-D, achieving state-of-the-art results. For CelebA (64 x 64), few attention-based NP variants (ANP and BANP) were not able to be trained due to the quadratic memory required during training. This issue was also previously reported by Garnelo et al. (2018a) Furthermore, we show that by increasing the number of latent vectors (see Section 3.2. and the Appendix) LBANP can perform competitively with TNP-D. For CelebA (128 x 128), due to its large number of datapoints, all existing attention- based NPs (CANP, ANP, BANP, and TNP-D) were not able to be trained. Our results show that LBANP outperforms all trainable baselines (CNP, NP, and BNP) by significant margin. This demonstrates the advantage of our proposed method in terms of scalability and performance. Let (N + M) Mmmar represent the maximum number of total points (including both context and target points). In the CelebA (128 128) task, L = 8 = 0.5%(N + M) mea and L = 128- 8%(N + M) mazz.
4.2 1-D REGRESSION
In this experiment, the model observes 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 16 functions are sampled from GP prior with an RBF kernel GP k) where m(z) 0 and k(xz,²' oxexp (a-z²)
The hyperparameters /
2
and o are randomly sampled per task. The hyperparameters are sampled as follows: I ~ L4(0.6, 1.0). a ~ 00,1, 1.0), N ~ L4(3, 47), and M ~ |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 ligeeiildoood of the targets. The number of context datapoints N and the number of evaluation points M are sampled according to the same distribution as in training.

---

Published as a conference paper at ICLR 2023
CelebA
EMNIST
Method
32x32
64x64
128x128 2.55 = 0.02
Seen (0-9)
Unseen (10-46
CNP
0.01
+.00
0.49 # 0.01
CANP
2.66
3.05
0.94
0.82 0.01
NP
2.08
2.60
2.67 = 0.01
0.79 0.01
ANP
2.00
0.98
0.00
BNP
2.06
0.00
0.01
0.73
BANP
3.09
1.01 1.06
0.94
TNP-D
0.09
5.41 0.01
1..0
EQTNP
3.91
5.29
1.00
1.27
LBANP (8)
0.00
4.39
4.94 + 0.03
1.34
1.05
LBANP (128)
0.02
5.09
5.84 + 0.01
1.39
1.17
TNP-ND
0.00
1.00
TNP-A
5.82 0.01
1.54
1.41 =
Table 3: Image Completion Experiments. Each method is evaluated with 5 different seeds according the log- lekellkdioood (higher better). The "dash" represents methods that could be run because of the large memory or time requirement.
Method
RBF
Matern 5/2
CNP
0.26 + 0.02
0.04 # 0.02
Results: As shown in Table 2] with latent vec- tors, LBANP outperforms all the baselines by large margin except TNP-D. Moreover, when us- ing 128 latent vectors instead of 8. LBANP has competitive performance with TNP-D. It is impor- tant note that the performance of LBANP can be further improved by increasing the number of latent vectors as shown Figure 3a In this task, L 8 16%(N + M)mar and L 128 256%(M + M) Mmmz
CANP
0.79
0.62 # 0.00
NP
0.27
0.07 = 0.01
ANP
0,81
0.63
BNP
0.18 0.002
BANP
0.01
0.66 0.00
TNP-D
1.00
0.01
EQTNP
0.92 1.02
LBANP (8)
1.20
0.75
LBANP (128)
1.02
0.85 + 0.02
TNP-ND
1.46
1.02 + 0.00
TNP-A
1.63 0.00
1.21 + 0.00
4.3 CONTEXTUAL BANDITS
Table 1-D Meta-Regresse Experiments with log- logelikeli metric (higher is better). All experiments are run with seeds.
In this problem previously introduced in Riquelme
et al. (2018), unit circle is divided into 5 regions: low reward region and high reward regions. scalar 5 determines the size of the low reward region.
The 4 high reward regions have equal sizes. The agent, however, has to select amongst 5 arms (each representing one of the regions) given its 2-D co-ordinate X and its actions in previous rounds. If || < 8, then the agent is within the low reward region. In this scenario, the optimal action is to select arm 1 which provides reward r ~ (1.2, 0.012). The other arms provide a reward r ~ (1100, 0.012). If the agent is a high-reward region (i.e., ||X|| >> 8), then the optimal arm would be to pull of the other arms corresponding to the region. If the optimal arm is pulled. then the agent receives reward of N (50.0, 0.012). Pulling arm will reward the agent with (1.2, 0.012). Pulling any of the other arms rewards the agent with r (1.00, 0.012).
In each training iteration, B different wheel problems are sampled ( according uniform distribution ~ (0, 1). For the iteration, M 50 points are sampled for evaluation and N 512 points are sampled context. Each datapoint tuple (X, r) where X is the co-ordinate and is the reward values for the 5 arms. The training objective is to predict the reward values for all arms given the co-ordinates. We evaluate LBANP and EQTNP on settings with varying values. In each of the experiments, different seeds for each value of where each run corresponds to 2000 steps. At each step, the agent selects the arm according to its Upper-Confident Bound (UCB) and receives the reward value from the arm. The performance is measured according to the cumulative regret.
report the mean and standard deviation of the cumulative regret for 50
Results: Table show that LBANPs outperform all methods except for TNPs by significant margin. For lower values, we see that LBANP performs similarly to TNP-D and occasionally outperforms TNP-D in the case of 5 0.95 even with 8 latent vectors. Interestingly, that EQTNP
7

---

Published as a conference paper at ICLR 2023
Method
8= 0.7
à == 0.95
à= 0.99
== 0.995
Uniform
100.00 + 1.18
100.00 + 3.03
100.00 44.16
100.00 7.52
8.0.00
CNP
4.08 0.29
4.34
8.01 0.440
26.78 0.85
38.25
CANP
#9.93
11.69
24.49
47.33 - 20.49
# 17.87
NP
0.16
2.96
4.24
2 1..4
ANP
0.12
4.35
+.59
0.67
27.65 # 0.95
BNP
1.07
2.11
5.22
62.50 4.85
BANP
16.58
31.10
52.59 + 18.11
49.35 = 14.52
TNP-D
0.94
1.70
0.45
3.57
4.68 = 1.09
EQTNP
0.15
1.37
2.09
7.87 0.45
11.15 0.13
LBANP (8)
1.28
+
0.16
7
10,41 = 0.20
LBANP (128)
0.36
+
1.25
6.13
8.76 + 0.15
TNP-ND
0.76
1.41 = 0.98
1.61
4.98
2.84
7.22 = 3.28
TNP-A
3.67 4.88
4.04 2.38
2.36
5.79 5.27
9.29 + 7.62
Table 4: Contextual Multi-Armed Bandit Experiments with varying Models are evaluated according to cumulative regret (lower is better). Each model is run 50 times for each value of 8.
outperforms TNP-D substantially for = 0.7, 0.9, 0.95. In this task, L = 8 = 0.4% + M) M)mo and T, 128 6.4%(N M) M)mar
4.4 ABLATION STUDIES
4.4.1 NUMBER OF LATENTS
In this experiment (see Figure 3ab CelebA (32 x 32), we vary the number of latent vectors from 8 until 256 to evaluate the effect of different number of latent vectors. The experiments show that by scaling the number of latent vectors, the performance increases. This inline with previous results that show that LBANP with 128 latent vectors outperforms LBANP with 8 latent vectors. Specifically, in this CelebA32 setting, we see that 128 latent vectors already achieves state-of-the- results but results can be further improved with 256 latent vectors. Furthermore, we can see that even with latent vectors LBANPs achieve strong results.
4.4.2 EMPIRICAL ANALYSIS OF COMPUTATIONAL COMPLEXITY
Empirical Time and Memory Cost: In these experiments, we show the empirical time cost to perform predictions with various Neural Process models. In Figure 36 and 3d we fix the number of target datapoints and vary the number context datapoints. We see that TNPs required memory and time scales quadratically. In contrast, we see that EQTNPs and LBANPs remain efficient regardless of the number of datapoints. Specifically for LBANPs, can be observed that LBANP queries utilize the same amount of time and memory as the number of context datapoints increases. These results confirm the big-O complexity analysis in Table [ As an example, when N = 1400, LBANP (L 128) uses 1172 MBs of memory and TNP uses 13730 MBs of memory (more than 11 times the memory LBANP used), showing a significant benefit of using LBANPs in practice.
5 RELATED WORK
The first NP model was Conditional Neural Processes (CNPs) (Garnelo et al. 2018a). CNPs encode the context dataset by encoding the context datapoints via deep sets encoder. NP (Gamelo 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 tackle underfit- ting (Kim et al.) 2019), transformer attention (Nguyen & Grover (022), bootstrapping (Lee et al.) 2020), modeling predictive uncertainty (Bruinsma et al. 2020), and tackling sequence data (Singh et al. 2019 Willi al.) 2019). Concurrent with work, Rastogi al. (2023) also proposed an attention based NP variant with linear complexity based on inducing points. From a practical perspective, NPs have also been applied 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 al. 2021). Mathieu al. (2021) propose
8

---

Published as a conference paper at ICLR 2023
CelebA (32
Empirical Memory Complexity
Empirical Complexity
12000
ЕОТNP
LBAN (128)
CAS
CPP
10000
LBAUP (32)
LBAN
DOTNP
TIP-D
134
LBAU (120)
LEAN (322)
25
715.0
6000
Time
22
230
2000
Conteat Ostancints 3020
Nkancher of Lalients
Context Ostaecints
(a) Number of Latents Analysis
(b) Memory Analysis
(c) Time Analysis
Figure 3: (a) Analyzing the relation between number of latent vectors in LBANP and the model performance. Since the other models do not have a latent bottleneck, their performance is consistent and are portrayed as a dotted line. (b) Comparison of TNPs, EQTNPs, and LBANPs in terms of the amount of memory to make prediction. (c) Comparison of TNPs, EQTNPs, and LBANPs in terms of the amount time to make a prediction with respect to the number of context datapoints.
contrastive learning framework to learn better representations for NPs. For an in-depth survey, refer the reader to the recent survey paper on Neural Processes (Jha et al. 2022).
Transformers have been shown achieve amazing results while being very flexible, but the quadratic scaling in terms of the number of inputs due to its self-attention layers renders it inapplicable to settings with large amounts of inputs e.g., an image. Several groups have proposed strategies image settings to reduce the size of the input to the Transformer e.g., subsampling the input (Chen et al.)
2020) and preprocessing the input using convolutions (Dosovitskiy et al., 2020 Wu al. 2020). Other groups have proposed to modify the internal self-attention mechanism of Transformers to obtain enhanced efficiency. Specifically, Set Transformers and related works (Lee et al. 2019 Goyal al. 2021) induce latent bottleneck by mapping an input array back and forth via cross-attention between an array with fewer elements or manipulating the size of the array Jaegle et al. (2021b a). Perceiver (Daegle et al. 2021b proposes a deep latent network which takes input simple 2D byte array (e.g.. pixels) and evaluates the proposed model on image classification, audio, video, and point clouds. Perceiver 10 (Jaegle et al. 2021a) extends Perceiver to inputs and outputs of different sizes.
CONCLUSION
Prior NP methods have generally been split between (I) computationally efficient (sub-quadratic) vari- ants of Neural Processes but poor overall performance and (2) computationally expensive (quadratic) attention-based variants of Neural Processes that perform well. In this work, we fix the slow querying of the previous state-of-the-ar method, TNPs, by proposing to apply self-attention only on the context dataset and retrieve information the target datapoints via cross-attention mechanisms. Taking the idea a step further, we introduce Latent Bottlenecked Neural Processes (LBANPs) a new member of the NP family that achieves results comparable with the state-of ffate-of-the-ar while being computationally efficient. Specifically, the deployed LBANPs have a linear cost in the number of context datapoints for conditioning on the context dataset. Afterwards, each query requires linear computation in the number of target datapoints. In our experimental results, we show that LBANPs achieve results comparable to state-of- ofrtre--rtttt while being more scalable. In our analysis, we show that LBANPs can trade-off the computational cost and performance according to a single hyperparameter (L: the number of latent vectors). This allows a user to easily select the expressivity of LBANPs depending on the computational requirement of the setting. Importantly, unlike with TNPs, the amount of computation (both time and memory) required query single target datapoint for LBANPs will always be constant regardless of the number of context datapoints.
REPRODUCIBILITY STATEMENT
The code is available at https///gthub.com/BBrellcc Sotreuliss//lttte--btllnneecccc bottlenecked-nng Hyperparameters and implementation details available in the Appendix, Section 3 and Section

---

Published as a conference paper at ICLR 2023
ACKNOWLEDGEMENTS
The authors acknowledge funding from CIFAR and the Quebee government.
REFERENCES
Wessel Bruinsma, James Requeima, Andrew YK Foong, Jonathan Gordon, and Richard E Turner. The gaussian neural process. In Third Symposium Advances in Approximate Bayesian Inference. 2020.
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1691-1703. PMLR, 13-18 Jul 2020. URL hnttps://proceed
mlr. wtwpppesssvvv//veeeeee
html
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. In 2017 international joint conference on neural networks (IJCNN), pp. 2921-2926. IEEE, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2020.
Alexandre Galashov, Jonathan Schwarz, Hyunjik Kim, Marta Garelo,, David Saxton, Pushmeet Kohli, SM Eslami, and Yee Whye Teh. Meta-learning surrogate models for sequential decision making. arXiv preprint 19X1 2019.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. Conditional neural processes, In International Conference on Machine Learning, pp. 1704-1713, PMLR, 2018a.
Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv: 1807.00622 2018b.
Jonathan Gordon, Wessel P Bruinsma, Andrew YK Foong, James Requeima, Yann Dubois, and Richard E Turner. Convolutional conditional neural processes. In International Conference on Learning Representations, 2019.
Anirudh Goyal, Aniket Rajiv Didolkar, Alex Lamb, Kartikeya Badola, Nan Rosemary Ke, Nasim Rahaman, Jonathan Binas, Charles Blundell, Michael Curtis Mozer, and Yoshua Bengio. Coordi- nation among neural modules through a shared global workspace. In International Conference on Learning Representations, 2021.
Peter Holderrieth, Michael Hutchinson, and Yee Whye Teh. Equivariant learning of stochastic fields: Gaussian processes and steerable conditional neural processes. In International Conference on Machine Learning. pp. 4297-4307, PMLR, 2021.
Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, et al. Perceiver io: A general architecture for structured inputs & outputs. In International Conference on Learning Representations, 2021a.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. International conference on machine learning, pp. (651-4664, PMLR, 2021b.
Saurav Jha, Dong Gong, Xuesong Wang, Richard E Turner, and Lina Yao. The neural process family: Survey, applications and perspectives. arXiv preprint 0rXiv:2209005//, 2022.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. 2019.
10

---

Published as a conference paper at ICLR 2023
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: framework for attention based permutation-invanne neural networks. In International Conference Machine Learning, pp. 7744-375.. PMLR, 2019.
Juho Lee, Yoonho Lee, Jungtaek Kim, Bunho Yang, Sung Ju Hwang, and Yee Whye Teh. Boot- strapping neural processes. Advances in neural information processing systems, 33:6606-6615, 2020.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In Proceedings of International Conference Computer Vision (ICCV), December 2015.
Emile Mathieu, Adam Foster, and Yee Teh. On contrastive representations of stochastic processes. Advances in Neural Information Processing Systems, 34:28823328888 2021.
Tung Nguyen and Aditya Grover. Transformer neural processes: ahocertainty-wware meta learning via sequence modeling. In International Conference on Machine Learning, pp. 16569-16594, PMLR, 2022.
Richa Rastogi, Yair Schiff, Alon Hacohen, Zhaozhi Li, lan Lee, Yuntian Deng, Mert R. Sabuncu, and Volodymyr Kuleshov. Semi parametric inducing point networks. In International Confer- Learning Representations, 2023. URL https:: https/////.ppeeee.e.eeeeet/////e
FE99- fDrWd5
Carlos Riquelme, George Tucker, and Jasper Snoek. Deep bayesian bandits showdown: An empirical comparison of bayesian deep networks for thompson sampling. arXiv preprint arXiv: 1802.09127, 2018.
Gautam Singh, Jaesik Yoon, Youngsung Son, and Sungjin Ahn. Sequential neural processes. Advances in Neural Information Processing Systems, 32, 2019.
Timon Willi, Jonathan Masci, Jürgen Schmidhuber, and Christian Osendorfer. Recurrent neural processes. arXiv preprint arXiv:1906.05915. 2019.
Bichen Wu, Chenfeng Xu, Xiaoliang Dai, Alvin Wan, Peizhao Zhang, Zhicheng Yan, Masayoshi Tomizuka, Joseph Gonzalez, Kurt Keutzer, and Peter Vajda. Visual transformers: Token-based image representation and processing for computer vision. arXiv preprint wrXiv:2006.03677 2020.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poezos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.

---

Published as a conference paper at ICLR 2023
A
APPENDIX: ADDITIONAL METHODOLOGY DETAILS
A.1 EFFICIENT QUERY TRANSFORMER NEURAL PROCESSES
Conditioning Phase: In the conditioning phase, we compute the embeddings representative of the context dataset
CEMB, = CONTEXT
CEMB; CalfAttentionCEEMBB
Since the number of context datapoints O(N), the computational cost per SelfAttention layer is O(N2). Let K be the number of self-attention layers. Then the embeddings [CEMBI,.... CEMB&} encode the model's 'ss knowledge of the context dataset.
Querying Phase: When making predictions, we make predictions retrieving information from the context dataset, specifically from [CEMB),..., CEMBB, }, via multiple cross-attention layers.
QEMB, QUERY
QEMB; = (attAttteentiinnMM..... CEMB,)
Output = Predictor(QEMMxx)
Each CrossAttention layer is linear in CEMB. Since there an embedding per datapoint, the computational cost is O(NM). The multiple cross- coos-aatenttion layers that retrieve information from the context dataset embeddings allows the predictor to retrieve higher-order information just like in TNPs.
A.2 COMPUTATIONAL AND MEMORY COMPLEXITY
In Table include the computational and memory complexity for the NP variants with N context datapoints and M target datapoints. Note that in the case of a complexity of O(M), the computational and memory cost for single target datapoint constant.
Computational (and Memory) Complexity (Big-O)
Method
Training
Evaluation
Step
Condition
Query
CNP
N + M
N
M
CANP
N + NM
N2
N
NP
N + M
N
M
ANP
N2 + NM
N2
NM
BNP
N + M
BANP
N2 NM
N2
NM
TNP-D
(N M)²
(N + M)²
EQTNP
N HNM
N
NM
LBANP
(N + M L)L
NL L
ML
TNP-ND
(N M)
N
(N M)²
LBANP-ND
N + M
N
M
LEANP-END
N + M
M
Table 5: Computational and Memory Complexity in Big-O notation the model with respect to the number of context datapoints N and number of target datapoints in batch M. L and K are prespecified hyperparameters. samples for BNP and BANP. When using NP at test-time, when making multiple predictions with the model, conditioning is one-time cost but querying is not. As such, it is most important that the cost of querying is low.
the number of latent vectors. K the number of bootstrapping
12

---

Published as a conference paper at ICLR 2023
B APPENDIX: ADDITIONAL EXPERIMENTAL DETAILS AND RESULTS
B.1 IMPLEMENTATION AND HYPERPARAMETER DETAILS
We use the implementation of the baselines from the official repository of TNPs (see wttpa/uwo.ooouuunnnn/n/aaia
and its respective paper (Nguyen & Grover 2022)) The hyperparameters for the CelebA experiments are the same regardless of the resolution. EQTNP ex- periments follow almost identically the TNP experiment hyperparameters. The only difference is that EQTNPs architecture slightly different than that of TNP. Specifically, EQTNP has a cross-attention mechanism retrieve information from the context dataset. The number of cross-attention layers in EQTNP equivalent to that of the number of self-attention layers in TNP. Similarly. LBANP hyperparameters are almost identical to that of EQTNP and that of TNP. LBANP consists of blocks each comprising of a cross-attention mechanism from the context dataset to the latent space and a self-attention mechanism in the latent space. In LBANP, the number of blocks for the conditioning phase is equivalent to the number self-attention mechanisms in the conditioning phase of EQTNP.
Due to its cross-attention mechanisms, LBANP allows for varying embedding sizes for the context dataset De. latent vectors D,, and query dataset Dc. For simplicity, we set Dc D2 DQ 64, following TNP's embedding size of 64. Similarly for EQTNP, we set D. DQ 64. We do not tune the number of latent vectors (L). Instead, we showed results for LBANP with L = 8 and == 128 latent vectors. The remainder of the hyperparameters is the same for LBANP, EQTNP, and TNP.
For the ND experiments, we follow TNP-ND (guuyen Grover 2022) and use cholesky decomposi- tion for our LBANP-NP and LBANP-ENP experiments. All experiments were either run on GTX 1080ti (12 GB RAM) or P100 GPU (16 GB RAM). When verifying computationally expensive models were trainable for CelebA64 and CelebA 128, used the P100 GPU (the GPU with larger amounts RAM).
B.2 ANALYSIS: NUMBER OF LATENTS
In this section, we include an additional plot (see Figure 4 that analyses the effect of varying the number of latent vectors. We clearly see that scaling the number of latent vectors improves the performance LBANP. 512 latent vectors, we see that IBANP achieves performance competitive with that of TNP-D.
CelebA (64 x 64)
5.5
5.0
Log-L likelllood
4.5
BNP
NN-U
3.5
LEAND
26
64
128
236
512
Number of Latents
Figure 4: Analysis regarding the number of latent vectors in LBANP.
B.3 ABLATION: INTERLEAVING CROSS-ATTENTION LAYERS
Naively an idea is to first compute embeddings for the context dataset and pass to a Trans- formerDecoder like layer, but this incurs quadratic complexity in the number of target datapoints. Alternatively, retrieving information via a single cross-attention layer. However, this limits the amount of information that can be retrieved which affects performance. We show that in this ablation. Specifically, we verify whether interleaving the cross-attention layers important for LBANP. In this experiment, consider LBANP-L (LBANN-LLatt where instead of multiple cross-attention layers
13

---

Published as a conference paper at ICLR 2023
for querying, the model only single cross-attention layer which attends over the last latent embedding. 6 we see that LBANP outperforms LBANP-L. As such, indeed interleaving the multiple cross-attention layers to retrieve higher-order information from the context dataset for the target datapoint is important.
CelebA
Method
32x32
64x64
LBANP-L (8)
3.35 # 0.01
4,10 0.04
LBANP (8)
3.50 0.05
4.39 0.02
LBANP-L (128)
3.78 # 0.03
4.90 0.05
LBANP (128)
3.97 + 0.02
5.09 = 0.02
Table 6: Ablation Image Completion Experiments.
B.4 ANALYSIS: COMPUTATIONAL. COMPLEXITY
B.4.1 EMPIRICAL TIME COMPLEXITY
In these experiments, we show the empirical time cost perform predictions with various Neural Process models. In Figure 5a we fix the number of target datapoints and vary the number of context datapoints. Figure 5b we fix the number of context datapoints and vary the number of target datapoints.
In both experiments, we that TNP's require significantly more time as the number of datapoints scale. In contrast, we see that EQTNPs and LBANPs remain efficient regardless of the number of datapoints. Importantly, we see that regardless of the number of context datapoints, LBANPs queries use the same amount of time.
Empirical Time Complexity
Empirical Time Complexity
Inference Time (in milliseconds)
LBANP (8)
Inference (in milliseconds)
LBANP (8)
LBANP (128)
LBANP (128)
35
TNP-D
TNP-D
EQTNP
EQTNP
30
20
15
10
10-
1000 1500 2000 # of Context Datapoints
3000
500
2500
# of Target Datapoints
(a) Varying of context datapoints
(b) Varying # of target datapoints
Figure 5: Comparison of TNPs, EQTNPs, and LBANPs in terms of the amount of time to make prediction.
B.4.2 EMPIRICAL MEMORY COMPLEXITY
In these experiments, we show the empirical memory cost perform predictions with various Neural Process models. In Figure 6a we fix the number of target datapoints and vary the number of context datapoints. The figure shows that with respect to the number of context datapoints: (1) TNP's memory cost scales quadratically, (2) EQTNP's memory cost scales linearly, and (3) LBANP's memory cost 1400, LBANP (L 128) 1172 MBs of memory and TNP uses 13730 MBs of memory (more than 11 times the memory LBANP used), showing a significant benefit of using LBANPs in practice.
stays constant. These results confirm the big-O complexity analysis Table ! When
In Figure 6b we fix the number of context and target datapoints and vary only the number of latent vectors for LBANPs. The plot shows that the memory cost of LBANPs indeed scales linearly with the number of latent vectors.
14

---

Published as a conference paper at ICLR 2023
Empirical Memory Complexity
LBANP's Empirical Memory Complexity
14000
TNP-D
2400
EQTNP
LBANP (B)
2200
10000
LBANP (32)
LBANP (128)
8000
LBANP (512)
1800
6000
4000
1400
2000
1200
200
400
600
800
1000
1200 1400
200
400
600
800
1000
# of Context Datapoints
of Latent Vectors
Varying of context datapoints
(b) Varying # of Latent Vectors
Figure 6: (a) Comparison of TNPs, EQTNPs, and LBANPs in terms of the amount of memory to make a prediction. (b) Plot of the amount of memory to make prediction relative the number of latent vectors.
In Figure 7 we include a scatterplot of the memory used when querying for varying number of context datapoints (N) and target datapoints (M). The plot shows that TNP-D uses significantly more memory than LBANPs when making predictions. Importantly, increases significantly with gains in performance as the number of datapoints increase. contrast, memory to achieve the large performance gains. The
notice that TNP's memory
LBANP does not require significantly memory used for querying in fact constant as the number of context datapoints increase.
Empirical Memory Complexity
8
N-1400,
N=1000, M-200
N=0000, M=400
N=1400,
M=200
N=1000, M-200
N=1000, M=400
N-600,
N=800, Ms400
N=800, M-200
N--800
M=400
Log- Li-eekklihood
6
N=400,
M=400
5
N=2000
N-400, M+400
N-200, M-200
LBANP (512)
LBANP (8)
TNP-D
3
2000
4000
6000
8000
10000
12000
14000
Memory Cost (in MBs)
Figure 7: Scatterplot of the memory used for varying number of context datapoints (N) and target datapoints (M). Since the performance of LBANP (8) is comparable as the number of datapoints increase. As such, for clarity, the majority of datapoints for LBANP (8) without text specifying N and M.
15

---

Published as a conference paper at ICLR 2023
B.4.3 ANALYSIS: CONVERGENCE
Following TNPs, we run our experiments for 200 epochs. In Fig we see that TNPs converge slightly quicker than that of LBANPs and EQTNPs. By comparing LBANP (L = 128) and LBANP (L 8), we see that increasing the number of latent vectors help LBANPs learn quicker.
Learning Curve
4.0
Log-Likellln Lieellkilloddd
3.0
2.5
2.0
LBANP (8)
(128)
1.5
TNP-D
EQTNP
1.0
100
125
175 200
# Ebochs
Figure 8: Learning Curve comparison of TNPs, EQTNPs, and LBANPs.
B.5 NOT-DIAGONA VARIANTS
In this section, we consider Not-Diagonal variants Neural Processes. Unfortunately the self- attention mechanism used in Not-Diagonal models (such TNP-ND LBANP-ND) require quadratic computation with respect to the number of target datapoints when querying. As a result, although LBANP-ND is efficient in terms of the number of context datapoints, the model would not be scalable in terms of the large number of target datapoints. Instead, we propose a more efficient variant called LBANP-END (Efficient Not Diagonal), Instead of a self-attention mechanism which has a quadratic computational cost, LBANP-END proposes latent bottlenecked attention mechanism that is linear in terms of computation. Specifically, the predictor information from the target datapoint embeddings via a cross-attention mechanism to fixed set of latent vectors LQEMB (Latent Query Embeddings). The initialization this set of latent vectors LQEMB, is meta-learned. We apply a self-attention mechanism to compute the higher-order information in latent space. After the higher-order information has been computed, we retrieve the information from the latent vectors via another cross- aoss-attention mechanism. Specifically, this is performed follows:
LQEMB; Crr.ttttnttiitooCoosostttttooino
QEMBx))
Output Predictor PredictorCCrossAttenntn (CossAAteenttionu(EEEEB, LQEMB,))
where Q is hyperparameter representing the number of layers.
We evaluate the performance of the non-diagonal variants the methods on the image completion task. Table 7 shows that LBANP-ND performs competitively with TNP-ND across both the CelebA and EMNIST datasets. In the table, we include TNP-D and LBANP (128) as a reference of the performance of the vanilla variants. In the results, see that both TNP-ND and LBANP-ND are incapable of scaling to CelebA64. In contrast, we see that, our proposed efficient non-diagonal more scalable than LBANP-ND and TNP-ND, and achieves tfate-of-the-n results on both the 32 x 32 and 64 x 64. Finally, similar to previous experiments, we see that scaling the number of latent vectors improves the performance of these variants as well.
variant, LBANP-ENS
16

---

Published as a conference paper at ICLR 2023
CelebA
EMNIST
Method
32x32
64x64
Seen (0-9)
Unseen (10-46)
TNP-D
3.89 = 0.01
#..01
1.46 0.01
LBANP (128)
3.97 # 0.02
5.09
0.02
1.39 0.01
1.17 0.01
TNP-ND
7.0820002
0.00
+.1000000
LBANP-ND (8)
0.01
1.38 0.01
1.06
LBANP-ND (128)
0.03
1.42
1.04
LBANP-END (8)
0.04
5.36 + 0.03
1.39 0.01
1.06
LBANP-END (128)
5.60 # 0.02
6.06 0.03
1.42 0.02
1.14
Table 7: Image Completion Experiments with Not Diagonal variants in terms of log-likelihood.
17


Reviews:
[{"review_id": "los-4aq8a8R", "reviewer": "Reviewer_vZAK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new member of neural processes. The aim is to improve the attention efficiency of transformer neural process where the attention strategy requires quadratic computation with respect to the number of context data points.  The idea is from iterative attention (Jaegle et al., 2021b) which changes the self-attentions between context data points to a cross-attention with latent variable and a self-attention between latent variable. When the size of the latent variable is significantly smaller than the one of context size, such an exchange could greatly reduce the total computational cost. ", "summary_of_the_review": "The paper is well written and easy to follow. The motivation is clear and method is reasonable and effective. The main weakness is the in-depth analysis on the new iterative attention to the base neural process. please see the `Strength And Weaknesses' for more details. \n\n\n---- after rebuttal \n\nThanks to the authors to clarify my concerns in the response. I do believe the proposed method could save a significant amount of computational time and memory according to the new results. Some other issues are resolved as well. I appreciate that. So I am happy to change my score from 5 to 6. The reason why I did not give 8 is the novelty reason, although I believe the proposed technique could improve the efficiency of TNP a lot and I recognize its value to the area, the technique itself is mainly from the existing work - iterative attention. I read the authors' response (to all reviewers) carefully but there is no explanation of the 'new' design and technique proposed compared with iterative attention. ", "strength_and_weaknesses": "Strength\n\n1. The motivation is clear and the proposed idea looks reasonable.\n2. There are extensive empirical evaluations and the performance looks good and comparable with transformer NP\n\nWeaknesses\n\n1. Considering the main contribution of this work is to reduce the computational cost, there should be some empirical evaluation on the time complexity in addition to table 6 so let the readers know how much time we can save from transformer NP when achieves similar results. \n\n2. The table 6 is a little unfair, because there is L in the new iterative attention. As shown in Figure 4, the size of L apparently heavily affects on the final performance. It would be better to include it in the time complexity. \n\n3. It is expected to see more discussions or analysis on the effect of this new attention comparing with previous one. For example, will it affect on the convergence rate? Will it break the process property of the neural processes because the proposed method introduce more `constraints' on the latent variables which are assumed to be independent before, like NP and CNP? like permutation invariant and marginal consistency? \n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FJ6pxYlKYl", "reviewer": "Reviewer_RFme", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "TNP uses GPT layers and shows good performance. But its memory complexity $O((N+M)^2)$ could be large if $N$ is large. The paper proposes a latent-bottleneck of $L$ vectors that mediate between the $N$ context points and the $M$ target points. By doing this, the authors are able to reduce the complexity to $O(NL + ML)$ where $L$ is the number of latent vectors and a human-specified hyperparameter.", "summary_of_the_review": "I believe in the promise and value of introducing the latent bottleneck and I tend to believe in the reduced big-O complexity. But in experiments, reporting the actual saving in memory should go hand-in-hand with the reporting of the performance. I am currently rating the paper as 6 in the hope that the authors will be able to add more justification. I am happy to revise my rating in the light of new information.", "strength_and_weaknesses": "### Pros\n\n1. The use of latent-bottleneck in the context of TNPs is novel and interesting.\n2. Shows better performance than pre-TNP baselines. Claims to show better memory complexity than TNP through a big-O analysis. Naturally, the performance is expected to be worse than TNP and this is empirically confirmed. But this performance drop can be traded-off with gains in memory saving shown via the big-O complexity.\n3. The comparison with EQTNP is useful as a natural variant of TNP. The fact that EQTNP performs a bit worse than TNP-D is new knowledge.\n\n### Weaknesses/Questions\n\n1. $L$ is treated as a pre-specified constant. However, if the tasks are indeed too complex/difficult such that they typically require a large $N$, then users are likely to choose a proportionately large  $L$ also to handle the underlying difficulty of the tasks. In other words, the task difficulty is a confounder that affects both $N$ and $L$ simultaneously. So, it might be tricky to say that the complexity is simply $O(N+M)$ while completely omitting $L$. It might be more justified to say $O((N+M)L)$?\n2. In experiments, the choice of $L$ should also be better justified given what distribution of $N$ values was used (perhaps reporting $L$ as a fraction of the mean value of $N$?). For instance, in the 1-D regression problem, $N$ takes values in $[3,47)$, and in this case, showing $L=128$ (which seems much larger than any $N$) seems to defeat the purpose of having latents? \n3. A way to address the above concerns might be to show the actual GPU memory consumption at test time as a function of $N$ and $L$ and compare LBANP’s memory savings relative to TNP. I am unsure what would be the best way though and I leave it up to the authors to decide the best way. Also, what was the memory cap that prevented running some experiments in Table 3?\n4. The interleaving seems to make the proposed model deviate from the standard GPT layers — something that, to me, was a big appeal of TNPs. Why is interleaving important in EQTNP and also in LBANP? Why not first do encoding of the context points with a standard transformer and then later do the prediction of target points with another standard transformer decoder? An ablation could be helpful. \n\n### Minor Comments/Questions\n\n1. Would be good to show a row for the computational complexity of EQTNP in Table 1 rather than only in Table 6.\n2. In Section 4.3, the Normal distribution may be shown with a calligraphic $\\mathcal{N}$.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "5MBpkue_QbO", "reviewer": "Reviewer_peD2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the computational bottleneck problem of Transformer Neural Processes (TNPs) that results from the quadratic complexity of transformers with respect to the input sequence length. The paper aims to combine the best of the worlds of state-of-the-art expressive power at the expense of high compute cost (TNPs) and scalable computation at the expense of suboptimal performance (other NP variants). The paper proposes a modification to the TNP pipeline to trade these competing goals using a fixed-sized attentive memory. The main idea is to embed the input into a fixed-dimensional latent space and model high-order interactions among data points as cross-attention to this embedding as opposed to self attention in the native observation space.", "summary_of_the_review": "While being a decent piece of work with some degree of originality and some promising results, I find the scientific contribution of the paper too slim for a main-track ICLR paper. I am also not yet convinced about the significance of the studied problem (see first bullet item under Weaknesses).", "strength_and_weaknesses": "**Strengths:**\n\n * The paper reports a comprehensive set of experimental results, where the proposed method appears to indeed approach the TNP performance in reduced computational cost.\n* The proposed idea is rather simple, easy to grasp and straightforward to implement. It is an original, though trivial, combination of a number of known architectural attention structures.\n\n\n**Weaknesses:**\n\n * It is not clear to me why quadratic complexity on the context set is such a severe limitation. Is it not the common assumption of NPs that the context set is rather small? Similarly, I assume the target queries are iid when conditioned on the context. Then how critical is it in reality for a target batch to be big? Can one simply not process small query batches even in parallel?\n* The proposed novelty is very incremental and straightforward. It does indeed improve performance, but I would categorize the presented method rather as a useful implementation trick than as a solid scientific finding.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster