Dimension Agnostic Neural Processes
Authors: Hyungi Lee
Chaeyun Jang
Dong Bok Lee
Juho Lee
Venue: ICLR 2025 Poster (2025)
Keywords: Neural Processes
varying dimension
regression
bayesian optimization
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: A new Neural Process model that can handle varying dimensional inputs and outputs.

Abstract: Meta-learning aims to train models that can generalize to new tasks with limited labeled data by extracting shared features across diverse task datasets. Additionally, it accounts for prediction uncertainty during both training and evaluation, a concept known as uncertainty-aware meta-learning. Neural Process (NP) is a well-known uncertainty-aware meta-learning method that constructs implicit stochastic processes using parametric neural networks, enabling rapid adaptation to new tasks. However, existing NP methods face challenges in accommodating diverse input dimensions and learned features, limiting their broad applicability across regression tasks. To address these limitations and advance the utility of NP models as general regressors, we introduce Dimension Agnostic Neural Process (DANP). DANP incorporates Dimension Aggregator Block (DAB) to transform input features into a fixed-dimensional space, enhancing the model's ability to handle diverse datasets. Furthermore, leveraging the Transformer architecture and latent encoding layers, DANP learns a wider range of features that are generalizable across various tasks. Through comprehensive experimentation on various synthetic and practical regression tasks, we empirically show that DANP outperforms previous NP variations, showcasing its effectiveness in overcoming the limitations of traditional NP models and its potential for broader applicability in diverse regression scenarios.

Published as a conference paper at ICLR 2025
DIMENSION AGNOSTIC NEURAL PROCESSES
Hyungi Lee, Chaeyun Jang, Dong bok Lee, Juho Lee KAIST [1hk2708, joy9911, markhi, jubolee)jkaist.acckk
ABSTRACT
Meta-learning aims to train models that can generalize to new tasks with limited labeled data by extracting shared features across diverse task datasets. Additionally, it accounts for prediction uncertainty during both training and evaluation, concept known as uncertainty-avve aware meta-learning. Neural Process (NP) is a well-known uncertainty-awant meta-learning method that constructs implicit stochastic pro- cesses using parametric neural networks, enabling rapid adaptation to new tasks. However, existing NP methods face challenges in accommodating diverse input dimensions and learned features, limiting their broad applicability across regression tasks. To address these limitations and advance the utility of NP models general regressors, we introduce Dimension Agnostic Neural Process (DANP). DANP in- corporates Dimension Aggregator Block (DAB) transform input features into fixed-dimensional space, enhancing the model's ' ability to handle diverse datasets. Furthermore, leveraging the Transformer architecture and latent encoding layers, DANP learns a wider range of features that are generalizable across various tasks. Through comprehensive experimentation on various synthetic and practical regres- sion tasks, we empirically show that DANP outperforms previous NP variations, showcasing its effectiveness overcoming the limitations of traditional NP models and its potential for broader applicability in diverse regression scenarios.
INTRODUCTION
In real-world datasets, there are many tasks that come with various configurations (such as input feature dimensions, quantity of training data points, correlation between training and validation data, etc.). However, each task has a limited number of data points available, making it difficult to train model capable of robust generalization solely based on the provided training data. To tackle this issue, meta-learning aims to train a model capable generalizing to new tasks with few labeled data by learning generally shared features from diverse training task datasets. In cases of limited labeled data for new target tasks, ensuring model trustworthiness involves accurately quantifying prediction uncertainty. which is critical as achieving precise predictions. A meta-learning strategy that considers prediction uncertainty during training and evaluation is known as uncertainty-awan meta-learning (Nguyen & Grover, 2022; Almecija et al., 2022).
One of the well-known incertainty-awar meta-learning methods is Neural Process (NP) (Garnelo al., 2018a;b). NP employs meta-learning understand the data-generation process governing the relationship between input-output pairs in meta-training and meta-validation data. Unlike the traditional approach to learning stochastic processes, where model selection from known class, e.g. Gaussian Processes (GPs), precedes computing predictive distributions based on training data, NP constructs an implicit stochastic process using parametric neural networks trained on meta-training data. It then optimizes parameters to maximize the predictive likelihood for both the meta-train and meta-validate validation data. Consequently, when NP effectively learns the Lata-generation process solely from data, it can quickly identify suitable stochastic processes for new tasks. Thus, NP can be viewed as a data-driven uncertainty-awe aware meta-learning method for defining stochastic processes.
However, previous works (Gordon et al., 2020; Foong et al., 2020; Lee et al., 2020; Nguyen & Grover, 2022; Lee et al., 2023) in NP literature lack two crucial attributes essential for broad applicability across different regression tasks: 1) the ability to directly accommodate diverse input and output dimensions, and 2) the adaptability of learned features for fine-tuning on new tasks that exhibit varying input and output dimensions. Due to the absence of these properties, it is necessary

---

Published as a conference paper at ICLR 2025
to train each NP model separately for different dimensional tasks. These limitations hinder the utility of NP models as general regressors across diverse datasets compared to traditional stochastic processes (Lee et al., 2021). Traditional stochastic processes naturally accommodate varying input dimensions, particularly in regression tasks involving high-dimensional input features with limited training data, such hyperparameter optimization tasks.
To tackle these limitations and advance the utilization of NP models as general regressors for a wide range of regression tasks, we introduce a novel extension of NP called Dimension Agnostic Neural Process (DANP). In DANP, we propose a module called Dimension Aggregator Block (DAB), which transforms imput features of varying dimensions into a sxedd-dimensionl representation space. This allows subsequent NP modules to effectively handle diverse datasets and generate predictive density for the meta-validation data. We also add the Transformer architecture (Vaswani al., 2017) alongside latent encoding layers based on the architecture of Transformer Neural Processes (TNP) (Nguyen Grover, 2022) to enhance the model's ability to learn wider range of features and effectively capture functional uncertainty, which can be applied across different tasks. Through experimentation on a variety of synthetic and real-world regression tasks with various situations, show that DANP achieves notably superior predictive performance compared to previous NP variations.
2 BACKGROUND
2.1 PROBLEM SETTINGS
Let X' be an input space defined Usen X, with each X; S R$ for all i E N. Similarly, let Usen 3 represent the output space, where each 3/ RR for all E N. Let T {Tyyben be task set drawn in i.i.d. fashion from task distribution Prack (7) Given two dimension mapping functions u, v - N. each task Tj comprises a dataset Dj = {d where dj,x = (Xj.n. E Xu(j XVvO) represents an input-output data pair, along with an index set Cj S [ny] where [m]] == {I....,.0 ........ m} for all m N. We assume elements in D, are i.i.d. conditioned
k=
some function fz- Here, the set of indices Cj defines the context set Dj.e := [dj, )kec, Similarly, the target set is defined as Dj,t (djA)kEt, where tj :- [ny] cj. We aim meta-learn collection of random functions f Xu(j) + Ve(s)) where each function within this set effectively captures and explains the connection between imput and output y pairs. For any given veta-training task Ty, we can regard its context set Dj.c as the meta- veta-training set and its target set Dj, as the meta- meta-validation set.
2.2 NEURAL PROCESSES
For the previous NP variants, their objective was to meta-learn a set of random functions fj Xá, Va... for some fixed dm. dout N, which is equal to the situation where dimension mapping functions u, v are constant functions, i.e., uuj) = du and v(z) = dout for all j E N. this context, to select suitable random function fj for the task Tj. NPS meta-learns how map the context set Dj.e to a random function fj that effectively represents both the context set Dj.e and the target set Dj.-- This entails maximizing the likelihood for both meta-training and validaiiint datasets within an uncertainty-awe aware meta-training framework. The process involves learning predictive density that maximizes the likelihood using the following equation:
X, Djcc) / p p(Y, X,
II plyanff,, x,,k) p(fjIDg.ddf kk[ng)
(I)
where X, ( and Y, = (y In line with discussion in Section 2.1, make the assumption that given the random function fz. the outputs collection Y, are i.i.d. Employing the Gaussian likelihood and parameterizing f3 with latent variable E Rdd, Eq. 1 reduces to,
/[
II N (yanler, (x3ul) (x.x), diag(o?, (xxx))) ( p(r,)D,,,ddr,,
XX, Djo))
(2)
kE(P,)
where M4y Xxa Vau and oz Xa. R Then different NP variants aim ef- fectively design the model structures of the encoder, denoted Senc, and the decoder, de-
2

---

Published as a conference paper at ICLR 2025
TNP
DANP
T
T
Misked Transformer Layers
Q
D
Zero Raddeed
Yorgert
Masked Transformer Layers
Decoder
Zero Packeed Vitagpe
DAB
MLP Decoder
Tranaformer
Self Actention
Layers
Concat
Latent path
Figure 1: Model comparison between TNP and DANP. While TNP (Nguyen & Grover, 2022) solely employs a deterministic pathway with Masked Transformer layers, DANP incorporates both DAB and an extra latent pathway alongside Transformer layers and Self-Attention layer.
noted as fasc. These components are responsible for describing the distributions p(r, (D,,) and N 1ysAlyy (xx,)), diag(o?, (x), respectively.
NP variations can be roughly categorized into classes based their approach modeling p(ry|Pj,c): 1) Conditional Neural Processes (CNPs) (Garnelo et al., 2018a; Gordon et al., 2020; Nguyen & Grover, 2022) and 2) (latent) NPS (Garnelo al., 2018b; Foong al., 2020; Lee al., 2023). CNPS establish a deterministic function, called deterministic path, from Dj.e to and represent p(ry p.,e)
as a discrete point measure, expressed as:
p(r,)Dj,c) òr, (r), Fy f (D
(3)
where @ is the parameter of sene- In contrast, NPs address functional uncertainty model uncertainty in modeling p(ry (D),)) Typically, they employ variational posterior ((ry/D),.), called latent path, to approximate p(r; (Dj.a) for any subset Dj.+ Dj, defined as:
q(r; j.,) = (N(rr (mp, di.g(s (mD,... s.,,)) = Senc (Djss Ø).
Then both of the classes decode the mean and variance of the input j as follows:
(Mr,(Xj,) o (x,,)) = (x rys
where Saee is another feedforward neural net vh.
Training CNPs involves maximizing the average predictive log-likelihood across meta-training tasks Tj. i.c log p( Y (X Dy,o]]. On the other hand, are typically trained by maximizing the Evidence Lower BOund (ELBO), which is expressed as:
E,, [ogloop(Y,,,,,, P5c) 2 Er,
Eg(ryID,) [log 1 KE(my)
(6)
where Ns is shorthand for NV Mylllry (x),)), diag(o?, (x)xa))).
There have been several attempts to enhance the flexibility of the encoder fene to improve the predictive performance (Gareeo et al., 2018a; Kim et al., 2018; Gordon et al., 2020; Nguyen & Grover, 2022). In this study, we adopt the state-of-the-art TNP model as our base structure, which leverages masked self-attention layers encoding layers and also belongs to the category of CNPS variants.
3 DIMENSION AGNOSTIC NEURAL PROCESS
As we mentioned in Section and Section 2.2, the limitation of the previous NP variants that they primarily handle scenarios where the input space and output space are confined to Xd, and Va for the fixed din dout E N. To address this constraint, we introduce a novel NP variant called DANP. Initially, we elucidate how the predictive density evolves when the dimension mapping functions u, are not constant in Section 3.1. Subsequently, we expound on how we can convert input features of varying dimensions into fixedddieession representation space utilizing the DAB module in Section 3.2. Finally, we detail the strategies employed to augment the model's capacity for learning diverse features in Section 3.3.
3

---

Published as a conference paper at ICLR 2025
3.1 A MODEL OF PREDICTIVE DENSITY WITH VARYING DIMENSIONS
Given our assumption that all the tasks may have varying input and output dimensions, write Dj {d - where Xj,k
[ E R²()) and j == [y ... j,k (o(1) RT)) (x), ... j.k
Given the context Dj,e, the equation for the predictive density p,,, D.e)) remains the same with Eq. However, due to the varying dimensions, the computation of both the likelihood NJJx and the context representation posterior p(r,| poses challenge. In a fixed dimension setting, only the size of the context varies across different tasks, and this could be processed by choosing fene as cermutation-invarian set functions (Zaheer et al., 2017). However, in our scenario, for two different tasks T1 and T2. single encoder should compute,
Average Pooling
(l)) Aei Φ) = fene (eno((((( ([((())()))...... 1,k 1,k (yi..................... N1,k (()
Miker)s)
Self Attention
(7)
Positional Encoding
Senc Φ) = Seno (1((x2,)))...
x ((2) (Y2 y2,k )eezi)))
Linear Projection
(8)
The existing permutation-invane -inaariinnnn-nniiinnnnn encoder can process when leil A |c2/. it cannot handle when (u(1),v(1) z (u(2),((2)) because this disparity happens at the lowest level of the encoder, typically implemented with normal feed-forward neural network. The standard architecture in the previous NP models to employ Multi-Layer Perceptron (MLP) taking the concatenated inputs, for instance,
Figure The overview of DAB module. A DAB can encode and decode inputs and outputs of varying dimensions.
feno (Djcci0)
Σ MLP(concat(xx,A. kEog
(9)
Cyl
and the MLP encoder can only process fixed-dimensional inputs. similar challenge also applies to the decoder fuee computing the predictive mean Hry (Xz,x) and variance of, (x,,).. To address this challenge, we need new neural network that is capable of handling sets of varying dimensions.
3.2 DIMENSION AGGREGATOR BLOCK
To enable our model to process sets with elements of varying dimensions, we introduce a module called Dimension Aggregator Block (DAB). The module can encode inputs with varying feature dimensions into fixed-dimensional representation and varying dimensional representations, which will describe individually below. The overall architecture is depicted in Fig. 2.
Encoding x into a fixed dimensional representation. Consider an input (x,y) where x = ... ydv] E Rdy, To maintain sufficient information for each dimension of the input data, even after it has been mapped to a fixed-dimensional representation, we initially expand each of the d, + dy dimensions of the input data to d, dimensions using learnable
[x1 xd=] E Rdd and y [y_...
linear projection w Rd, as follows:
[ = w wcconeat(x, y)] E Rdd x(d++
(10)
When encoding point in the target set Dj, without the label y, simply encode the zero- eeto-padded value concat(x, 0). Next, following Vaswani et al. (2017), we incorporate cosine and sine positional encoding to distinguish and retain positional information for each input dimension follows:
PEX = sin(j/P(i)) PEX = cos(j/P(i)). PEY (26.1) = cos(l/P(i), PEY (21+1.1))) sin(l/P(i)).
(11)
(12)
where P(i) 20/4 and PEX, PEY represent the positional encoding for x and y respectively.
Here, [dm] and 1 E [d,] respectively denote the position indices of & and y, while i E L
4

---

Published as a conference paper at ICLR 2025
represents dimension in the representation space. Since the dimensions of both x and y vary and the representation must be divided between corresponding positions in x and y, e.g., x and y we use distinct positional embeddings, PEX for and PEY for y. After adding positional encoding to (), we further compute,
(x, y) SelfAttn(concat(* + (PEX, PEY)),
(13)
where SelfAttn indicates self-attention layer (Vaswani et al., 2017). Here, we can regard (x, y) before the self-attention layer as akin to context-free embeddings (Rong, 2014), because they remain unchanged when position and value are fixed. Conversely, (x. y) after the Self-Attention layer can be likened to contextual embeddings (Devlin et al., 2018), which is known as advanced embedding compared to context-free embedding, as the final representation may vary depending on alterations in value and position across other dimensions due to the interaction through the self-attention layer. Then, we employ average pooling for the X to integrate all the information across varying dimensions in into fixed-dimensions representation, i.c., x AvgPool (X) E R with AvgPool representing the average pooling operation across the feature dimension.
Handling variable number of outputs for y. The DAB should produce representations that can be used in the decoder later to produce outputs with varying dimensions. To achieve this, unlike for the x, we keep the sequence ý without average pooling. Instead, for each f = 1, ...... ... dy, we concatenate & and y' and put into the decoder to get the predictive mean and variances. The dimension of the encoding for y from the DAB would be the same as original dimension of y. Note that this is not like the sequential decoding in autoregressive models, and is possible because know the dimension of y before we actually have to decode the representation.
3.3 LEARNING MORE GENERAL FEATURE UTILIZING LATENT PATH
Let D, := (X x))))= k= be the representations obtained by DAB for the dataset Dj- The next step involves computing the predictive density using the encoder and decoder structure. Here, we employ TNP (Nguyen & Grover, 2022), a variant of CNPs, as our base model structure. In TNP, Masked Transformer layers are utilized the encoder for the deterministic path, while a simple MLP structure serves as the decoder. To improve the model's 's capacity to learn generally shared features across various tasks and effectively capture functional uncertainty, we introduce new latent path comprising Transformer layers and a Self-Attention layer alongside the single deterministic path encoder. In specific, we pass the entire Dj into Masked Transformer layers to make deterministic del as follows:
parameter )
j.k= concatt F/AA) E TR2dd.
(14)
det (IFFL (r() 1)),), E Ru())yx2dd
(15)
where MTFL denotes Masked Transformer layers with mask M₁, and concat( ( (()),.) ( } indicate concatenation operation which concatenate z),8 for all I E [u()]] and h E [mz]. In this context, for all 1, E [v(j)], the mask M, E RRo0), xu()), assigns a value of to the index (1, 12k:2) if both K1 and h2 are elements of cj. or if ky is in ty and k2 is in ej; otherwise, it assigns a value of O.
(())
For the latent path, we only pass context set Dr through Transformer layers, followed by one self-attention and MI. operation to determine the latent parameter riu as follows:
pla = (NgPoo(SSeAAttt((((Connnn ((z) 12 )kee)))),
(16)
Dr DD ) MLP(T!"),
(17)
Lat (( (r (D,,) = -N(1)" |o,, diag(s?
(18)
where TL denotes Transformer layers. Finally, we concatenate the deterministic parameter gdt and latent parameter r make the final parameter r, before forwarding them the decoder module. Then, by utilizing& a variational posterior with the latent path, our training objective transforms into
R., (olog((Y; Y, Djc)) Er,
Σ kE(mg]
BuGYDD, (ggg (I.(allag.))
(19)
5

---

to observe the Published as a conference paper at ICLR 2025
where 9 and lie denotes (Dr) and q(r)) (Dj.e), respectively. Refer to Fig. 1 contrast in the architecture between TNP and DANP.
4 RELATED WORKS
The first NPS model, called CNP (Garnelo et al., 201 8a), utilized straightforward MLP layers for both its encoder and decoder. Similarly, NP (Garnelo et al., 20188 2018b) adopted MLP layers but introduced a global latent variable to capture model uncertainty, marking an early attempt address uncertainty in NP frameworks. Conditional Attentive Neural Process (CANP) and Attentive Neural Process (ANP) (Kim et al., 2018) are notable for incorporating attention mechanisms within the encoder, enhancing the summarization of context information relevant to target points. Building on these ideas, TNP (Nguyen & Grover, 2022) employs masked transformer layers in its encoder, delivering state-of-the-arr performance among NPs across multiple tasks. Louizos et al. (2019) introduced variant that used local latent variables instead of global latent variable improve the model's s ability to capture uncertainty. Following this, Bootstrapping Attentive Neural Process (BANP) (Lee et al., 2020) proposed the residual bootstrap method (Efron, 1992), making NPs more robust to model misspecification. Lastly, Martingale Posterior Attentive Neural Process (MPANP) (Lee al., 2023) addressed model uncertainty with the martingale posterior (Fong al., 2021), offering modern alternative to traditional Bayesian inference methods. Refer to Appendix B to see a more detailed review of previous Neural Processes works.
Neural Processes
Neural Diffusion Process Similar to DANP, there are prior works (Liu al., 2020; Kossen et al., 2021; Dutordoir et al., 2023) that utilize bi-dimensional attention blocks to facilitate more informative data feature updates to ensure the permutation invariance property both at the data-instance level and the dimension level. Specifically, Neural Diffusion Process (NDP) (Dutordoir et al., 2023) employs bi-dimensional attention blocks to guarantee permutation invariance both the data and dimension levels, naturally leading to dimension-agnostir properties. However, NDP has a structural limitation in that it is only partially dimension-agnostie for when y 1, and is not dimension- agnostic for other combinations. This makes it difficult to use as a general regressor. Additionally, the use of diffusion-based sampling approximate the predictive distribution leads significantly high computational costs during inference and results in limited likelihood performance. Refer to Appendix D.1 to see the empirical comparison between DANP and NDP.
EXPERIMENTS
In this section, we carry out series of experiments to empirically showcase the efficacy of DANP across different situations, especially in various regression tasks and Bayesian Optimization task. To establish a robust experimental foundation, we employ five distinct variations of NP, encompassing state-of-the-art model: CANP, ANP, BANP, MPANP, and TNP. For a fair comparison, we maintain an identical latent sample size in the latent path across all models, except for deterministic models such as CANP and TNP. We marked the best performance value with boldfaced underline, and the second-best value with underline in each column in all tables. All the performance metrics are averaged over three different seeds and we report I-sigma error bars for all experiments. Refer to Appendix C for experimental details containing data description and model structures.
5.1 GP REGRESSION
To empirically verify the effectiveness of DANP, we initially conducted GP regression experiments under various conditions: From-scratch, Zero-shof, and Fine-tuning. In the From-scratch scenario, we compared DANP against other baselines using fixed input dimensional GP data for both training and testing. In the Zero-shot scenario, we demonstrated the ability of DANP to generalize different dimensional input GP data without direct training on that data. Lastly, in the Fine-tuning scenario, we conducted experiments where we fine-tuned on unseen dimensional GP data, using limited training data points, utilizing pre-trained DANP alongside other baseline models.
From-scratch To validate the capability of DANP to effectively learn generally shared features and capture functional uncertainty across tasks, we first compare DANP against other baseline models in the

---

Published as a conference paper at ICLR 2025
Zero-shot
Full Img
Context
Mean
Std
1.0
0.5
0.0
ground Brutn
-0.5
contect points
Figure 3: Posterior samples of DANP in (Left) the Zero-shot scenario with I-dimensional GP dataset and (Right) the Image completion task using the EMNIST and CelebA datasets. (Left) Black stars represent the context points and the dashed line indicates the ground truth for the target points. Each color represents different posterior samples generated from the latent path. (Right) Displays the full image, context points, predictive mean, and standard deviation of DANP for both the EMNIST and CelebA datasets. Outputs for both images are produced by a single model.
Table 1: Results of the context and target log- log-likelihood for the GP regression task in the From-scrated scenario. nD A in the first row denotes the n n-dimensional GP dataset using the A kernel.
ID RBP
ID Matern
2D RBP
2D Moter
Model
context
target
crntest
targer
contest
target
crntest
target
CANP
1.372 200000
0,839
+:0.002
1.377
1-0000
0.663
=:00007
1.377
20.001
0.165
2:00015
1.373
200.011
-0.066
=00.007
ANP
1.322 1.0000
0.855
1:0.004
1,377
1.0000
0.681
1100003
1.378
10.000
0.170
1:00044
1.346
1.00035
-0.107
1.0.006
BANP
1377 10.000
0.861
100001
1.317
1.0.000
0.689
1.0004
L378
10.000
0.228
10.004
1.378
10.030
-0.033
1.0.013
MPANP
1.376 100000
0.856
1:00006
1.376
1.0 000
0.679
1.0005
1.378
100001
0.242
10.001
1.376
1.0 00.002
-0.029
1.0.007
TNP
1.381
2.0000
0.904
1:0.003
1,381
2:0 000
0.710
100000
1,383
10.000
0.362
10.001
1,383
1.0.000
0.060
200002
DANP Cours)
1.381 +0.000
0.921
+:0.003
1.382
++0.000
0.723 =0.000
1.383 +0.000
0.373
+0.001
1,383 +0.000
0.068 =-0.001
From-scratch scenario in diverse fixed dimensional GP regression tasks. In this experiment, the meta- training datasets are produced using GP under four distinct configurations: either one-dimensional or two-dimensional input, utilizing either the RBF or Matern kernels. The results presented in Table 1 demonstrate that DANP consistently surpasses other baseline models across various settings, particularly excelling on the target dataset in terms of log- likelihood. These results prove that DANP effectively grasps common features and captures functional uncertainty, outperforming other baseline models even with various settings in fixed-dimensional GP regression tasks.
Zero-shot In the Zero-shot scenario, we train a single NP model using various dimensional GP datasets and then evaluate the performance of the model on a range of GP datasets with different dimensions. Specifically, and 4-dimensional GP datasets, and another where the training datasets include 2, 3, and 4-dimensional GP datasets. After training, we assess our pre-trained model on GP datasets ranging from to 5 and 7 dimensions. We validate results for our model as DANP, in distinction from other baselines, is capable of simultaneously training on and inferring from datasets with diverse dimensions. Table 2a demonstrates that DANP successfully learns the shared features across different dimensional GP datasets and generalizes effectively test datasets with the same dimensions the training datasets. Remarkably, DANP also generalizes well to test datasets with previously unseen dimensions. For instance, the shot log- likelihood for the 1-dimensional GP dataset, when DANP is trained on and 4-dimensional datasets, is nearly comparable to the log-likelihood of CANP with From-scratch training in Table These findings suggest that DANP efficiently captures and learns general features across various tasks, allowing it to explain tasks with unseen dimensions without additional training. For further results using 2 and 3-dimensional GP datasets or different kernels during training, see Appendix D. The trends are consistent, showing that DANP generalizes well across various tasks. Refer to Fig. to see the zero-shot posterior samples for the 1-dimensional GP regression task. And also refer to Appendix D.2.3 see the results the additional zero-shot scenarios, especially extrapolation scenarios.
consider two different cases: one where the training datasets include 2
Fine-tuning In the fine-tuning scenario, we fine-tuned pre-trained NP models on a limited set of 160 I-dimensional GP regression tasks. For the baselines, we used pre-trained models that were trained on 2-dimensional tasks as described in the From-scratch experiments. For DANP, we used models pre-trained on 2, 3, and 4-dimensional tasks mentioned in the Zero-shot experiments. In Table 2b, 'Full fine-tuning' refers to the process where all pre-trained neural network parameters adjusted during fine-tuning, while FFreeze fine-tuning' means that the shared parameters in the
7

---

Published as a conference paper at ICLR 2025
Table 2:
Log-likelihood results for the GP regression task in (a) the Zero-shot and (b) the Fine-
tuning scenarios using RBF kernel. For (a), nD in the first column denotes the outcomes for the n-dimensional GP dataset. The colored cell indicates the data dimension used to pre-train DANP.
(a) shot scenario
(b) Frinetuning scenario
DANY trained 2D
DANI trained 00
Full fic-tuning
line-ianing
Dimension
Methad
bargs
ctrlect
bagg
1D RBF
1,336
=0,047
0.806
10.048
1,366
20,004
0.826
10.018
CANP
-0.305 20.045
-0,455 20000s
-0.061
200 66.236
-0.386 20.112
2D RBF
3D KBF
1,333
1,352
1,377
1.79
00.01
-0,007
10,0000
00.00
-0.699
-0589
00.300
0.30
750001
100007
3500066
15.063
1,38
1,383
1,383
1,383
20,000
20.053
20.000
20.000
-0.66
-0.56
-0.261
0:35
10011
113001
1:3042
1:3025
TNP
ANP
MPANP
BANP
-0.254 -0.339
-0.042
-0.273 =0.121
0.292 20.004
=0.016
-0.414
-0.448
.03665 200,033
0.339
200
+02255
=:0226
-0481
0.131
-0.311
+0.372 0.557
26 230,034
20 129
10.032
-0. 0..63 =0.020
-0.087 =0.205
-0.369 20.007
0,199 200.211
TD RBF
1,348
20,016
0.726
2.0.26
1,355
30.022
-0.723
10.022
DANF(OuES)
1.376 20,000
0 4:0004
1,376 360.00
0,809 200005
Table 3: Log-likelihood results for context and target values were obtained for (a) image completion tasks using the EMNIST and CelebA datasets, and (b) fine-tuning on video completion tasks. For (a),
DANP was trained concurrently on both the EMNIST and CelebA datasets. For (b), indicates the zero-shot performance of DANP.
(a) Image completion
(b) Fine- Fine-tuning on CelebA video data
Model
EMNIST
CelebA
Model
context
target
context
larget
context
target
CANP
-1.013 10.1
-1.053
00.076
CANP
L378 10.001
0.87
00.00
4.129
10.004
1,495
10.004
ANP
-0.498 -0.143
-0.517
+0.128
ANP
1,372 1:1000
0383
=0.011
4,131
+0,05
1.93
20.016
BANP
-0.037 00.334
-0.099
20.273
BANP
1,333 10.001
0,901
_0.00
4,127
10.005
2.292
10.021
MAPNP
-1.341 -0.132
-1,336
10.
MAPNP
1,365 100000
-0.05
4.127
10.004
1,505 20.011
TNP
-1.574 10.471
2,747
L0.501
TNP
1377 1.0.001
0,945 00.004
1.400
10.005
1.630
00.
DANP' (ours)
4.086
20.036
0.503
1:1 2:00603 063
DANP (ours)
1,382 1.0.011
0.969 _0,002
4.149
10.000
2,077
10.008
DANP (ours)
4.094
=0.011
0,560
+0.086
encoder layers remain unchanged during the fine-tuning process. The results in Table 2b clearly show that all the NP models, except for DANP, fail to achieve high generalization performance. Furthermore, the performance of DANP shows clear improvement over the zero-shot log- likelihhod result in Table 2a following the fine-tuning with the limited data. This indicates that the features from the pre-trained baselines are not effectively applied to unseen dimensional downstream datasets with a limited amount of data. In contrast, DANP is able generalize well on these unseen dimensional downstream datasets with only a small amount of downstream data. Refer to Appendix D.2.4 to see the results on additional fine-tuning scenarios.
5.2 IMAGE COMPLETION AND VIDEO COMPLETION
Image Completion To validate model's capability to meta-train implicit stochastic processes for varying output dimensions, we perform image completion tasks on two different datasets: EM- NIST (Cohen et al., 2017) and CelebA (Liu et al., 2015). In these tasks, we randomly select some pixels as context points and use NP models to predict the selected target pixels. Here, we use the 2-dimensional position value as input and the channel value as output. Previous NP models were unable to formulate the predictive distribution for varying output dimensions, failing to learn image completion tasks with different numbers of channels. However, our DANP model can handle varying output dimensions, allowing it to simultancously learn various image completion tasks with different numbers of channels. The experimental results for EMNIST and CelebA reported in Table were measured using models trained separately for each dataset for the baselines, whereas ours were ob- tained using a single model trained simultaneously for both datasets. Table demonstrates that DANP successfully learns both image completion tasks, validating its ability to formulate the predictive density for outputs with varying dimensions. Refer to Fig. 3 and Appendix D for the visualizations of predicted mean and standard deviation of completed images.
Fine-tuning on Video Completion To further validate the capability of utilizing pre-trained features in DANP for tasks with unseen dimensions, we created a simple video dataset based on the CelebA dataset. Specifically, we used the original CelebA data the first frame at time t 0 and gradually decreased the brightness by subtracting from each channel for each time t E [9]. This process resulted in an input dimension of 3, combining the time axis with position values. We fine- tunet pre- trained NP models on only 5 video data, simulating a scenario with limited data for the downstream task. For the baselines, used models pre-trained on the CelebA dataset. For DANP, used
8

---

Published as a conference paper at ICLR 2025
dim BO (Inter.
dim BD (Extra.)
10 dim BD (Extra.)
16 dim BD (Extra.)
012
2.00
0.3
025
2.51
0.11
013
0013
0.52
201
220
BD
ANP
CANP
BANP
MPANP
TNP
DANP
Figure 4: Results for Bayesian Optimization (BO) on 6-, 9-, 10-, and 16-dimensional hyperparameter tuning tasks in the HPO--B benchmark. Note that DANP is pre-trate pre-trained on 2-, 3-, 8-dimensional tasks. Table 4: Ablation results for ID, 2D GP regression and image completion tasks. In this table, we report the log-likelihood results for only the target dataset, excluding the context dataset.
Model
ID RBF
ID Matern
2D RBF
2D Matern
EMNIST
CelebA
TNP
0.000
0.710 : 0.001
0,362 = 0.001
0.060 :: 0.002
0.945 a 0.004
1.632 a 0.005
+ DAB
0,907 0.0001
0.713 0.001
0,365 0.001
0.061 = 0.000
0,949 = 0.004
1.645 0.014
Latent
0.923=0.00m
0.722 0.000
0.371 = 0.001
0.064 = 0.001
0.967 0.000
1.933
0.023
DANP
0,921 = 0.003
0.723 a 0.003
0.373 a 0.001
0.068 a 0.001
0.969 a 0.002
2,027 a: 0.006
Pos
0.922 = 0.002
0.724 # 0.001
-0.395 = 0.022
-0 -0.446 + 0.006
0.376 # 0.012
0.631 # 0.030
PMA
0.921 0.000
0.721 0.001
0.372 0.004
0.067 0.002
0.975 0.007
2.025 + 0.007
models pre-trained on both the EMNIST and CelebA datasets. Table 3b demonstrates that, while other baseline methods fail generalize the increased input dimensional dataset, our method successfully learns and generalizes well with scarce amount of training data. Refer to Appendix D for the example of video data and the predicted mean and standard deviation.
5.3 BAYESIAN OPTIMIZATION FOR HYPERPARAMETE TUNING
To illustrate the real-world applicability DANP, we conducted BO (Brochu et al., 2010) experiments on 6, 9, 10, and 16-dimensional hyperparameter tuning tasks in HPO-B Ffieda-AAange al., 2021) benchmark. HPO-B is a large-scale black-box hyperparameter optimization benchmark, which is assembled and preprocessed from the OpenML. repository with 8 different hyperparameter dimensions (2, 3,6.6, 6, 8, 9, 10, 16, and 18-dim) and evaluated sparsely on 196 datasets with a total of 6.4 million hyperparameter evaluations. We target 6, 9, 10, and 16-dimensional hyperparameter tuning tasks. To do so, we pre-trained baselines on the 2-dimensional tasks in the meta-train split of the HPO-B benchmark and then fine-tune the baselines on a limited of meta-batches sampled from each target task. In contrast, for DANP, we follow the zero-shot setting, where we use single model pre-trained on 2, 3, and 8-dimensional tasks in the meta-train split without fine-tuning on target tasks. Please see more detailed setups in Appendix C.7. We use Expected Improvement (Jones et al., 1998) as an acquisition function for all experiments and measured performance using normalized regret, -V where ymax and Ymin denotes the global best and worst value, respectively. We run 200 you Dyois iterations for all the BO experiments and report the average and standard deviation of normalized regrets over 10 different random seeds. The results in Fig. demonstrate that DANP outperforms other baselines in terms of regret with the same iteration numbers. This demonstrates that DANP is capable of serving as a surrogate model for different BO-based hyperparameter tuning tasks using only a single model without additional training on new BO tasks, and it also effectively learns generalized shared features across a wide range of tasks. Surprisingly, the gap between DANP and baselines is even larger for the dimension extrapolation settings (9, 10, 16-dim), which empirically validates that the DANP is capable of handling unseen, varying- loyynn-diiensionall data, including cases where extrapolation required. Refer to Appendix 1.10 to see the results on synthetic BO tasks. The trends are similar.
5.4 ABLATION STUDY
To analyze the roles of each module in DANP through ablation experiments, we conducted series of ablation experiments main parts: 1) the roles of the DAB module and Latent path, 2) the role of positional encoding in the DAB module, and 3) experiments replacing mean pooling with attention-based averaging in the
various modules. The ablation experiments were categorized into three

---

Published as a conference paper at ICLR 2025
DAB module. The experiments were conducted on 1D GP regression, 2D GP regression, and image completion tasks. In Table 4, we analyze the log-likelihood results only for the target, excluding the context. For the full results, including the context, please refer to Appendix D.3. It can be observed that the context exhibits similar trends to the target.
The Role of the DAB Module and Latent Path As mentioned in Section 3.2 and Section 3.3, the DAB module is used as module for handling varying dimensional input, and the latent path is added to capture more accurate model uncertainty, thereby improving model performance and increasing the capacity to learn generally shared features among tasks. Table 4 show that the performance trends align well with the direction we described and intended in the Section 3. In Table 4, the DAB and Latent rows show the performance when DAB and Latent paths are added to TNP, respectively. We can observe that in all data and experiments, adding only the DAB to TNP results in performance close to TNP, while adding the latent path results in performance closer to DANP. This demonstrates that adding only the DAB module allows for the handling of varying dimensional data, but there limitations in improving model performance. However, adding the latent path improves model performance but still has the issue of not being able to handle varying dimensional data.
The Role of Positional Encoding in the DAB Module When treating the diverse dimensional tasks, permuting the orders of the features should not affect the result, but note that the permutation should apply simultaneously for all inputs. For instance, for a model having three features, say we permute the features to (3,1,2) for the first input and (1,3,2) for the second input. Then there is no way for the model to distinguish different feature values. Removing positional embeddings from the DAB is effectively doing this; since treat all the features as elements in set, it allows different permutations applied for different inputs, so the model does not distinguish different features.
We've tested the necessity of positional encoding through additional experiments, which confirmed its importance. In Table 4, "Pos" indicates the case when we extract the positional encoding from the DAB module. For the 1D GP regression tasks, because there is only one dimension for the input æ. the existence of positional encoding does not affect the final performance. However, as seen in 2D regression tasks and image completion tasks, the existence of positional encoding is crucial for the final performance, Refer to Appendices D.5 to D.7 to see additional ablation results using Rotary Position Embedding (RoPE; Touvron al., 2023).
Attention-based averaging To verify if the mean pooling in the DAB module can be enhanced by using attention -based averaging, we employed the Pooling by Multihead Attention (PMA; Lee et al., 2019) module. This module uses a learnable token that is updated through cross-attention layers by pooling the input tokens using attention. In Table 4, the PMA row shows the results when mean pooling in the DAB is replaced with the PMA module. The results consistently indicate that mean pooling and attention-based averaging yield similar performance across nearly all tasks. Refer to Appendix D to extensive additional empirical analysis and additional experiments.
CONCLUSION
In this paper, we present novel NP variant that addresses the limitations of previous NP variants by incorporating a DAB block and a Transformer-based latent path. Our approach offers two key advantages: 1) the ability to directly handle diverse input and output dimensions, and 2) the capacity for learned features to be fine-tuned on new tasks with varying input and output dimensions. We empirically validate DANP across various tasks and scenarios, consistently demonstrating superior performance compared to the baselines. Conducting various experiments only with a single model can be a starting point for the utilization of NP models as general regressors for a wide range of regression tasks.
Limitation and Future work In this study, DANP concentrated on the regression task, but it can naturally be extended to other tasks, such as classification. A promising direction for future work would be to pre-train the encoder, which includes the DAB module to handle diverse dimensional data, using various datasets from different tasks and then fine-tuning with small amount of downstream data for various tasks using appropriate decoders.
10

---

Published as a conference paper at ICLR 2025
Reproducibility Statement. We provide the architecture of proposed model along with the architectures of other baseline models in Appendix C.1. Additionally, the hyperparameters used in the experiments, metrics, and detailed information about the data utilized in each experiment are described in Appendix C. Furthermore, additional experimental results, including ablation experiments and additional visualizations, presented in Appendix D.
Ethics Statement. Our research introduces new NP variants that calculate predictive density for context and target points tasks with varying inputs. is unlikely that our work will have any positive or negative societal impacts. Also, we utilize openly accessible standard evaluation metrics and datasets. Furthermore, we refrain from publishing any novel datasets or models that may pose a potential risk of misuse.
Acknowledgements This work was partly supported by Institute of Information & communi- cations Technology Planning & Evaluation(IITP) grant funded by the Korea (MV)) (No.RS-2019 11990775, Artificial Intelligence Graduate School Program(KAIST) Institute of In- formation & communications Technology Planning & Evaluation(IITP) grant funded by the Korea govermmen((MSTT) N0oRS-2002-00500277 Global AI Frontier Lab), and Institute of Information & communications Technology Planning & Evaluation(IITPP) grant funded by the Korea govern- ment(MSIT) (N..RS-0022-12207733 Meta-learning Applicable to Real-world Problems).
REFERENCES
Cesar Almecija, Apoorva Sharma, and Navid Azizan. Uncertainty-aware meta-learning for multi- modal task distributions. arXiv preprint arXiv:2210.0188// 2022.
Jason Ansel, Edward Yang, Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesensky,
Bin Bao, Peter Bell, David Berard, Evgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary Vetioo, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang, Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, CK Luk, Bert Maher, Yunjie Pan, Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi, Helen Suk, Michael Suo, Phil Tillet, Eikan Wang, Xiaodong Wang, William Wen, Shunting Zhang. Xu Zhao, Keren Zhou, Richard Zou, Ajit Mathews, Gregory Chanan, Peng Wu, and Soumith Chintala. PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation. In 29th ACM International Conference Architectural Support for Programming Languages and Operating Systems, Volume (ASPLOS '24). ACM, April 2024. doi: 0..1145/3620655.6660066 pptpoccttcoch2o pdf.
URL https: / http://yytttcccccoco
Matthew Ashman, Cristiana Diaconu, Junhyuck Kim, Lakee Sivaraya, Stratis Markou, James Re- queima, Wessel P Bruinsma, and Richard E Turner. Translation equivariant transformer neural processes. arXiv preprint arXiv:2406. 12409, 2024.
Thomas Back. Evolutionary algorithms in theory and practice: evolution strategies, evolutionary programming, genetic algorithms. Oxford university press, 1996.
Maximilian Balandat, Brian Karrer, Daniel R. Jiang, Samuel Daulton, Benjamin Letham, An- drew Gordon Wilson, and Eytan Bakshy. BoTorch: A Framework for Efficient Monte-Carlo Bayesian Optimization. In Advances in Neural Information Processing Systems 33, 2020. URL
http:///rxiz org/ abs / 111//991110 06403.
Eric Brochu, Vlad M Cora, and Nando De Freitas. A tutorial on bayesian optimization of expensive cost functions, with application to active user modeling and hierarchical reinforcement learning. arXiv preprint arXiv: 1012.2599 2010.
Wessel P Bruinsma, Stratis Markou, James Requiema, Andrew YK Foong, Tom R Andersson, Anna Vaughan, Anthony Buonomo, J Scott Hosking, and Richard E Turner. Autoregressive conditional neural processes. arXiv preprint arXiv:2303. 14468, 2023.
Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extending mnist to handwritten letters. 2017 international joint conference on neural networks (IJCNN), pp. 2921-2926. IEEF, 2017.
11

---

Published as a conference paper at ICLR 2025
Jacob Devlin, Ming-
Ming-Wei Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04805,, 2018.
Vincent Dutordoir, Alan Saul, Zoubin Ghahramani, and Fergus Simpson. Neural diffusion processes. In International Conference on Machine Learning, pp. 8990-9012. PMLR, 2023.
Bradley Efron, Bootstrap methods: another look at the jackknife. In Breakthroughs in statistics, pp- 569-593 Springer, 1992.
Lco Feng, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Latent bottle- necked attentive neural processes. arXiv preprint arXiv:22/ 1.08558, 2022a.
Leo Feng, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Efficient queries transformer neural processes. In Sixth Workshop on Leta-Learning at the Conference on Neural Information Processing Systems, 2022b.
Leo Feng, Frederick Tung, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Memory efficient neural processes via constant memory attention block. arXiv preprint arX1v:2305..4567, 2023.
Edwin Fong, Chris Holmes, and Stephen G Walker. Martingale posterior distributions. arXiv preprint 1oX.v:2003.1561,, 2021.
A. Y. K. Foong, W. P. Bruinsma, J. Gordon, Y. Dubois, J. Requeima, and R. E. Turner. Meta-learning stationary stochastic process prediction with convolutional neural processes. Advances in Neural Information Processing Systems 33 (NeurIPS 2020), 2020.
Jacob R Gardner, Geoff Pleiss, David Bindel, Kilian Q Weinberger, and Andrew Gordon Wilson Gpytorch: Blackbox matrix-matrix gaussian process inference with gpu acceleration. In Advances in Neural Information Processing Systems, 2018.
M. Garnelo, D. Rosenbaum, C. Maddison, Ramalho, D. Saxton, M. Shanahan, Y. W. Teh, D. J. Rezende, and S. M. A. Eslami. Conditional neural processes. In Proceedings of The 35th International Conference on Machine Learning (ICML 2018), 2018a.
M. Garnelo, J. Schwarz, D. Rosenbaum, F. Viola, D. J. Rezende, S. M. A. Eslami, and Y. W. Teh. Neural processes. ICML Workshop Theoretical Foundations and Applications of Deep Generative Models, 2018b.
J. Gordon, W. P. Bruinsma, A. Y. K. Foong, J. Requeima, Y. Dubois, and R. E. Turner. Convolutional conditional neural processes. In International Conference on Learning Representations (ICLR), 2020.
Alistair EW Johnson, Tom J Pollard, Lu Shen, Li-wei H Lehman, Mengling Feng, Mohammad Ghassemi, Benjamin Moody, Peter Szolovits, Lco Anthony Celi, and Roger G Mark. Mimic-iii, a freely accessible critical care database. Scientific data, 3(1):1-9, 2016.
Donald R Jones, Matthias Schonlau, and William J Welch. Efficient global optimization of expensive black-box functions. Journal of Global optimization, 13:455-492, 1998.
H. Kim, A. Mnih, J. Schwarz, M. Garnelo, S. M. A. Eslami, D. Rosenbaum, and V. Oriol. Attentive neural processes. In International Conference on Learning Representations (ICLR), 2018
Jungtaek Kim and Seungjin Choi. BayesO: A Bayesian optimization framework in Python. Journal of Open Source Software, 8390),5320, 2023.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations (ICLR), 2015.
Jannik Kossen, Neil Band, Clare Lyle, Aidan N Gomez, Thomas Rainforth, and Yarin Gal. Self- attention between datapoints: Going beyond individual input-output pairs in deep learning. Ad- vances in Neural Information Processing Systems, 34:28742-28756. 2021.
Alex Krizhevsky, Geoffrey Hinton, al. Learning multiple layers of features from tiny images, 2009.
12

---

Published as a conference paper at ICLR 2025
Tuan Anh Le, Hyunjik Kim, Marta Gareel,, Dan Rosenbaum, Jonathan Schwarz, and Yee Whye Teh. Empirical evaluation neural process objectives. In NeurIPS workshop on Bayesian Deep Learning, pp. 71, 2018.
Yann LeCun, Bernhard Boser, John S Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1(4):541-53) 1989.
Hyungi Lee, Eunggu Yun, Hongseok Yang, and Juho Lee. Scale mixtures of neural network gaussian processes. International Conference on Learning Representations (ICLR), 2021.
Hyungi Lee, Eunggu Yun, Giung Nam, Edwin Fong, and Juho Lee. Martingale posterior neural processes. In International Conference on Learning Representations (ICLR), 2023.
J. Lee, Lee, J. Kim, E. Yang, S. J. Hwang, and Y. W. Teh. Bootstrapping neural processes. In Advances in Neural Information Processing Systems 33 (NeurIPS 2020), 2020.
Juho Lee, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: framework for attention-based permutation-invarian neural networks. In Proceedings of The 36th International Conference on Machine Learning (ICML 2019), 2019.
Sulin Liu, Xingyuan Sun, Peter J Ramadge, and Ryan P Adams. Task-agnoste amortized inference of gaussian process hyperparameters. Advances in Neural Information Processing Systems, 33: 21 21440-21452, 2020.
Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
C. Louizos, X. Shi, K. Schutte, and M. Welling. The functional neural process. In Advances in Neural Information Processing Systems 32 (NeurIPS 2019), 2019.
Stratis Markou, James Requeima, Wessel P Bruinsma, Anna Vaughan, and Richard E Turner. Practical conditional neural processes via tractable dependent predictions. arXiv preprint arXiv:2203.08775, 2022.
Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-awarc meta learning via sequence modeling. In Proceedings The 38th International Conference on Machine Learning (ICML 2022), 2022.
Sebastian Pineda- Pineda-Arango, Hadi S. Jomaa, Martin Wistuba, and Josif Grabocka. HPO-B: A large-scale reproducible benchmark for black-box HPO based on openml. Neural Information Processing Systems (NeurIPS) Track on Dalasets and Benchmarks, 2021.
Ofir Press, Noah A Smith, and Mike Lewis. Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108. /2409, 2021.
Leonard Andreević Rastrigin. Systems of extremal control. Nauka, 1974.
Xin Rong. word2vee parameter learning explained. arXiv preprint arXiv:1411.2736, 2014.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. /397/, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30 (NIPS 2017), 2017.
Qi Wang and Herke Van Hoof. Learning expressive metu-rrepresentations with mixture of expert neural processes. Advances in neural information processing systems, 35:26242-2615 2022.
13

---

Published as a conference paper at ICLR 2025
Qi Wang, Marco Federici, and Herke van Hoof. Bridge the inference gaps of neural processes via expectation maximization. The Eleventh International Conference on Learning Representations, 2023.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. In Advances in Neural Information Processing Systems 30 (NIPS 2017), 2017.
14

---

Published as a conference paper at ICLR 2025
A ADDITIONAL DISCUSSION FOR THE FUTURE WORK
As demonstrated in our experiments with the MIMIC-III dataset in Appendix D.12, the ultimate goal in the Neural Processes field should be developing a general foundation regressor model capable of handling wide range of data structures and realistic scenarios, such as diverse time series data and cases with missing features. We view the DANP research as the initial step toward achieving this ambitious objective.
key focus of this future work direction will be to extend the model's ability to appropriately process inputs with varying dimensions, numbers of context and target points, and diverse data structures (for example there can be lots of different tasks with the same dimensional inputs such as EMNIST image completion and 2d GP regression task). Developing model that can flexibly adapt such variability without specific data processing based on inductive bias while providing accurate and reliable inferences across these scenarios remains critical challenge and exciting direction for further exploration.
B ADDITIONAL RELATED WORKS
The first Neural Process (NP) model, the Conditional Neural Process (CNP) (Garnelo al., 201 8a), utilized straightforward MLP layers for both the encoder and decoder. Neural Process (NP) (Garnelo al., 20) extended this by incorporating global latent variable to capture model uncertainty. Enhancements followed with Attentive Neural Processes (ANP) (Kim et al., 2018) and Conditional Attentive Neural Processes (CANP), which introduced attention mechanisms the encoder for better context summarization. Transformer Neural Processes (TNP) (Nguyen & Grover, 2022) replaced MLPs with masked transformer layers, achieving otate-of-the-art performance.
Further advancements include Functional Neural Processes (Louizos et al., 2019), which employed local latent variables to improve uncertainty capture, and Bootstrapping Attentive Neural Processes (BANP) (Lee et al., 2020), which utilized a residual bootstrap approach to address model misspec- ification. Martingale Posterior Attentive Neural Processes (MPANP) (Lee et al., 2023) addressed uncertainty with the martingale posterior, offering Bayesian alternative.
Recent developments have expanded NPs' scalability and expressiveness. For example, Translation Equivariant Transformer Neural Processes (TE-TNP) (Ashman al., 2024) enhance spatio-temporal modeling with translation equivariance, which leverages symmetries in posterior predictive maps common in stationary data. Latent Bottlenecked Attentive Neural Processes (LBANP) (Feng et al., 2022a) and Mixture of Expert Neural Processes (MoE-NPs) (Wang & Van Hoof, 2022) improve latent variable modeling through bottlenecks and dynamic mixtures, improving computational efficiency and generalization across various meta-learning tasks. Autoregressive Conditional Neural Processes (AR CNPs) (Bruinsma et al., 2023) address temporal dependencies by autoregressively defining joint predictive distribution, while Self-normalized Importance weighted Neural Process (SI-NP) (Wang et al., 2023) refine inference through iterative optimization.
Other contributions include Constant Memory Attentive Neural Processes (CMANPs) (Feng et al., 2023), which reduce memory usage with constant memory attention blocks, and Gaussian Neural Processes (GNPs) (Markou et al., 2022), focusing tractable dependent predictions by modeling the covariance matrix a Gaussian predictive process. Efficient Queries Transformer Neural Processes (EQTNPs) (Feng al., 2022b) improve TNPs by applying self-attention only to the context points, retrieving information for target points through cross-attention. Together, these advancements address key limitations in uncertainty modeling, inference, and computational efficiency, forming the basis for further progress in NP research.
C EXPERIMENTAL DETAILS
To ensure reproducibility, we have included our experiment code in the supplementary material. Our code builds upon the official implementation of TNP (Nguyen & Grover, 2022). We utilize PyTorch (Ansel et al., 2024) for all experiments, and BayesO (Kim & Choi, 2023), BoTorch (Balandat
15

---

Published as a conference paper at ICLR 2025
al., 2020), and GPyTorch (Gardner et al., 2018) packages for Bayesian Optimization experiments. All experiments were conducted on either a single NVIDIA GeForce RTX 3090 GPU or an NVIDIA RTX A6000 GPU. For optimization, used the Adam optimizer (Kingma & Ba. 2015) with cosine learning rate scheduler. Unless otherwise specified, we selected the base learning rate, weight decay, and batch size from the following grids: {5x 10-5, 7 10-5 9 10 1 x 10-100030000 x 10-1, 5 10-1) for learning rate, {0, 1 } for weight decay, and {16, 32} } for batch size, based on validation task log- ligeliik.i.odd
C.I DETAILS MODEL STRUCTURES
Table 5: Model structure details CANP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
DETERMINISTIC PATH HIDDEN DIMENSION
128
MLP DEPTH FOR VALUE IN CROSS ATTENTION LAYER
4
MLP DEPTH FOR KEY AND QUERY IN CROSS ATTENTION LAYER
2
MLP DEPTH FOR SELF-ATTENTION INPUT LAYER
4
MLP DEPTH FOR SELF-ATTENTION OUTPUT LAYER
2
DECODER DEPTH
3
NUMBER OF PARAMETERS FOR ID GP REGRESSION
331 331906
Table 6: Model structure details ANP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
DETERMINISTIC PATH HIDDEN DIMENSION
128
LATENT PATI IIDDEN DIMENSION
128
MLP DEPTH FOR VALUE IN CROSS ATTENTION LAYER
4
MLP DEPTH FOR KEY AND QUERY IN CROSS ATTENTION LAYER
2
MLP DEPTH FOR SET.--ATTNNION INPUT LAYER
4
MLP DEPTH FOR CELE-AITENTION OUTPUT LAYER
2
DECODER DEPTII
3
NUMBER OF PARAMETERS FOR ID GP REGRESSION
348418
In this section, we summarize the structural details of CANP, ANP, BANP, MPANP, TNP, and DANP. important to note that while we report the number of parameters for the baseline models in I-dimensional GP regression scenario, their parameter counts increase as the input and output dimensions increase. In contrast, the number parameters our model remains constant regardless of the input and output dimension combinations. This proves that model is structurally efficient compared to other baseline models. Also, following Lee et al. (2023), we model MPANP without the Self-Attentin layer in the deterministic path.
16

---

Published as a conference paper at ICLR 2025
Table 7: Model structure details of BANP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
DETERMINISTIC PATH HIDDEN DIMENSION
128
MLP DEPTH FOR VALUE IN CROSS ATTENTION LAYER
4
MLP DEPTH FOR AND QUERY IN CROSS ATTENTION LAYER
2
MLP DEPTII FOR SELF-ATTENTION INPUT LAYER
4
MLP DEPTH FOR SELF-ATTENTION OUTPUT LAYER
2
DECODER DEPTH
3
NUMBER OF PARAMETERS 1D GP REGRESSION
364674
Table 8: Model structure details of MPANP
CATEGORY
DETAILS
MODEL. SPECIFICATIONS
DETERMINISTIC PATH HIDDEN DIMENSION
128
HIDDEN DIMENSION FOR EXCHANGEABLE GENERATIVE MODEL
128
DEPTH FOR EXCIIANGEABLE GENERATIVE MODEL
DEPTH FOR VALUE IN CROSS ATTENTION LAYER
MLP DEPTH FOR KRY AND QUERY IN CROSS ATTENTION LAYER
MLP DEPTH FOR SELF-ATTENTIO INPUT LAYER
4
DEPTH FOR SELF-ATTENTION OUTPUT LAYER
2
DECODER DEPTH
3
NUMBER OF PARAMETERS FOR ID GP REGRESSION
892415
C.2 EVALUATION METRIC FOR THE TASKS
Following Le al. (2018), we used the normalized predictive lioeliich:hood:
log 1o Dj.c)
(20)
for the CNP variants CANP and TNP, where {cj, ( denotes context or target points. For the other models, we approximated normalized predictive likeeiidl as follows:
K
log Pog Dj.e)
log
Σ
p(ygMlaj, 0 k=1
(21)
K
kEO
kEO
where (K)) are independent samples drawn from q(0,| (D.)) for E [K]. Again, o E {eg, tt} indicates context target points.
17

---

Published as a conference paper at ICLR 2025
Table 9: Model structure details of TNP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
HIDDEN DIMENSION FOR EMBEDDING LAYERS
64
NUMBER OF LAYERS FOR EMBEDDING LAYERS
4
HIDDEN DIMENSION FOR MASKED TRANSFORMER LAYERS
128
NUMBER OF LAYERS FOR MASKED TRANSFORMER LAYERS
6
NUMBER OF HEADS FOR MASKED TRANSFORMER LAYERS
4
DECODER DEPTH
NUMBER OF PARAMETERS FOR 1D GP REGRESSION
222082
Table 10: Model structure details of DANP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
HIDDEN DIMENSION FOR LINEAR PROJECTION IN DAB
32
HIDDEN DIMENSION FOR
DAB
32
HIDDEN DIMENSION FOR TRANSFORMER LAYERS LATENT PATH
64
NUMBER OF LAYERS FOR TRANSFORMER LAYERS IN LATENT PATH
HIDDEN DIMENSION FOR ATTENIINNN IN LATENT PATII
64
HIDDEN DIMENSION FOR MLP LAYERS IN LATENT PATH
128
NUMBER OF LAYERS FOR MLP LAYERS IN LATENT PATH
HIDDEN DIMENSTON FOR MASKED TRANSFORMER LAYERS
128
NUMBER LAYERS FOR MASKED TRANSFORMER LAYERS
NUMBER OF HEADS FOR MASKED TRANSFORMER LAYERS
DECODER DEPTH
2
NUMBER OF PARAMETERS FOR ID GP REGRESSION
334562
C.3 DATASET DETAILS N-DIMENSIONAL GP REGRESSION TASK
In n-dimensional Gaussian Process (GP) regression task, start by sampling the context and target inputs. Specifically, we first determine the number of context points |c] by drawing from uniform distribution Unif(n² x 5, n² 50 n² 5). The interval for this uniform distribution scaled by n² to ensure that as the input dimension increases, the number of context points also increases, which is necessary for constructing accurate predictive density for the target points.
Next, we sample the number of target points |t] from uniform distribution Unif(n? x 5, n² 50- |el) to keep the total number of points within manageable range. After determining the number of
18

---

Published as a conference paper at ICLR 2025
context and target points, sample the input for each context and target point from the uniform distribution Unif(-2,2) 2) independently for each dimension in [n]]
We then generate the outputs y using the corresponding kernel functions. We employ the RBF
kernel k(x, x') = exp exp =xxxxx|1"" and the Matern 5/22 kernel (xx,x') = s² (1 v
5d
? + 32
2
where d = ||x - x'1 || For these kernels, the parameters are sampled as follows: s ~ On1, 1.0), 8~ Unif(0.1,0.6)) and p ~ Unif(0.1,0.5).
C.4 EMNIST DATASET
We employed the EMNIST Balanced dataset (Cohen et al., 2017), which consists of 112,800 training samples and 18,800 test samples. This dataset encompasses 47 distinct classes, from which selected 11 classes for use. Consequently, our training and test datasets comprise 26,400 and 4,400 samples, respectively. Each image is represented by 28 x 28 grid with single channel. We mapped the pixel coordinates range from -0.5 0.5 and normalized the pixel values lie within [-0.5, 0.5]. We sample the number of context points |c] ~ Unif(5, 45) and the number of target points [c)).
[t] Unif(5, 50
C.5 CELEBA DATASET
We utilized the CelebA dataset (Liu et al., 2015), which includes 162,770 training samples, 19,867 validation samples, and 19,962 test samples. The images were center-cropped to 32x32 pixels, resulting in 32 x 32 grid with 3 RGB channels. As with the EMNIST dataset, we scaled the pixel coordinates to a range of -0.5 to 0.5 and normalized each pixel value within [-0.5, 0.5). We sampled the number of context points |c| from Unif(5, 45) and the number of target points [t] from Unif(5, 50 - lc)).
C.6 CELEBA VIDEO DATASET
For the CelebA Video dataset, we generated simple video dataset using the CelebA image dataset. Specifically, after normalizing the pixel coordinates and values according to the pre-processing steps for the CelebA dataset, set the original CelebA data as the initial frame at time = 0. We then gradually decreased the brightness by subtracting 5/255 from each channel for each time step t E [9], concatenating each generated image the previous ones. This resulted in simple video with 32 x 32 grid, 3 RGB channels, and 10 frames. Consequently, the input dimension was 3, combining the time axis with pixel coordinates. As with the CelebA dataset, we sampled the number of context points |c| from Uni 45) and the number of target points |t] from Unif(5, 50 - Icl).
C.7 BAYESIAN OPTIMIZATION
Except for the BO experiments on HPO-B benchmark, we adjust the objective function to have the domain of [-2.0, 2.0]. We evaluated our method using various benchmark datasets and real-world scenarios. Below, we provide details of these experiments.
Hyperparameter Tuning on HPO-B benchmark We utility the HPO-B benchmark (Pineda- Arango et al., 2021), which consists of 176 search spaces (algorithms) evaluated sparsely on 196 datasets with a total of 6.4 million hyperparameter evaluations. In this benchmark, the continuous hyperparameters normalized in [0, 1). and the categorical hyperparameters are one-hot encoded. We use all the search space except for the 18-dimensional space, i.e., 2-, 3-, 6-, 8-, 9-, 10-, and 16-dimensional search spaces are used for the experiments. To construct meta-batch from each task, we sample the number of context points |c| from Unif(5, 50) and the number of target points |t] from Unif(5, 50 - Ic]); therefore, we exclude tasks which lesser than 100 (= 50 + 50) hyperparameter evaluations. For baselines,
first pre-train them on all the tasks collected from all the 2-dimensional search spaces of meta-train split. We then fine-tune them on meta- meta-batches randomly sampled from each target task (6-,5 9., 10,, 16-dim). To prevent overfitting the limited data, early-stop the
hhttp::/mmmaa..c.cuhhef edu www.kkwwww..../oppeeeccccccccaaa
19

---

Published as a conference paper at ICLR 2025
training with respect to the likelihood meta-validation split of target task. For DANP, pre-train it on the tasks collected from 2-, 3-, and 8-dimensional search spaces meta-trait split. We then target tasks, which corresponds to zero-shot setting. For BO,
evaluate it without fine-tuning run 200 iterations for each hyperparameter tuning task in meta-test split; therefore, we exclude tasks which have lesser than 210 hyperparameter evaluations. Furthermore, the order of hyperparameter dimensions is randomly shuffled to make the task diverse.
dimensional BO with GP generated objective functions evaluate basic BO performance when using each model surrogate for the black-box objective function, first create oracle sample using a GP with an RBF kernel and evaluate how well each model approximates these samples. We conducted ID BO for 100 iterations across 100 tasks using the expected improvement acquisition function.
and dimensional Bo benchmarks
We utilize three benchmark objective functions:
Ackley (Back, 1996) function
d
21 cos(czi) ) + exp(1)
f(x) = exp
2
exp
(22)
d
1=1
where v; E [-332.768 (1322768832.768) for all 1, ... d and global minimum is = (0, , ,0.
Cosine function
0.1 lad -1)
f(x) = cos(22 (2)
(23)
2nT
where x; E [-2x, 2nr, for all 1, and global minimum is is Ro (0.....)
Rastrigin (Rastrigin, 1974) function
f(x) = 10d + Σ [u] c0cos
(24)
1=1
where x; E [15.112, 5.12), for all i = 1, and global minimum is x² =o (0,...,))
To evaluate the models in multi-dimensional scenarios, we conduct experiments with cases of d = and d = for all the aforementioned benchmark functions. We perform evaluations over 100 iterations for each of the 100 tasks, utilizing the expected improvement acquisition function.
CNN BO For evaluation in a real-world BO scenario, we utilized the CIFAR-IO 10 dataset (Krizhevsky al., 2009) and Convolutional Neural Network (CNN) (LeCun et al., 1989). The CIFAR-10 dataset consists of 50,000 training samples and 10,000 test samples across 10 classes. In this setting, we generated 1,000 samples by creating combinations of weight decay, learning rate. and batch size, and trained the model for 20 epochs using the Adam optimizer (Kingma & Ba, 2015). The range for each hyperparameter is [1e 05, 1e 01] for learning rate, [1e 04, 1e 01] for weight decay, and [128, 256] for batch size, with 10 values selected uniformly within each range. These 1,000 samples were pre-generated, and we evaluated each BO task with initial sample and conducted 50 iterations for each of the 10 tasks using the expected improvement acquisition function.
D ADDITIONAL EXPERIMENTS
D.1 COMPARISON BETWEEN NEURAL DIFFUSION PROCESS AND DIMENSION AGNOSTIC NEURAL PROCESSES
it has structural limitation, being only As discussed in Section 4, NDP has two major issues: it
partially dimension-agnosti: for when y = 1, and not dimension aggnosic for other combinations, and 2) its reliance on diffusion-based sampling approximate the predictive distribution results in
20

---

Published as a conference paper at ICLR 2025
Table 11: Additional results of the context and target log- likellid for the GP regression task in the From-scratch scenario on NDP and DANP.
ID RBF
ID Matern
2D RBP
2D Matern
Model
langei
tangei
langet
langet
NDP
5.914 1.0.007
-0.376 1.0.077
5.924 10.0
503
016
5.945 13.041
-0.570
006
6.079 10.114 114
-0. 0.704 100002
DANP (ours)
1.381
0.921
1.383
0.068 +0
Table 12: Additional results of zero-shot scenario. The colored cell used pre-train DANP and NDP.
indicates the data dimension
NDP trained on 2D & 3D & 4D
DANP trained on 21) & 3D & 4D
Dimension
context
target
context
target
ID RBF
5.5664 #0.001
-0.5665 +0.097
1.366 #0.004
0.826 +0.018
2D RBF
5.9409 #0.002
-1.5654
+0.092
1.383 #0.000
0.335 #0.014
3D RBF
5.5935 #0.001
10.098
1.383 10.000
-0.261 =0.025
4D RBF
5.9792 #0.005
-7.8666
+0.095
1.383 #0.000
-0.568
#0.012
5D RBF
5.3512
#0.008
-8.4127
#0.103
1.359 #0.032
-0.676
+0.004
7D RBF
5.4938 +0.009
-14.6106 #0.101
1.355 +0.022
-0.723 =0.022
significantly high computational costs during inference and limited likelihood performance. Table 11 and Table clearly show that while NDP outperforms DANP in terms of context likelihood, significantly underperforms in target likelihood. This discrepancy arises because NDP relies diffusion- sampling method to generate possible outputs and then calculates the empirical posterior distribution from the gathered samples. This approach leads the model to predict the context points with high accuracy and low variance, thus achieving high context likelihood. However, for target points, the model struggles accurately predict the distribution, resulting in a lower target likelihood. Moreover, in most of our tasks, it is more important to achieve high likelihood predictions for unseen target points rather than focusing the observed context points. Therefore, having higher target point likelihood is more crucial than having a high context likelihood. Specifically, as shown in Table 12, NDP struggles to simultaneously learn across diverse dimensional inputs. demonstrating that it cannot function effectively as general regressor for unseen dimensions.
D.2 GP REGRESSION TASK
D.2.1 OTHER METRICS
Table 13: Additional evaluation of CRPS metric and confidence interval coverage on the 1D GP regression task with RBF kernel. For the CRPS metric, a smaller value indicates better performance.
Model
context CI
target CI
context CRPS (L)
target CRPS (4)
ANP
0.999 0.000
0.889 + 0.014
0.024 0.000
0.075 0.004
BANP
0.999 # 0.000
0,904 0.001
0.024 = 0.000
0.075 + 0.000
CANP
0.999 0.000
0.001
0.024 0.000
0.076 0.000
MPANP
0.999 + 0.000
0.898 0.001
0.024 = 0.000
0.075 + 0.000
TNP
0.999 0.000
0.909
0.001
0.024 0.001
0.071 0.000
DANP
0.999 + 0.000
0,912 + 0.002
0.024 + 0.000
0.068 + 0.000
CRPS and Empirical Confidence interval coverage Here, we further evaluate DANP and other baselines using additional metrics like continuous ranked probability score (CRPS) and empirical confidence interval coverage. We have measured these metrics for the ID and 2D GP regression tasks. For the fair comparison, we used the checkpoints from the From-scratch experiment in Section 5.1 for all models. The results are presented in Table 13.
21

---

Published as a conference paper at ICLR 2025
First, regarding the confidence interval coverage, it is observed that the context confidence interval tends to be too wide for all models. This issue arises because Neural Process models set a minimum standard deviation of 0.1 during inference to account for training stability and data noise. Additionally, the value of 0.1 is simply a conventional choice when designing Neural Process models. Therefore, setting this value lower during model construction can help ensure appropriate confidence interval. On the other hand, the target confidence interval tends to be relatively narrow, with the DANP model showing the best results. Additionally, when looking at CRPS scores, clear that for context points, the models generally perform at similar level, but for target points, DANP shows better scores compared to the other models.
Table 14:
Results on additional metrics containing MAE, RMSE, R², RMSCE, MACE, and MA on Id GP regression task. Except for R², lower values for all these metrics indicate better alignment with the target and improved calibration performance.
Model
MAE (1)
RMSE (1)
R (*)
RMSCE (1)
MACE (L)
MA (1)
ANP
0.126 0.001
0.176 0.003
0.788 0.012
0.273 #.000
0.238
0,240 0.003
BANP
0.125 0.001
0.175 0.003
0.811 = 0.001
0.273 # 0.007
0,237 = 0.001
0.239 = 0.002
CANP
0.127 + 0.001
0.178 1
0.801 : 0.005
0.267 # 0.008
0.239 + 0.002
0.237 :: 0.005
MPANP
0.124 + 0.001
+
0.807 # 0.005
0.274
0.242 = 0.007
0.244 + 0.008
TNP
0.122 + 0.002
0.173 + 0.001
0.808 # 0.002
0.287 = 0.003
0.251 = 0.005
0.253 = 0.006
DANP
0.120 0.001
0.165 + 0.002
0.816 + 0.002
0.259 = 0.000
0.230 = 0.003
0.228 = 0.002
Metrics related to calibration We additionally measure some other metrics related to calibration. We measured and reported the following 6 additional metrics: 1) Mean Absolute Error (MAE), 2) Root Mean Square Error (RMSE), 3) Coefficient of Determination ( (R²2,, 4) Root Mean Square Calibration Error (RMSCE). 5) Mean Absolute Calibration Error (MACE), 6) Miscalibration Area (MA). Except for R², lower values for all these metries indicate better alignment with the target and improved calibration performance. We conducted the evaluation using models trained on Id GP regression task, comparing our method with the baselines. The results, summarized in Table 14, demonstrate that DANP achieves the best performance across a range of metrics. This observation reaffirms that DANP not only outperforms in terms of NLL but also achieves improved performance in calibration-relates metries compared to the baselines. These additional evaluations highlight the robustness of our method across diverse aspects of model performance.
D.2.2 FINE-GRAINED EVALUATION ON THE ID GP REGRESSION TASKS
Table 15:
Additional fine- fine-grained evaluation on the ID GP regression task. Here, we evaluate for the less context and more context scenarios.
Less context
More context
Model
context
target
context
target
ANP
1.380 + 0.000
0.323 # 0.006
1.375 = 0.000
1.165 = 0.002
BANP
1.380 0.000
0.334 0.002
1.375 + 0.000
1.172 # 0.001
CANP
1.380 0.001
0.291 0.005
1.374 + 0.000
1.156 + 0.001
MPANP
1.380 + 0.000
0.317 + 0.007
1.374 + 0.001
1.165 # 0.003
TNP
1.382 + 0.000
0.363
0.005
1.379 + 0.001
1.209 + 0.004
DANP
1.383 0.000
0.396 0.003
1.380 + 0.000
1.214 + 0.002
Here, we evaluate how uncertainty behavior changes under various conditions for each method. To explore these changes, we considered three settings in GP regression tasks: 1) scenarios with a small number of context points versus large number, 2) situations with high noise scale, and 3) cases where the training kernel differs from the evaluation kernel. The experimental results can be found in Table 15 and Table 16.
First, in the ID GP regression experiment reported in Section 5.1, the number of context points ranged randomly from minimum of 5 to maximum of 45 for evaluation. For the first setting, 'small
22

---

Published as a conference paper at ICLR 2025
Table 16: Additional fine-grainee evaluation on the ID GP regression task. Here, we evaluate the increased noise scale and kernel change scenarios.
Noise
Kernel change
Model
context
target
context
target
ANP
1.377 + 0.000
0.855 # 0.004
1.373 + 0.000
0.667 + 0.003
BANP
1.377 0.000
0.864 + 0.001
1.373 + 0.000
0.688 + 0.003
CANP
1.377 0.000
0.837 0.003
1.372 + 0.000
0.641 + 0.003
MPANP
1.376 0.001
0.856 = 0.005
1.372 = 0.001
0.663 = 0.005
TNP
1.381 + 0.000
0.906 0.005
1.380 = 0.000
0.697 = 0.003
DANP
1.381
0.001
0.922 0.002
1.381 = 0.000
0.717 = 0.008
Table 17: Additional experimental results showing the context and target log-likelihoods for the GP regression task Zero-shot scenarios. The first column labeled nD represents the outcomes for the n-dimensional GP dataset with an RBF kernel. Cells highlighted in indicate the data dimension used to pre-train DANP. The initial context and target log-likelihood results are derived from the DANP model trained on 2 and 3-dimensional GP datasets with an RBF kernel. The subsequent results are obtained from the DANP model trained on 2, 3, and 4-dimensional GP datasets with both RBF and Matern kernels.
2D & 3D with RBF
2D & 3D & 4D with RBF and Matern
Dimension
context
target
context
target
ID RBF
1.354 +:0023
0.826 +0.048
1.360 +0.022
0.830 ++0.330
2D RBF
1,383 #0.000
0.341
#0.008
1.383
#0.001
0.318 #0.021
3D RBF
1.383 +0.001
-0.251
#0.008
1.383
+0.001
-0,296 +0.050
4D RBF
1.368 +0.003
0.661:0.011
1.383
#0.001
-0.601 +0.063
5D RBF
1.293 +0.050
-0.711 #0.012
1.378 #0.003
-0.679 #0.005
context' refers to using 5 to 15 context points, while 'large context" involves 30 to 45 context points for evaluation. In the second setting, the variance of the Gaussian noise used was increased to 2.5 times the original value, and the model was evaluated using this adjusted evaluation set. Lastly, for the third setting, we evaluated the model, trained on an RBF kernel, with an evaluation set generated using a Matern 5/2 kernel.
As shown in Table 15, DANP clearly outperforms other baselines in both small and large context scenarios. Notably, DANP demonstrates superior performance compared the other baselines in the small context scenario, indicating its ability to accurately predict the predictive distribution with limited number observations. Moreover, as illustrated in Table 16, DANP excels in both the increased noise scale and differing kernel scenarios. These results confirm that DANP can effectively adapt unseen tasks by learning generally shared features across different tasks.
D.2.3 ADDITIONAL. RESULTS FOR THE ZERO-SHOT SCENARIO
In the Zero-shot scenario for the GP regression task, we conducted two additional experiments utilizing DANP pre-trained with: 1) 2 and 3-dimensional GP datasets using the RBF kernel, and 2) 2, 3, and 4-dimensional GP datasets using both RBF and Matern kernels. The log-likelihood results in Table 17 for the first experiment indicate that DANP effectively predict the density of unseen dimensional GP datasets, though performance slightly declines for higher-dimensiee
datasets that farther from the trained dimensions, compared to the results in Table 2a. This demonstrates that while DANP can perform zero-shot inference on unseen dimensional datasets, training on a diverse range of dimensions enhances predictive performance.
In the second experiment, the log- ligeliikl results show that DANP can be trained on diverse tasks with different kernels. Notably, DANP was able to simultaneously train on 3, and 4-dimensional GP
23

---

Published as a conference paper at ICLR 2025
Table 18: Additional experimental results showing the context and target logeliiels for the GP regression in Fine-tuning scenarios. Here, we utilize 1-dimensional GP regression task with the Matern kernel as downstream task.
Full fine-tuning
Freeze fine-tuning
Method
context
target
context
target
CANP
-0.352
=0.053
-0.512
+0.034
-0.233
#0.108
-0.408
#0.084
ANP
-0.280 =0.094
-0.348
+0.080
-0.343
#0.044
-0.422
=0.023
BANP
-0.320 =0.033
-0.401
=0.037
-0.025
+0.145
-0.207
#0.154
MPANP
-0.128 0.246
-0.259
+001255
-0.260
+00092
-0.490
+0.142
TNP
-0.086
00004
-0.476
+0.139
0.336
+-0.188
-0.243
+0.162
DANP(ours)
1.372
-0.00
0.689
+-0.004
1.372 +-0.001
0.684
+-0.04
Table 19: Ablation results for ID, 2D GP regression and image completion tasks.
I0
ID Manura
280
3D Maam
DNNES
CIMbA
TXP
1.58
11710 wi8
1.363
1.35 11:00
1:10 Suen
+ Lakoa
1.38 BARG
8.929 6 0.100
L 41.00
0.723 94 063
1.383 2.000
01371 00.31 0..00
1086 & 0001
1,399 g (000)
Eden DURD
4,140 DUKC
1513 5/C3
DANP
1,388
1.363
13:33
1.383
18333 4..00
1,385 (0:90)
DUBC
2,07
PMAA
17BE 9106
11917
+ (69
+8 606
+8
1
+0
0.37 0..0
1 00000
B/K2 e 0/00
1.961 00000
B 0007
1118- (100)
2.00 6/001
datasets with both RBF and Matern kernels without increasing model size. By increasing the model size to accommodate more features and additional structural layers, DANP can generalize to wider variety of tasks with different generating processes.
D.2.4 ADDITIONAL RESULTS FOR THE FINE-TUNING SCENARIO WITH DIFFERENT KERNEL.
In the Fine-tuning scenario for the GP regression task, we fine-tuned on 160 I-dimensional regression tasks using the Matern kernel. For baselines, we utilized pre-trained models that had been trained on 2-dimensional GP regression tasks with the RBF kernel, as detailed in Section 5.1. For DANP, we used models pre-trained on 2, 3, and 4-dimensional GP regression tasks with the RBF kernel, also as described Scction 5.1. The results in Table 18 clearly demonstrate that while the baselines fail to generalize, DANP generalize the 1-dimensional GP regression task with the Matern kernel almost as effectively as the From-scratch results in Table with only a few datasets. This indicates that DANP not only generalize well to unseen dimensional GP tasks with known kernel but also to unseen dimensional GP tasks with an unknown kernel, compared to other baselines.
D.3 FULL EXPERIMENTAL RESULTS FOR THE SECTION 5.4
In this subsection, we report the full log-likelihood results from the ablation study on both the context and target datasets. In Table 19, it can be easily observed that the context exhibits similar trends to the target as we discussed in Section 5.4.
D.4 ABLATION RESULTS ON DIFFERENT OBJECTIVES
As highlighted in Foong et al. (2020), the ELBO loss we used for DANP does not provide the exact ELBBO for the - log (w (r,, D.). because we use D.) instead of p.)) More precisely, the Maximum Likelihood Loss is a biased estimator of log Po(ye (2,, Da), and the ELBO we used is lower bound of the same quantity. Therefore, both losses still share the same issue, and the effectiveness of each loss depends on the model.
Typically, the maximum likelihood loss tends to exhibit larger variance compared to variational inference, so, given our model's 's need to handle multiple varying dimensional tasks simultaneously, opted for variational inference to ensure stability. However, it is worth experimenting with other loss functions. Therefore, we conducted additional experiments and included the results from training with the maximum likelihood loss well.
24

---

Published as a conference paper at ICLR 2025
Table 20:
Comparison of zero-shot performance between DANP trained with the variational loss and the maximum likelihood loss. Here, each method trained with 2 and 4D GP datasets with RBF kernel while performing inference on the 3, 4, and 5D GP datasets with RBF kernel
Variational Inference
Marginal Likelihood
Method
context
target
context
target
ID RBF
1.336
+:0047
0.806
+0.048
1.340
::0.225
0.790 +0.008
2D RBF
1.383
#0.000
0.340
=0.007
1.383
10.000
0.330 #0.012
3D RBF
1.377
+0.007
-0.360
+0.063
1.381 +00001
-0.420
+0.112
4D RBF
1.379
=0.007
-0.589
#0.056
1.383 =0.000
-0.614
#0.045
5D RBF
1.357 +-0.012
-0.689 ++0.004
1.356 +-0.040
-0.701 +-0.223
Table 21: Comparison of zero-shot performance between DANP trained with the variational loss and the maximum likelihood loss. Here, each method trained with 2, 3, and 4D GP datasets with RBF kernel while performing inference on the 1, 2, 3, 4, and 5D GP datasets with RBF kernel
Variational Inference
Marginal Likelihood
Method
context
target
context
target
ID RBF
1.366
#0.004
0.826 #0.018
1.360
#0.006
0,805 00.021
2D RBF
1.383 #0.000
0.355
#0.014
1.382
+0.000
0.285 #0.012
3D RBF
1.383 #0.000
-0.261
#0.025
1.383
#0.001
-0.320
#0.044
4D RBF
1.383
+:0.00
+00042 -0.568
1.381
+0.002
-0.658
+.0039
5D RBF
1.359 #0.032
-0.676 #0.004
1.364 10.021
-0.742 #0.006
We conducted ablation experiments on the ML loss and VI loss using DANP trained on and 4d GP data, as well as DANP trained on 2d, 3d, and 4d GP data. These experiments were performed in zero-shot scenario by inferring on 1, 2, 3, 4, and 5d GP regression data. The results, presented in Table 20 and Table 21, show that while ML loss occasionally yields better log- loglliielisoods for context points, the VI loss consistently provides superior performance for the target points, which are of greater interest during inference. This trend is particularly evident in experiments trained on 2, 3, and 4d GP data. These findings demonstrate that using the VI loss for training DANP is generally more beneficial for improving generalization compared to the ML loss.
D.5 ABLATION EXPERIMENTS ON POSITIONAL EMBEDDING IN DAB MODULE
Many previous works have shown that sinusoidal positional encoding tends to perform poorly (Press al., 2021) in terms of generalization when extrapolating longer sequence lengths for Large Language Models. In response to this, approaches like Rotary Position Embedding (RoPE: Touvron et al., 2023; Su et al., 2024) have been proposed and used to address these limitations. While sinusoidal positional encoding successfully handled interpolation and extrapolation in our experimental settings, RoPE could potentially improve this performance. Therefore, we conducted additional experiments using a modified RoPE-based encoding tailored for the DAB module.
In our implementation, we retained the basic formulation of RoPE while ensuring different positional encodings for the each x and dimensions, similar to the approach we used with DAB. Specifically, we distinguished the embeddings added to queries and keys from x and y by alternating the cosine and sine multiplications for each. For example, if for x calculate qıx cos((po)) + q2 sin(pos), then for y. we compute 919 sin(pos) + 42y (poss.
Using this modified positional encoding, we conduct additional experiments the zero-shot and the fine-tune scenario in Gaussian Process regression tasks using the same settings in the main paper to evaluate the impact of RoPE the performance of our model.
We conducted ablation experiments on sinusoidal PE and RoPE a zero-shot scenario by inferring on 1D, 2D, 3D, 4D, and 5D GP regression data using DANP models trained 2D and 4D GP
25

---

Published as a conference paper at ICLR 2025
Table 22: Comparison of zero-shot performance between DANP trained with the sinusoidal positional embedding and RoPE. Here, each method trained with 2, and 4D GP dataset with RBF kernel while performing inference on the 2, 3, 4, and 5D GP dataset with RBF kernel
sinusoidal PE
RoPE
Positional Embedding
context
target
context
target
ID RBF
1.336 +0.047
0.806
+0.048
1.352 +0.012
0.777
10.035
2D RBF
1.383 =0.000
0.340
=0.007
1.383
=0.000
0.348 #0.003
3D RBF
1.377 +0.007
-0.360
+0.063
1.381 +0.001
-0.360
#0.013
4D RBF
1.379 #0.007
-0.589
+0.056
1.383
=0.000
-0.577 =0.008
5D RBF
1.357 +-0.012
-0.689 +-0.004
1.351
+-0.224
-0.704 +00019
Table 23: Comparison of zero-shot performance between DANP trained with the sinusoidal positional embedding and RoPE. Here, each method trained with 2, 3, and 4D GP dataset with RBF kernel while performing inference on the 4, and 5D GP dataset with RBF kernel
sinusoidal PE
RoPE
Positional Embedding
context
target
context
target
ID RBF
1.366 #0.004
0.826 =0.018
1.367
=0.002
0.787
00.021
2D RBF
1.383 +:0000
0.355
+0.014
1.382
+00000
0.334
0.007
3D RBF
1.383 #0.000
-0,261
+0.025
1.383
+0.001
-0.256
0.006
4D RBF
1.383 +0.000
-0.568
+0.042
1.383
+0.002
-0.576
00.036
5D RBF
1.359
#0.032
-0.676
=0.004
1.367 #0.014
-0.679 =0.007
regression data, well as on 2D, 3D, and 4D GP regression data. The results, presented in Table 22 and Table 23, indicate that while sinusoidal PE consistently outperforms RoPE in the ID case, their performance largely similar across other dimensions. This suggests that for these scenarios, both sinusoidal PE and RoPE exhibit comparable interpolation and extrapolation capabilities.
We also conducted experiments using the trained models perform few-shot learning on ID GP regression, following the setup in the main paper. As shown in Table 24, while there were some performance differences in the zero-shot setting for the ID GP regression task, these differences largely disappeared after few-shot fine-tuning. This indicates that the choice of positional embed- ding whether sinusoidal PE or RoPE has minimal impact on performance once the model is fine-tuned.
D.6 ABLATION ON GP REGRESSION SETUP AND ZERO-SHOT EVALUATION
Because most of the models have trouble with extrapolation rather than interpolation, it is important to analyze our method's 's extrapolation capabilities as compared to performance interpolation settings. To address this, we conducted additional experiments by training on the {1,2), 2]. and {3.4] dimensional cases, then evaluating the results on [1, 2, 3, 4, 5) dimensional test data.
Here, we train DANP utilizing both sinusoidal PE and RoPE to further analyze their generalization ability. Table 25 and Table 26 present the performance of DANP when trained on data from {1, 2} dimensions and (3, 4} dimensions, respectively.
From Table 25, we observe that when trained on the limited range of (1,2) dimensions, both positional embedding methods fail to learn sufficient general features, leading to lower generalization performance compared to training on (2,4] or {2, 3, 4} dimensions. This result emphasizes the importance of training on ligees-dimensional data to capture general features that enable better generalization to unseen dimensions. A similar pattern is evident in Table 26.
However, distinct trend emerges in Table 25 compared Table 22 and Table 23. While both sinusoidal PE and RoPE performed similarly when sufficient general features could be learned from more diverse training dimensions, RoPE demonstrates noticeably weaker generalization ability than
26

---

Published as a conference paper at ICLR 2025
Table 24: Comparison of fine-tune performance between DANP trained with the sinusoidal PE and
the RoPE. Here, each method trained with 2, and 4D GP datasets or 2, 3, and 4D GP datasets with RBF kernel while performing few-shot training on the 1D GP dataset with RBF kernel. Here, we report the performance for both the full fine-tuning and freeze finetuning
Full finetuning
Freeze finetuning
Positional Embedding
context
target
context
target
2,4D sinusoidal PE
1.375 +:0.001
0.890
+0.004
1.375
+00001
0.889 +00002
2,3,4D sinusoidal PE
1.375 +=0.00
0.893 +0.004
1.376
+00001
0.890 +00005
2,4D RoPE
1.375 10.001
0.886 =0.020
1.374
10.001
0.884
#0.015
2,3,4D RoPE
1.376 #0.000
0.882 =0.006
1.376 #0.001
0.882 #0.007
Table 25:
Comparison of zero-shot performance between DANP trained with the sinusoidal PE and the RoPE. Here, each method trained with I, and 2D GP datasets with RBF kernel while performing inference on the 1, 2, 3, 4, and 5D GP datasets with RBF kernel.
sinusoidal PE
RoPE
Positional Embedding
context
target
context
target
ID RBF
1.381 +0.000
0.916 #0.003
1.381 #0.012
0.916
#0.002
2D RBF
1.383 +0.000
0.346
#0.001
1.383 #0.000
0.350
10.006
3D RBF
1.307 #0.004
-0.633
#0.030
1.056 #0.204
-0.919
#0.172
4D RBF
1.138
+-0.012
+0.005 -0.817
0.101 +0.676
-1.685
+0.416
5D RBF
0.885
#0.022
-0.961 #0.069
-1.223 #0.758
-2.899
10.360
sinusoidal PE when the training data is limited the narrow dimensional range of {1,2]. This result highlights the dependency of RoPE on richer training data which contains richer general features to achieve high generalization ability.
D.7 ADDITIONAL EXTRAPOLATION RESULTS FOR THE FINE-TUNING SCENARIO
We conducted additional fine-tuning experiments on 5 GP regression data to analyze the extrapo- lation ability of our method. In this experiment, aim compare not only the performance of single DANP model against the baselines but also evaluate and compare multiple variants of DANP trained on different dimensional GP data. Specifically, we include DANP models trained on (1,2), f3. 4). {2, 4}. and {2, 3, }} dimensional GP data, as well as the corresponding DANP models where sinusoidal PE is replaced with RoPE.
The results in Table 27 clearly demonstrate that DANP outperforms the baselines in extrapolation few-shot scenarios, showcasing its robustness in handling these challenging tasks. Additionally, we observe that the DANP trained with 1,2d RoPE shows notable improvement in generalization perfor- mance when provided with a few-shot setting. However, despite this improvement, its performance on the target data remains inferior compared to other DANP training settings, such as those utilizing higher-dimensiooall data ((3, 4). (2, 4), or {2, 3, 43) or sinusoidal PE.
D.8 TRAINING BOTH GP REGRESSION AND IMAGE COMPLETION
To further demonstrate the ability of DANP to learn various tasks simultaneouslyy, we conducted an experiment involving both GP regression tasks and image completion tasks. Specifically, we trained our model on 2 and 3-dimensional GP regression tasks with the RBF kernel, well as on EMNIST and CelebA image completion tasks. We then evaluated our model using an additional 1-dimensional GP regression task. As shown in Table 28, although the performance slightly decreased compared to training each task separately, DANP successfully learned all training tasks and generalized well to the unseen task. This demonstrates that DANP capable of simultaneously training diverse
27

---

Published as a conference paper at ICLR 2025
Table 26: Comparison of zero-shot performance between DANP trained with the sinusoidal PE and the RoPE. Here, each method trained with 3, and 4D GP datasets with RBF kernel while performing inference on the 1, 2, 3, 4, and 5D GP datasets with RBF kernel.
sinusoidal PE
RoPE
Positional Embedding
context
target
context
target
ID RBF
1.130 +0.042
0.501 +0.016
1.239
+0.021
0.472
100019
2D RBF
1.301 #0.008
0.178 =0.010
1.369
#0.001
0.248
=0.012
3D RBF
1.383 +0.000
-0.278
+0.005
1.383 00.001
-0.265
=0.00
4D RBF
1.383 +0.00
-0.582 =0.014
1.383
=0.000
-0.556 #0.006
5D RBF
1.359 +-0.012
-0.701 +-0.115
1.242 +-0.224
-0.726 +0.044
Table 27: Comparison of fine-tuning performance between DANP with various settings and the baselines. Here, use the few-shot 5d GP regression task with RBF kernel for the evaluation. We also compare the performances for both full finetuning and freeze finetuning for all models.
Full fine-tuning
Freeze fine-tuning
Method
context
target
context
target
ANP
-0.851 =0.017
-0.552
#0.016
-0.837
#0.024
-0.837
#0.025
BANP
-0,817 #0.012
-0.813
=0.011
-0,830
#0.013
-0.828
#0.016
CANP
-0.854 +0.026
-0.856
+0.022
-0.847
#00057
-0.851
+:0050
MPANP
-0.862 +0.081
-0.863
+0.883
0..110+0.006 +0.016
-0.911
+:0015
TNP
-0.825 +0.081
-0.831
+0083
-0.830
+00021
+00023 -0.831
2,4D sinusoidal PE
1.382 +-0005
-0.674
+0.003
1.382
+-0.001
-0.674
++0.003
2,3,4D sinusoidal PE
1.382 +0.001
0.672
+0.004
1.382 +00001
+0.671 +00006
1,2D sinusoidal PE
1.301
#0.020
-0.772
00.33
1,300
#0.021
-0.774 +0.030
3,4D sinusoidal PE
1.377
#0.006
-0.683
10.004
1,377
10.006
-0.684 #0.004
2,4D RoPE
1.381 #0.001
-0.672
#0.001
1.382 #0.001
-0.672
#0.001
2,3,4D RoPE
1.382 +0.000
-0.672
0.003
1.382 +0.001
-0.672 +0.004
1,2D RoPE
1.126 +0.010
-0.901
00000
1.124 +0.009
-0.903 +0.005
3,4D RoPE
1.371 #0.009
-0.693
=0.023
1.374 #0.006
-0.691 #0.021
tasks and generalizing across different tasks. The model's performance could be further improved by increasing its capacity, either by expanding the feature space or adding more layers.
D.9 ADDITIONAL EXTRAPOLATION EXPERIMENTS FOR THE IMAGE COMPLETION TASK
We conducted an additional experiment on the CelebA landmark (Liu et al., 2015) task to further demonstrate the capabilities of our method. In the standard CelebA landmark task, the goal is to predict the locations of five facial landmarks: left eye, right eye, left mouth corner, right mouth corner, and nose, based on single image. However, since Neural Processes predict distribution over the target points using a given context, we adapted the CelebA landmark task to better fit this approach. We modified the task by combining the image's RGB values with the corresponding coordinates for each landmark, creating 5-dimensional input. The output was restructured as 5-dimensional label representing which of the five facial regions the prediction corresponds to. This setup allowed us to train and evaluate the model in way that aligns with the predictive distribution framework of Neural Processes.
For the experiment, we used pre-trained models for the baselines, specifically the CelebA image completion models, while we trained DANP on both the EMNIST dataset and CelebA image completion tasks. This approach allowed to assess the performance of DANP under a slightly modified but challenging setup, testing its ability to generalize across different types tasks. Table 29 validates that DANP still performs well on the different types of tasks compared other baselines.
28

---

Published as a conference paper at ICLR 2025
Table 28: Additional results for training both GP regression tasks and image completion tasks. We
trained DANP with 2 and 3-dimensional GP dataset and EMNIST and CelebA image completion tasks.
Dataset
context
target
ID RBF
1.299 +0.023
0.710 =0.032
2D RBF
1.381 ++0.000
0.294 +:0005
3D RBF
1.381
#0.000
-0.313 #0.020
EMNIST
1.382
#0.000
0.888 =0.004
CelebA
4.148 #0.000
1.895 =0.024
Table 29: Experimental results on the modified CclebA landmark task. Here, we fine-tuned baselines with 100-shot CelebA landmark dataset.
Full fine-tuning
Freeze fine-tuning
Method
context
target
context
target
ANP
0.572 10.024
0.557 10.027
0.568
#0.022
0.554 10.027
BANP
0.636 +0.031
0.574 +0.020
0.628
#0.027
0.568
#0.023
CANP
0.525 #0.030
0.506 #0.028
0.523
#0.031
0.504
=0.028
MPANP
0.536 #0.036
0.485 +0.023
0.535
#0.034
0.487
#0.024
TNP
0.658 +00020
0.557 +0.035
0.653
+0.021
0.554 +0.033
DANP(ours)
1.354 #0.001
0.674 #0.007
1.340 #0.002
0.672 +0.005
For the zero-shot scenario, DANP achieves 1.171 + 0.020 for the context dataset and 0.252 + 0.003 for the target dataset. These results demonstrate that although the target likelihood of zero-shot DANP is lower compared to that of fine-tuned aaeelins--rimariii due to variations in both input and output dimensions from the training data-DANP quickly surpasses other baselines after fine-tuning. This highlights DANP's robust ability to generalize effectively in challenging zero-shot scenarios while rapidly improving with minimal fine-tuning.
D.10 BAYESIAN OPTIMIZATION
To illustrate the wide-ranging applicability of DANP, we conducted BO (Brochu et al., 2010) experi- ments across various scenarios: 1) a 1-dimensional bo experiment using objective functions derived from GPs with an RBF kernel, 2) 2 and 3-dimensional Bо benchmarks, and 3) hyperparameter tuning for a 3-layer CNN (LeCun et al., 1989) on the CIFAR-10 (Krizhevsky et al., 2009) classification task. Following Nguyen & Grover (2022), we utilized Ackley, Cosine, and Rastrigin benchmark functions as the objective functions for the 2 and 3-dimensional Bо experiments. For the hyperparameter tuning of the 3-layer CNN, we initially trained 1000 CNN models with varying hyperparameters. including learning rate, batch size, and weight decay, and then identified the optimal hyperparameter combination using NP models. We measured performance using best simple regret, which measures the difference between the current best value and the global best value. And, we run 50 iterations for all the BO experiments. For detailed information about the objective functions in the 2 and 3-dimensional BO and CNN training, see Appendix C. As baselines, we employed pre-trained models tasks, In contrast,
for each n-dimensional GP regression task corresponding the n--dimensional for DANP, we used a single model pre-trained with 2, 3, and 4-dimensional GP regression tasks in the Zero-shot scenario. The results in Fig. 5 demonstrate that DANP outperforms other baselines in terms of regret with same iteration numbers. This demonstrates that DANP is capable of serving as surrogate model for different BO tasks using only a single model without additional training using bo datasets. In Fig. 5, we only report BO results with 2-dimensional Cosine and 3-dimensional Ackley objective function among various 2 and 3-dimensional BO benchmarks.
Full results for the synthetic Bayesian Optimization Here, we present the comprehensive experi- mental results for 2 and 3-dimensional BO benchmark objective functions, including Ackley, Cosine,
29

---

Published as a conference paper at ICLR 2025
1D BO
2D Cosine
3D Ackley
CNN BO
Regret
413
643
642
ANP
BANP
CANP
MPANP
TNP
DANP
Figure 5: Results for bo with various tasks. These four figures, from left right, show the regret results for I I-dimensional GP with RBF kernel, 2-dimensional cosine, 3-dimensional Ackley, and the CNN BO experiments.
Rastrigin
Ackley
Cumulative
8
Iterations
Reeations
Berations
Rerations
ANP
BANP
CANP
MPANP
TNP
DANP
Figure 6: Full Results for BO with 1-dimensional GP generated BO tasks and 2-dimensional bench- mark BO tasks. Here, we present cumulative regret results in addition regret results.
and Rastrigin. Additionally, we report cumulative regret results alongside regret results for all Bо experiments. Similar to the BO experiments outlined Appendix D.10, we employed pre-trained models for each n-dimensional GP regression task corresponding to the n-dimensional BO tasks as baselines. In contrast, for DANP, we utilized a single model pre-trained with 3, and 4-dimensional GP regression tasks in the Zero-shot scenario. The results depicted in Fig. 6 and Fig. 7 demonstrate that DANP is proficient in serving as a surrogate model for various BO tasks using only single model, without requiring additional training BO datasets.
D.11 IMAGE COMPLETION AND VIDEO COMPLETION
Additional visualization examples for the Image completion and Video completion In this section, we provide additional visualization examples for the image completion task using the EMNIST and CelebA datasets, well as for the video completion task with the CelebA video datasets. First, in Fig. 8, we display true video examples generated using the process described in Appendix C. It is visually apparent that the images gradually become darker over time. Next, visualize 10 example images from the EMNIST and CelebA datasets. Fig. and Fig. 10, the full images shown in the first column and the context points in the second column. Following that, we sequentially visualize the predicted mean and variance of CANP, ANP, BANP, MPANP, TNP, and DANP. And finally, in Fig. 11, report predictive mean and variance of video trained DANP.
30

---

Published as a conference paper at ICLR 2025
3D Cosine
Rastrigin
3D Ackley
CNN BO
Rogret
Cumulative Regret
Iterations
Iterations
teeations
Itorations
ANP
BANP
CANP
MPANP
TNP
DANP
Figure 7: Full Results for with 3-dimensional benchmark Bo tasks and CNN BO tasks. Here, present cumulative regret results in addition to regret results.
Table 30: Empirical results on the time series blood pressure estimation task.
Full fine-tuning
Freeze fine-tuning
Method
context
target
context
target
CANP
0.964 +0.030
0.875
+00024
0.962
#0.031
0.870 #0.022
ANP
1.037 #0.021
0.950
10.017
1.035
+0.021
0.947 10.019
BANP
1.104 +0.018
0.968
10.011
+0.017 1.100
0.966 +0.012
MPANP
1.012 #0.016
0.938
#0.018
1.010
0.930 #0.010
TNP
1.165 #0.020
0.987
10.013
1.160 #0.022
0.986 #0.011
DANP(ours)
1.235 +0.001
1.184 10.006
1.230 #0.002
1.180 +0.005
D.12 TIME-SERIES EXPERIMENT
To further validate the practicality, conducted additional experiments time series data using the blood pressure estimation task from the MIMIC-III dataset (Johnson et al., 2016). Specifically, we assumed real-world scenarios where certain features from patient data might be missing, entirely different sets of features could be collected. Under this assumption, we trained the model using only subset of features from the MIMIC-III dataset and evaluated its performance when additional or different sets of features became available.
Specifically, we considered five features: T, Heart Rate, Respiratory Rate, SpO2, and Temperature. For pre-training, we utilized T and Heart Rate features, while for the fine-tuning scenario, we assumed only Respiratory Rate, SpO2, and Temperature features were available (this scenario can happen we assume that we trained our model with data from hospital A and want to evaluate the data in hospital B). We pre-trained the models with 32,000 training samples and fine-tuned them with only 320 samples. And here, considered observations from time O..., & as context points and 1..,..., T as target points. As shown in Table 30, our DANP achieved strong performance in the time series blood pressure estimation task, demonstrating robustness and adaptability in this real- world scenario. These results are consistent with the findings presented in the main paper, further validating DANP's effectiveness in handling diverse and practical challenges.
31

---

Published as a conference paper at ICLR 2025
t=
a
t=2
t=3
t=4
t=5
t=
t=7
t=8
Figure 8: Examples of video data constructed following Appendix C.
Table 31: Wall clock time evaluation for the TNP and DANP in various settings. Here, utilize RTX 3090 GPU for the evaluation.
Model ID regression
2D regression
EMNIST
CelebA
TNP
min 30 sec
min 50 sec min sec
min min 20 sec
min 20 sec min 40 sec
DANP
min sec
D.13 DISCUSSION ON THE RESOURCE REQUIREMENTS
Here, we will analyze the time complexity compared to the TNP both theoretically and practically.
First theoretically, let denote B, N, dz, dy. d,, Ldt- and L4 denote the batch size, the number of data points (union of context and target), the dimension of input x. the dimension of output the representation dimension, the number of layers in the deterministic path, and the number of layers in the latent path, respectively. The additional computational cost for the DAB module is (dB + d,,)2) dd), and for the latent path, it is OLL Since the computational cost for TNP ( B.,NN Nd,,,, the overall computational cost of DANP be expressed as O((,,
32

---

Published as a conference paper at ICLR 2025
Img
CANP
BANP A
TNP 8
8
g
8
8
0
8
/
/
/
/
/
/
/
3
3
3
3
3
3
3
6
6
7
7
7
7
7
7
7
7
7
7
4
4
4
Figure 9: Predicted mean and variance of EMNIST dataset with baselines and DANP,
Full Img
Context
CANP
CANP
ANP μ
ANP
BANP J
BANP
MPANP y MPANP
TNP
DANP
DANP
Figure 10: Predicted mean and variance of CelebA dataset with baselines and DANP.
BN) + O(BN(d, + dd). Generally, since N >> (dz + d2,)² holds, the dominant term in the computational cost can be approximated as O( + T.a) BN
For the practical time cost, measure the time cost train 5000 steps for the GP regression tasks and image completion tasks for TNP and DANP. The results are shown in Table 31.
33

---

Published as a conference paper at ICLR 2025
Full img
context
mean
std
Figure 11: Predicted mean and variance of video data with DANP when training with video dataset.
34


Reviews:
[{"review_id": "LZXRPO6FKs", "reviewer": "Reviewer_LwkF", "summary": "The authors introduce Dimension Agnostic Neural Process (DANP) that incorporates Dimension Aggregator Block (DAB) to transform input features into a fixed-dimensional space, in an attempt to enhance the model's ability to handle diverse datasets. Leveraging a transformer architecture and latent encoding layers, the proposed approach learns a wide range of features, generalizable across various tasks. \n\nTo evaluate the approach the authors present a comprehensive evaluation, including synthetic and practical regression tasks. The empirical results, consisting of comparisons with exiting state-of-the-art methods and ablations show that the effectiveness of the proposed approach.  The authors outperforms past existing Neural Process methods, demonstrating advantages and improvements on GP regression, Image and Video Completion and Bayesian Optimization tasks.", "questions": "In Table 1, it seems that the proposed approach has only marginal advantage over the proposed methods, in the GP Regression in the from-scratch case. Since, I'm able to see and it appears that between TNP and DANP (proposed approach) across 1D RBF, 1D Matern, 2D RBF and 2D Matern at the target column the only difference (improvement) is at the second (third) decimal. \n\nWhat is the reason for that?\n\nI would have expected higher improvement (e.g. on the first decimal). I'm not sure I can consider this to be statistically significant. Can I ask also whether confidence intervals are available and would it be possible to share them?  \n\nSimilar performance is for the image completion task.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Originality \n - The paper seem it has an evident level of novelty, tackling the diverse input and output dimensions challenge in the uncertainty aware meta-learning methods such as neural processes. Two novelties seems to be the case here, the dimension aggregation block and the latent path, in a transformer-like arhitecture.\n\nQuality\n- The paper is well motivated, structured and presented, the problem is well introduces and connected to existing work. The writing is good. There is extensive and diverse evaluation over synthetic and publicly available data sets for the GP regression, Image and Video Completion, and Bayesian Optimization tasks. The Ablation study is also useful.\n\nClarity\n- The proposed approach is nicely presented and explained.\n\nSignificance\n- The approach shows effectiveness of such methods under evaluated dataset and seems to hold potential for applicability in diverse regression scenarios", "weaknesses": "Presentation of the tasks/problems that the method addresses\n- The presentation is sufficiently clear. I find that more on the actual task considered here can help to appreciate more the significance and the benefits of this approach. In particular related how the GP regression and the image completion tasks benchmarks help to validate the broad applicability of the approach? \n\nEvaluations\n- Non consistent result and seem to be marginal improvements in the GP Regression (from-scratch case) and the image completion task.", "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": "", "extra_scores": {}}, {"review_id": "bs7ZHstpZL", "reviewer": "Reviewer_8LUr", "summary": "This paper introduces a new model, Dimension Agnostic Neural Process (DANP), which addresses the limitations of current Neural Processes (NP) in handling inputs with varying dimensions. DANP includes a Dimension Aggregator Block (DAB) that converts input features into a fixed-dimensional space, while also incorporating Transformer architecture and latent encoding layers to enhance adaptability across different tasks. Experimental results show that DANP performs well on multiple regression tasks, highlighting its potential as a versatile, general-purpose regressor.", "questions": "I mainly have concerns about the practical aspects of this work. It would be helpful if the authors could provide more concrete examples and relate them to more complex, real-world applications.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Introduces the DAB module, enabling the model to handle inputs and outputs of varying dimensions, adding flexibility.\n2. Covers multiple tasks and scenarios, demonstrating the model's stability across different conditions.\n3. Performs well in regression, hyperparameter tuning, and other tasks, showing promise for broad applications.", "weaknesses": "1. Model Complexity. The design is complex, making replication and understanding challenging. \n2. Lacks Analysis of Computational Costs. There’s no discussion of the model's time and resource requirements, impacting assessments for practical use.\n3. Limited Application Scope. Primarily validated on regression tasks, with little exploration of classification or other tasks.", "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": "", "extra_scores": {}}, {"review_id": "ZoTcmwFKmk", "reviewer": "Reviewer_P7fT", "summary": "This work introduces a novel approach to meta-learning, specifically addressing the challenges of accommodating diverse input dimensions and learned features in Neural Process (NP) methods. \n- The authors propose the Dimension Agnostic Neural Process (DANP), which incorporates a Dimension Aggregator Block (DAB) to transform input features into a fixed-dimensional space, enhancing the model's ability to handle varied datasets. \n- By leveraging the Transformer architecture and latent encoding layers, DANP is capable of learning a broader range of features that are generalizable across different tasks. \n- Through extensive experimentation on various synthetic and practical regression tasks, the authors demonstrate that DANP outperforms previous NP variations, effectively overcoming the limitations of traditional NP models and showcasing its potential for broader applicability in diverse regression scenarios.", "questions": "- How does the authors' proposed DANP model perform in tasks outside of regression, such as classification or time-series forecasting? Are there any modifications needed for effective application in these domains?\n- Are there any plans to conduct longitudinal studies to evaluate the long-term performance and stability of DANP in dynamic environments?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- DANP is a novel extension of NP, that addresses the limitations of existing NP methods in handling diverse input dimensions and learned features.\n- This work not only points out the shortcomings of current NP methods but also proposes a robust solution through the DAB and the integration of Transformer architecture.\n- The paper is clear in its structure and presentation.", "weaknesses": "- The paper focuses on regression tasks, but its applicability to other tasks such as classification is not thoroughly explored. It could benefit from additional experiments or a theoretical discussion on how DANP might perform in non-regression tasks.\n- While DANP shows promising results, the paper lacks a detailed discussion on the model's interpretability. The paper should include an analysis or discussion on how the components of DANP contribute to its predictions, especially given its complex architecture involving the DAB and Transformer-based latent path.", "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": "", "extra_scores": {}}, {"review_id": "CMLrVEfw4G", "reviewer": "Reviewer_utkY", "summary": "This paper works on neural processes and studies the case when there exist diverse input dimensions and learned features. Tot this end, the Dimension Agnostic Block (DANP) is developed to transform input features into a fixed embedding and then combined with neural process modules. It conducts experiments in zero-shot and few-shot scenarios.\n\n--post rebuttal--\n\nAfter attending the discussion, I pose some revision suggestions on the notation system:\n\n(1) Both NP and CNP employ a **global latent variable (in NP, $z$ is parameterized with an encoder to obtain the Gaussian distribution by inputting $D_c$, in CNP $z$ is a deterministic embedding of $D_c$)**, and this means **marginalization the global latent variable $z$ out is standard instead of local latent variables** in NPs' family. In CNP, there is no marginalization out operation due to the deterministic global one.\n\n(2) Let's use the author's notation $\\theta_j$ as the latent variable for the encoder, it is right to use $p(\\theta\\vert D_c)$ rather than the data-point specific $\\theta_j$ as the latter means the local latent variable. \n\nI accept this work conditionally on revising the notation system to make it misleading.", "questions": "See the above", "limitations": "", "rating": 5, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "I can easily follow this work and the layout is clear. However, there are severe writing and examination issues.", "weaknesses": "(1) Overall, this work includes several engineering tricks to handle diverse input dimension cases and lacks theoretical analysis to examine the proposal. \n\n(2) It seems the motivation of this work also considers the uncertainty, however, I did not see sufficient results to illustrate this part when the dimension of the output is high.\n\n(3) In the related work [1], it has been demonstrated the Eq. (19) is not a valid ELBO. Hence, Line 266-269 should be revised. In line279, I disagree that NP is the earliest to address uncertainty as CNP can also achieve this in experiments. It seems several related works [2-9] are not discussed in the literature.\n\n(4) In line 316 and the following experiments, I am afraid that the zero-shot and finetune scenarios are not appropriate in evaluation as NP families require the context and amortize the few-shot adaptation without gradient updates. I am not convinced by the meaning of fine-tune or zero-shot in the task concept.\n\n(5) The computational complexity towards the modules and other ablations are required in examining the performance. Meanwhile it seems a lot of results in the context are nearly the same in scales in Table2-3. Details about the number of shots are missing in experiments, further weakening the results.\n\n---\n\nReference:\n\n[1] Foong A, Bruinsma W, Gordon J, et al. Meta-learning stationary stochastic process prediction with convolutional neural processes[J]. Advances in Neural Information Processing Systems, 2020, 33: 8284-8295.\n\n[2] Ashman, Matthew, et al. \"Translation Equivariant Transformer Neural Processes.\" arXiv preprint arXiv:2406.12409 (2024).\n\n[3] Feng, Leo, et al. \"Latent bottlenecked attentive neural processes.\" arXiv preprint arXiv:2211.08458 (2022).\n\n[4] Bruinsma, Wessel P., et al. \"Autoregressive conditional neural processes.\" arXiv preprint arXiv:2303.14468 (2023).\n\n[5] Wang Q, Federici M, van Hoof H. Bridge the inference gaps of neural processes via expectation maximization[C]//The Eleventh International Conference on Learning Representations. 2023.\n\n[6] Tailor D, Khan M E, Nalisnick E. Exploiting inferential structure in neural processes[C]//Uncertainty in Artificial Intelligence. PMLR, 2023: 2089-2098.\n\n[7] Markou S, Requeima J, Bruinsma W P, et al. Practical conditional neural processes via tractable dependent predictions[J]. arXiv preprint arXiv:2203.08775, 2022.\n\n[8] Wang Q, Van Hoof H. Learning expressive meta-representations with mixture of expert neural processes[J]. Advances in neural information processing systems, 2022, 35: 26242-26255.\n\n[9] Feng, Leo, et al. \"Memory efficient neural processes via constant memory attention block.\" arXiv preprint arXiv:2305.14567 (2023).\n\n[10] Feng L, Hajimirsadeghi H, Bengio Y, et al. Efficient Queries Transformer Neural Processes[C]//Sixth Workshop on Meta-Learning at the Conference on Neural Information Processing Systems.", "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": "", "extra_scores": {}}, {"review_id": "cPRhRSdNj9", "reviewer": "Reviewer_1CSd", "summary": "The paper describes the existing challenges of Neural Processes (NP) when using a variable number of input dimensions and learned features. The authors proposed to tackle this problem by extending the Transformer Neural Process architecture with a Dimension Aggregator Block, using Positional Embeddings to take into account the different input dimensions before transforming the features into a fixed dimensional space. The paper tests the performance of the proposed method on zero-shot and fine-tuning settings using synthetic regression tasks and image completion datasets well-known in the NP literature.", "questions": "- In the 0-shot evaluation on GP regression (section 5.1), why is dimension 6 not included in the validation? It seems straightforward to do and would paint a complete picture. Can this experiment be included? If it’s a matter of space in the manuscript or computational resources, I believe dimension 6 is better than dimension 7.\n- What are your thoughts on the inherent limitations of Positional Embeddings extrapolation? Is the lack of extrapolation a concern in the dimension-agnostic setting?\n- Did you experiment with Zero-Shot on image tasks? Given that you had results for fine-tuning, I imagine it is possible to get zero-shot results, unless I’m missing something.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- The paper has a good and concise summary of the Neural Processes setting and it highlights clearly the problem of fixed dimensions, as opposed to variable ones.\n- The authors submitted the source-code to reproduce their experiments, which is a neat way to connect the ideas exposed on the main body with practical details of the implementation. Great work here.\n- The idea of leveraging positional embeddings on the dimensions axis is I believe novel, and interesting on itself.", "weaknesses": "Post rebuttal: The authors have addressed many of the weaknesses I've listed below, added new ablations that show case the limitations of extrapolation and RoPE vs Positional Ebeddings among others. Hence I've increased my score. Reviewer utkY still has some methodological and novelty concerns, which I think are valid points, but I still think the empirical value of this work is enough to be accepted. Thanks to the authors for engaging on the rebuttal process.\n\n**Limitations of positional embeddings on the proposed DAB module:**\n\nPositional Embeddings is a well-studied part of the transformer architecture, specially in Language Models. Generally speaking, while there has been progress since the original Transformer, the community agrees that length extrapolation is an open-problem and positional embeddings do not extrapolate on further sequence lengths, see [1] for example, where the sinusoidal embeddings used by DAB are the worst in extrapolation. In LLMs (arguably the biggest current application of transformers and positional embeddings) the community has largely moved on from Sinusoidal embeddings and into other approaches such as RoPE [2] [3]. Thereby, while the setting here is very different I see the following weaknesses with the current paper given that positional embeddings is a crucial part for the DAB module to be dimension-agnostic:\n\n- Concretely, Positional Embeddings have a hard time in transformers when extrapolating from small to large context lengths. This warrants a discussion on whether they’re applicable to the setting on this paper.\n- There has been a substantial amount of work on newer and better positional embeddings, I believe this warrants at least a discussion and acknowledgement of recent work, and a stronger argument for using Sinusoidal Embeddings.\n- At best this warrants an ablation which justifies the choice in the architecture.\n- I believe highlighting the failure modes of using Positional Embeddings (sinusoidal or otherwise) do not diminish the contributions of the paper, quite on the contrary. But in my view, it’s important to highlight the limitations of the proposed approach, and if there’s evidence that these failure modes do not exist on the setting exposed here, it makes even an stronger paper.\n\n**Experimental weaknesses (dimension generalisation):**\n\nConnecting to the previous section, I believe that the experiments need to be more robust to analyse more carefully the failure modes of the proposed approach. In some cases, this is already in the appendix, but unfortunately this is not properly referenced from the main body. I urge the authors to present the failure modes of the approach in the main body more clearly, and when pushing results to the appendix, to reference them from the main body **explicitly,** calling out from the main body what are the strengths or weaknesses of the appendix results.\n\n- On the 0-shot scenario of GP regression (section 5.1), the model is either trained on {2,4} or {2,3,4} dimensions and evaluated on {1,2,3,4,5,7}. Given the above discussion about positional embeddings, I think that this setting benefits evaluating on dimensions {1,3} since they’re covered by interpolation. A better experimental setting would involve fully non-overlapping dimensions in train and validation, such as {1,2,3} and {4,5,6}. Even better, would be to try different configurations and report the threshold where the performance breaks down or when the performance is better. I hypothesise this model is better when it is evaluated during interpolation. At the very least, there should be an open discussion about this in the paper; however, as this is core to the contribution of this work, I believe it’s important to address these experimental weaknesses. Some of this is already in Appendix B.2 and in Table 16, however it is not clear if the kernels are the same (names are present in table 16 but not in table 2a); irrespective of this, given the arguments above, I strongly believe this discussion on the limitation of the approach belongs to the main body and should be **clearly stated on the conclusions** **and contributions sections** since it paints a full picture of the failure modes. The whole appendix B is barely referenced (it’s very big with very different subsections), and it’s very easy to miss the discussion of table 16.\n- For the fine-tuning section of GP regression, I have similar concerns. In the positional embedding LLMS literature, typically what you do to extend to longer contexts is to finetune on context lengths which are larger. Here, however, the finetuning happens on 1-dimensional tasks, after being pretrained on 2,3,4. Again, a likely failure mode is that when fine tuning on dimensions which are bigger than in pretraining the model won’t extrapolate very well. Unfortunately, as opposed to the zero-shot section, the appendix does not have experiments which reflect this setting.\n\n**Experimental weaknesses (practical regression tasks)**\n\n- While I understand the engineering challenges of video datasets, I believe the authors should tone down section 5.2. The proposed adaptation of CelebA does not make it a Video dataset — hence I believe it’s wrong to call this a Video Completion task, I suggest this is rephrased throughout the paper to reflect that this is just an image completion task with a synthetically generated extra dimension. Alternatively, it’d be really interesting to test this on an actual short video dataset, such as [4], however I understand this can be fully out of scope. Perhaps it’s useful to reference it as future work.\n- I believe the extra dimension added to CelebA is a rather weak task, since it’s a very simple subtraction over the brightness. I wonder if using off-the-shelf dimensions such as CelebA landmarks [5] [6] would be a better task (there are 5 different landmarks locations for each face, tagged as 2d coordinates). This has the benefit of being an already established benchmark in the literature, and arguably a more real-world and practical regression task. If the proposed method does not work well on these harder tasks, it’s still interesting to highlight the limitations of the model.\n- The core results of section 5.2 rely on pre-training on both the CelebA and EMNIST datasets, while interesting to see some positive transfer, after reading the framing of the paper and the GP regression section, I would have expected some zero-shot results on image tasks (for both pretraining and validation), which arguably is more relevant to meta-learning than fine-tuning. A Video Completion experiment in a zero-shot setting would make a stronger contribution, which is more consistent with the GP regressions section.\n\n**Experimental weaknesses (other)**\n\n- The log-likelihood is likely an incomplete picture of the results. As done in [7] I believe a stronger analysis should include other metrics such as calibration error.\n- The context provided about confidence interval coverage is worth mentioning on the main body, instead of being in the appendix where it’s hardly accessible and referenced directly from the main body.\n- In general, Appendix B is too broad and many interesting conclusions that should be in the main body are there without a direct reference.\n- In my opinion, the details that have been mentioned to be in the appendix, deserve to be in the main body to have a better complete picture of the strengths and limitations of the proposed approach. While I understand the page-limit is a limitation, I’d argue that what I have mentioned is of more relevance compared to the mostly positive results of section 5.4 — these can be in the appendix and mentioned briefly with a direct reference for interested readers. Another alternative is to leave the tables in the appendix but still mention the relevant conclusion in the main body.\n- NDP is mentioned as a relevant baseline and briefly tried on Appendix B - Table 11/12. Since this is the most relevant baseline according to the authors, I believe results with this relevant baseline from prior work should be in the main body as much as possible (it’s clear that it’s only comparable with dimension-agnostic methods for x when y=1) with a discussion about the quantitative results, not just qualitative as in section 4. There’s clearly at least two settings where it’s possible to do this, but they’re in the appendix.\n\n**Misc:**\n\n- The Dimension Aggregator Block implementation uses a linear projection **with bias [8].** This renders equation 10 incorrect since the bias term is missing there. I believe this is a common blind spot because the bias is set by default to True in Pytorch, so I urge the authors to revise all linear layers and either update the manuscript or the code/results. Note that I do not expect this to change the results much, but it’s better to address this for reproducibility and clarity.\n- Table 2a does not have the 2nd or 1st best performance marked with underline or boldfaced underline. This makes it harder to read.\n- Table 2a does not state which kernels were used, as compared to table 1 and table 16.\n- It’s hard to read table 1 and table 2 separately, and as the authors mentioned, it’s helpful to read them in context.\n- While I appreciate the detail problem setting on section 2, if there are concerns of page limits, I’d suggest the authors to prioritise space for experimental results and discussion as suggested in the previous section — this can be either added as an Appendix and referenced directly from the main body, or just cite from previous work as appropriate.\n\n[1] https://openreview.net/forum?id=R8sQPpGCv0\n\n[2] https://arxiv.org/abs/2104.09864\n\n[3] https://arxiv.org/abs/2302.13971\n\n[4] https://svdbase.github.io/\n\n[5] https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html\n\n[6] https://www.kaggle.com/code/danmoller/inspecting-the-celeba-dataset-face-landmarks\n\n[7] https://arxiv.org/abs/2207.04179\n\n[8] https://pytorch.org/docs/stable/generated/torch.nn.Linear.html#torch.nn.Linear", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)