Practical Equivariances via Relational Conditional Neural Processes
Authors: Daolang Huang
Manuel Haussmann
Ulpu Remes
S. T. John
Grégoire Clarté
Kevin Sebastian Luck
Samuel Kaski
Luigi Acerbi
Venue: NeurIPS 2023 poster (2023)
Keywords: neural processes
equivariance
Gaussian processes
TL;DR: RCNP is a new class of equivariant neural processes, with improved performance across tasks with inherent equivariances.

Abstract: Conditional Neural Processes (CNPs) are a class of metalearning models popular for combining the runtime efficiency of amortized inference with reliable uncertainty quantification. Many relevant machine learning tasks, such as in spatio-temporal modeling, Bayesian Optimization and continuous control, inherently contain equivariances – for example to translation – which the model can exploit for maximal performance. However, prior attempts to include equivariances in CNPs do not scale effectively beyond two input dimensions. In this work, we propose Relational Conditional Neural Processes (RCNPs), an effective approach to incorporate equivariances into any neural process model. Our proposed method extends the applicability and impact of equivariant neural processes to higher dimensions. We empirically demonstrate the competitive performance of RCNPs on a large array of tasks naturally containing equivariances.

Practical Equivariances via Relational Conditional Neural Processes
Daolang Huang'
Manuel Haussmann'
Ulpu Remes?
ST John'
Grégoire Clarté3
Kevin Sebastian Luck Samuel K
Luigi Acerbi
Department of Computer Science, Aalto University, Finland
2 DDpartmeent of Mathematics and Statistics. University of Helsinki
Department of Computer Science, University of Helsinki
Department of Electrical Engineering and Automation (EEA), Aalto University, Finland
'Department of Computer Science, University of Manchester Department of Computer Science, Vrije Universiteit Amsterdam, The Netherlands
(daolang.Ni) huang, manuel. haussmann, ti.john, samuel (taski)@aalla fi k.s. k.s.kuck@v..nl (ulpu.remes) gregoire. clarte, luigi. acerbi)ennkih
Abstract
Conditional Neural Processes (CNPs) class of metalearning models pop- ular for combining the runtime efficiency of amortized inference with reliable uncertainty quantification. Many relevant machine learning tasks, such in spatio- temporal modeling. Bayesian Optimization and continuous control, inherently contain equivariances for example to translation which the model can exploit for maximal performance. However, prior attempts to include equivariances in CNPs do not scale effectively beyond two input dimensions. In this work, we propose Relational Conditional Neural Processes (RCNPs), an effective approach incorporate equivariances into any neural process model. Our proposed method extends the applicability and impact of equivariant neural processes to higher di- mensions. We empirically demonstrate competitive performance of RCNPs on large array of tasks naturally containing equivariances.
Introduction
Conditional Neural Processes (CNPs; (10]) have emerged as a powerful family of metalearning models, offering the flexibility of deep learning along with well-calibrated uncertainty estimates and tractable training objective. CNPs can naturally handle irregular and missing data, making them suit- able for wide range of applications. Various advancements, such as attentive (ACNP: [20]) and Gaus- sian (GNP; [30]) variants, have further broadened the applicability of CNPs. In principle, CNPs can be trained on other general- purppoee stochastic processes, such Gaussian Processes (GPs; [34]), and be used as an amortized, drop-in replacement for those, with minimal computational cost at runtime.
However, despite their numerous advantages, CNPs face substantial challenges when attempting model equivariances, such as translation equivariance, which are essential for problems involving spatio-temporal components for emulating widely used GP kernels in tasks such as Bayesian Optimization (BayesOpt; [12]). In the context of CNPs, kernel properties like stationarity and isotropy would correspond to, respectively, translational equivariance and equivariance to rigid transformations. Lacking such equivariances, CNPs struggle to scale effectively and emulate (equivariant) GPs even in moderate higher-dimensional input spaces (i.e., above two). Follow-up work has introduced Convolutional CNPs (ConvCNP: [15]), which leverage a convolutional deep sets construction to
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

INT
OOID
(c)
(d)
RCNP
GP
-2
0
2
-2
Figure 1: Equivariance in 1D regression. Left: Predictions for a CNP (a) and RCNP (b) in an interpolation (INT) task, trained for 20 epochs emulate GP (c) with Matérn- kernel and noiseless observations. The CNP underfits the context data (black dots), while the RCNP leverages translation equivariance learn faster and yield better predictions. Right: The CNP (d) fails to predict in sngf--innndddisibitttte (OOID) task, where the imput context is outside the training range (note the shifted x axis); whereas the RCNP (e) generalizes by means of translational equivariance.
induce translational-cquivarinn embeddings. However, the requirement of defining an input grid and performing convolutions severely limits the applicability of ConvCNPs and variants thereof (ConvGNP [30]; FullConvGNP [2]) to one- wwo-dimensional equivariant inputs; both because digherrdimensional implementations of convolutions are poorly supported by most deep learning libraries, and for the prohibitive cost of performing convolutions in three more dimensions. Thus, the problem of efficiently scaling equivariances CNPs above two input dimensions remains open.
In this paper, we introduce Relational Conditional Neural Processes (RCNPs), novel approach that offers a simple yet powerful technique for including a large class of equivariances into any neural process model. By leveraging the existing equivariances of problem, RCNPs can achieve improved sample efficiency, predictive performance, and generalization (see Figure 1). The basic idea in RCNPs is to enforce equivariances via relational encoding that only stores appropriately chosen relative information of the data. By stripping away absolute information, equivariance is automatically satisfied. Surpassing the complex approach of previous methods (c.g.. the ConvCNP family for translational equivariance). RCNPs provide a practical solution that scales to higher dimensions, while maintaining strong performance and extending to other equivariances. The cost to pay is increased computational complexity in terms of context size (size of the dataset we are conditioning
runtime); though often not bottleneck for the typical metalearning small-context setting of CNPs. Our proposed method works for equivariances that can be expressed relationally via comparison between pairs of points (e.g., their difference distance); in this paper, we focus on translational equivariance and equivariance to rigid transformations.
Contributions. In summary, our contributions in this work are:
We introduce simple and effective way w encoding to encode exact equivari- ances directly into CNPs, in way that easily scales to higher input dimensions.
We propose two variants of relational encoding: one that works more generally ("Full"); and one which simpler and more computationally efficient ('Simple"). ), and is best suited for implementing translation cquivariance.
We provide theoretical foundations and proofs to support our approach.
We empirically demonstrate the competitive performance of RCNPs variety of tasks that naturally contain different equivariances, highlighting their practicality and effectiveness.
Outline of the paper. The remainder this paper is organized as follows. Section 2, we review the foundational work of CNPs and their variants. This is followed by the introduction of our proposed relational encoding approach to equivariances at the basis of our RCNP models (Section 3). We then provide in Section 4 theoretical proof that relational encoding achieves equivariance without losing essential information; followed in Section by a thorough empirical validation of our claims in various tasks requiring equivariances, demonstrating the generalization capabilities and predictive performance RCNPs. We discuss other related work in Section 6, and the limitations of our approach, including its computational complexity, in Section 7. We conclude in Section 8 with overview of the current work and future directions.
Our code is available
com http://hhw..coooccectaaaaarrrrra nwwraaaaaa-eeeaaaaeeee processes.

---

2 Background: the Conditional Neural Process family
In this section, we review the Conditional Neural Process (CNP) family of stochastic processes and the key concept of equivariance at the basis of this work. Following [30], we present these notions within the framework of prediction maps [8]. We denote with x E X' C R imput vectors and y y/ S Re output vectors, with du dy 2 1 their dimensionality. If f fz) is a function that takes as input elements a set Z. we denote with fZ) the set [z)z)xzzz
Prediction maps. A prediction map a is a function that maps (I) a context set (X, Y) comprising input/output pairs to distribution over the corresponding target outputs Y = "i.....;'
(xx, y yN)} and (2) a collection of target inputs X* = (xi, xM)
(X, Y). x = |rY
where r - r((, Y),X* is the representation vector that parameterizes the distribution over Y via the representation function r. Bayesian posteriors are prediction maps, well-known example being the Gaussian Process (GP) posterior:
π TY"I(X, Y) = N (Y" m, K),
(2)
where the prediction map takes the form of multivariate normal with representation vector r (m, K). The mean m = (X, Y), X*) and covariance matrix K kpoat (X, X*) of the multivariate normal are determined by the conventional GP posterior predictive expressions [34].
Equivariance. A prediction map with representation function r is Teqquivariant with respect to a group T of transformations' of the input space, X - X, if and only if for all T E T:
r((X, Y), x*) r((r(((()))) X, Y),,) rX))
where rX = rxx) and +X the set obtained by applying T to all elements of X. Eq. 3 defines equivariance of a prediction map based on its representation function, and can be shown to be equivalent to the common definition of an equivariant map; see Appendix A. Intuitively, equivariance means that if the data (the context inputs) target inputs) transform correspondingly. Common groups of transformations include translations, rotations, reflections- all examples of rigid transformations. In kernel methods and specifically in GPs, equivariances are incorporated in the prior kernel function k(x, x') For example, translational equivariance corresponds to stationarity has k(x-x")
transformed in certain way. the predictions (the
x*) ), and equivariance all rigid transformations corresponds to isotropy, kiss = &[[[x x" 2122), where || ||2 denotes the Euclidean norm of a vector. crucial question address in this work how implement equivariances in other prediction maps, and specifically in the CNP family.
Conditional Neural Processes. A CNP [10] uses an encoder? fe to produce an embedding of the context set, fe (X, Y). The encoder uses DeepSet architecture [48] to ensure invariance with respect to permutation of the order of data points, a key property of stochastic processes. We denote with rm (e, x the local representation of the m-th point of the target set X for m M. CNPs yield a prediction map with representation r = (ri...., IM):
. Y),
II a (ym A(rm))
where q-|lA) belongs to a family of distributions parameterized by X, and À fr(r..) is decoded in parallel for each r.. In the standard CNP, the decoder network fa is multi-layer perceptron. A common choice for CNPs a Gaussian likelihood, q(y", [A) = A (y (((r)) E(rm)), where μ and Σ represent the predictive mean and covariance of each output, independently for each target (a mean field approach). Given the closed-form likelihood, CNPs are easily trainable via maximum- likelihood optimization of parameters of encoder and decoder networks, by sampling batches of context and target sets from the training data.
group of transformations is a family of composable, invertible functions with identity Tux = x. We use purple to highlight parametrized functions (neural networks) whose parameters will be learned.
3

---

Gaussian Neural Processes. Notably, standard CNPs do not model dependencies between distinct target outputs yi and y for m # m'. Gaussian Neural Processes (GNPs [30]) are a variant of CNPs that remedy this limitation, by assuming a joint multivariate normal structure over the outputs for the target set, x (X, )*)
N (,,, Σ). For ease of presentation, we consider now scalar outputs (dy = 1), but the model generalizes to the multi-output case. GNPs parameterize the fx) (x:( (r. f. for target points x X where fu. fx. and f. neural networks with outputs, respectively, in R. Rds, and RR+, A:(:;" -) is a positive-definite kernel function, and dy: E N+ denotes the dimensionality of the space in which the covariance kernel is evaluated. Standard GNP models use the linear covariance (where fo = 1 and k is the linear kernel) or the kvv covariance (where is the exponentiated quadratic kernel with unit lengthscale), as described in [30].
mean as Mm = Jp. (Im) and covariance matrix Em.m² & (Fm))
Convolutional Conditional Neural Processes. The Convolutional CNP family includes the ConvCNP [15], ConvGNP [30], and FullConvGNP [2]. These CNP models built to imple- ment translational equivariance via a ConyDeepSet architecture [15]. For example, the ConvCNP is
prediction map p( (XX, X*) Im= m= (y' [xx (x,,,), where 2x, (.) is a CenvDeepSet. The construction of DeepSees involves, among other steps, gridding of the data if not already on the grid and application of d₂ de-dimensional convolutional neural networks (2dz for FullConv- GNP). Due to the limited scaling and availability of convolutional operators above two dimensions, ConyCNPs do not scale in practice for dz > 2 aaaussaationalyy-quivvrain input dimensions.
Other Neural Processes. The neural process family includes several other members, such as latent NPs (LNP: [1IJ) which model dependencies in the predictions via a latent variable - however, LNPs lack a tractable training objective. which impairs their practical performance. Attentive (C)NPs (A(C)NPs; [20]) implement an attention mechanism instead of the simpler DeepSet architecture. Transformer NPs [31] combine a transformer-batt based architecture with a causal mask to construct an autoregressive likelihood. Finally, Autoregressive CNPs (AR-CNPs [3]) provide novel technique to deploy existing CNP models via autoregressive sampling without architectural changes.
Relational Conditional Neural Processes
We introduce now our Relational Conditional Neural Processes (RCNPs), an effective solution for embedding equivariances into any CNP model. Through relational encoding, we encode selected relative information and discard absolute information, inducing the desired equivariance.
Relational encoding. In RCNPs, the (full) relational encoding of a target point x E X with respect to the context set (X, Y) is defined as:
Pran(x",, (X,Y)) = fo (g(x, Xin), Benr),
Ryn = (g(x,,,
(5)
n,p'==1
where g X x -> TR. dhop is a chosen comparison function² that specifies how a pair x, x' should be compared: R. is the relational matrix, comparing all pairs of the context set; fr Rdomm x R2d, -> da is the relational encoder, a neural network that maps a comparison vector and element of the relational matrix into high high-dimensional space gda; b is commutative aggregation operation (sum in this work) ensuring permutation invariance of the context set [48]. From Eq. 5, point x* is encoded based on how it compares to the entire context set.
Intuitively, the comparison function o(:-)) should be chosen remove all information that does not matter to impose the desired equivariance. For example, if we want to encode translational equiv- ariance, the comparison function should be the difference of the inputs, gdint ( x²m) = x" X2, (with deomp d.). Similarly, isotropy (invariance to rigid transformations, i.e. rotations, transla- tions, and reflections) can be encoded via the Euclidean distance gaist (xn, xx,) = | x,.l12 (with deomp- 1). We will prove these statements formally in Section 4.
We use green to highlight the selected comparison function that encodes specific set of equivariances.

---

Full RCNP. The full-context RCNP, or FulIRCNP, is prediction map with representation r r (p..... Par), with Pm = Piull (x (X, Y)) the relational encoding defined in Eq. 5:
M II q (y MPm)))
x
where q(- belongs family of distributions parameterized by A, where X = fa(P is decoded from the relational encoding pm of the m-th target. As usual, we often choose a Gaussian likelihood, whose mean and covariance (variance, for scalar outputs) are produced by the decoder network.
Note how Eq. (FulIRCNP) nearly identical to Eq. 4 (CNP), the difference being that we replaced the representation m (e, xm) with the relational encoding Pm from Eq. 5. Unlike CNPs, in RCNPs there is no separate encoding of the context set alone. The RCNP construction generalizes casily to other members of the CNP family by plug-in replacement of Im with pm For example, relational GNP (RGNP) describes a multivariate normal prediction map whose mean is parameterized as jmm = (Pm) and whose covariance matrix is given by Lm,m' ff (m)): fx((m))) (.
Simple RCNP The full relational encoding Eq. is cumbersome as it asks to build and aggregate over full relational matrix. Instead, consider the simple 'diagonal' relational encoding:
N
Pding (x X Y))
fr (g(xn: xm), I(x,, yn).
===1
Eq. 7 is functionally equivalent to Eq. 5 restricted to the diagonal n = n', and further simplifies the common case g(x,; Xa. O, whereby the argument of the aggregation becomes (x yn).
We obtain the simple RCNP model (from now on, just RCNP) by using the diagonal relational encoding Pdiag instead of the full one, Piull- Otherwise, the simple RCNP model follows Eq. 6. We will prove, both in theory and empirically, that the simple RCNP is best for encoding translational cquivariance. Like the FulIRCNP, the RCNP casily extends to other members of the CNP family.
In this paper, we consider the FuIIRCNP, FulIRGNP, RCNP and RGNP models for translations and rigid transformations, leaving examination of other RCNP variants and equivariances to future work.
4 RCNPs are equivariant and context-preserving prediction maps
In this section, we demonstrate that RCNPs are Teequívariant prediction maps, where T transformation group of interest (e.g., translations), for an appropriately chosen comparison function x X + deoe Then, we formalize the statement that RCNPs strip away only enough information achieve equivariance, but no more. We prove this by showing that the RCNP representation preserves information the context set. Full proofs are given in Appendix A.
4.1 RCNPs are equivariant
Definition 4.1. Let be a comparison function and T group of transformations 7 X - X. We say that g is T Tinvarriant if and only 9(x, x') = g(Tx, rx') for any x, x' E X' and T.
Definition 4.2. Given comparison function g, define the comparison sets:
((X, Y), (X, Y)) = ((((x. X Vn, (isnnnsn 9 ((X, x*) = ((((xn x²m) y (w) henSNNSMMM
(8)
9 "X,X") = (x(x x Jsmmm's
If g is not symmetric, we can also denote g (X', (X, Y)) ((o(x) x.). y Nwwhhhhssssssss
Definition 4.3. A prediction map and its representation function are relational with respect to comparison function if and only if can be written solely through set comparisons:
r((X, Y), x*) (((((X)))))) Y), (X, ()),(((X, Y), X), 9X", (X, gXX, X*)).
(9)
Lemma 4.4. Let n be a prediction map, T a transformation group, and g a comparison function. If is relational with respect to g and g is T -invarriant, then π is T -equivariant.
5

---

From Lemma 4.4 and previous definitions, we derive the main result about equivariance of RCNPs.
Proposition 4.5. Let be the comparison function used in a RCNP, and T group of transformations. If g is Tiivaariant, the RCNP T- -equivariant.
As useful examples, the difference comparison function (x, x*) x" x invariant to transla-
tions of the inputs, and the distance comparison function Gaiaa (x.x*) x²) )²=== = ||x x||2 invariant to
rigid transformations; thus yielding appropriately equivariant RCNPs.
4.2 RCNPs are context-presst context-preserving
The previous section demonstrates that any RCNP is T Tequivvrrian, for an appropriate choice of g. However, a trivial comparison function g(x, x') = O would also satisfy the requirements, yielding trivial representation. We need to guarantee that, at least in principle, the encoding procedure removes only information required to induce T -eequivariance, and no more. A minimal request is that the context set is preserved the prediction map representation r, modulo equivariances.
Definition 4.6. A comparison function g is context-prestt context-rreeerving with respect to transformation group T if for any context set (X, Y) and target set X', there is a submatrix Q' CQQ of the matrix ), Yn Yne). ), a reconstruction function y, and transformation T E 'T such that (Q') (rX,Y) Y).
Qman' = (g( X,: km).9/ (X,, )n),,
For example, Gdist is context-preserve preserving with respect to the group of rigid transformations. For any m, Q' = Qm: is the set of pairwise distances between points, indexed by their output values. Reconstructing the positions of a set of points given their pairwise distances is known as the Euclidean distance geometry problem [38], which
be solved uniquely up rigid transformations with traditional multidimensional scaling techniques [39]. Similarly, Jai is context- pratexxtpreservigg with respect translations. For any and n', Q' Qmia' can be projected to the vector ((x) T.m(.) = X
x yi),................... (XN x )), which is equal to (Tm X, Y) with the translation
Definition 4.7. For any (X, Y). X*, family of functions ho (X((X Y), X*) - The is context- preserving under transformation group T if there exists θ drep E N, reconstruction fenttruuutioouunttiioooo function y, and a transformation r E 'I such that y(ho ((X, Y), X*)) ((r) = Y).
Thus, an encoding is context-preserving if it is possible at least in principle to fully recover the context set from r, implying that no relevant context is lost. This is indeed the case for RCNPs.
Proposition 4.8. Let T be transformation group and the comparison function used in FullRCNP. with respect to T, then the representation function r of the FullRCNP is
If g is connexx-preeservee context context-preserving with respect to T.
Proposition 4.9. Let T be the translation group and gan the difference comparison function. The representation of the simple RCNP model with 9an is context-presenee with respect to T.
Given the convenience of the RCNP compared to FuIIRCNPs, Proposition 4.9 shows that we can use simple RCNPs to incorporate translation-equiivrraacee with no loss of information. However, the simple RCNP model is not context- pentext-rreservig for other equivariances, for which ought to use the FuIIRCNP. Our theoretical results are confirmed by the empirical validation in the next section.
Experiments
In this section, we evaluate the proposed relational models on several tasks and compare their perfor- mance with other conditional neural process models. For this, we used publicly available reference implementations of the neuralprocesses software package [1, 3]. We detail our experimental approach in Appendix B, and we empirically analyze computational costs in Appendix C.
5.1 Synthetic Gaussian and non-Gaussian functions
We first provide thorough comparison our methods with other CNP models using a diverse array of Gaussian and non-Gaussian synthetic regression tasks. We consider tasks characterized by functions derived from (i) range of GPs, where each GP sampled using one of three different kernels (Exponentiated Quadratic (EQ), Matémn and Weakly-Periodie); (ii) non- non-Gaussian sawtooth

---

process; (iii non-Gaussian mixture task. In the mixture task, the function is randomly selected from either one of the three aforementioned distinct GPs or the sawtooth process, each chosen with probability 1 Apart from evaluating simple cases with d, {1, 2}, we also expand our experiments to higher dimensions, dz = (3. 5, 10). In these higher-dimensional scenarios, applying ConvCNP and ConvGNP models is not considered feasible. We assess the performance of the models in two distinct ways. The first one, interpolation (INT), uses the data generated from a range identical to that employed during the training phase. The second one, out-of-input- inputddistribution (OOID), uses data generated from a range that extends beyond the scope of the training data.
Results. We first compare translation-eqquiarraan ( ('stationary") versions of RCNP and RGNP with other baseline models from the CNP family (Table 1). Comprehensive results, including all five regression problems and five dimensions, are available in Appendix D. Firstly, relational encoding of the translational equivariance intrinsic to the task improves performance, as both RCNP and RGNP models surpass their CNP and GNP counterparts in terms of INT results. Furthermore, the OOID results demonstrate significant improvement of our models, as they can leverage translational- equivariance to generalize outside the training range. RCNPs and RGNPs are competitive with convolutional models (ConvCNP, ConvGNP) when applied to ID data and continue performing well in higher dimension, whereas models in the ConvCNP family are inapplicable for d, > 2.
Table 1: Comparison of interpolation (INT) and ces.oof-nnuuddstrriiii
(OOID) performance of our RCNP models and CNP baselines on synthetic regression tasks with varying input dimensions. We show mean and (standard deviation) across 10 runs with different seeds. "F" denotes failed attempts that yielded very bad results. Missing entries could not be run. Statistically significantly best results are bolded. Our methods (RCNP and RGNP) perform better than their CNP and GNP counterparts in terms of both INT and OOID, and scale to higher dimension compared to ConvCNP and ConvGNP.
Weskly-periodi divergencc(()
Sootooth log-likclioook(t)
Mixture log-likelibboolt))
d_
dg 3
d,
d_
d, - 3
4,
d,
de -
d, 5
RCNP
0.24 (0.00)
0.28 (0.00)
0.31 (0.00)
3.03 (0.06)
0.85 (0.01)
0.44 (0.00)
0.20 (0.01)
-0.10 (0.00)
-0.31 (0.03)
RGINP (sta)
0.03 (0:00)
0.05 (0.00)
0.08 (0.00)
3.90 (0.09)
1.09 (0.01)
1.13 (0.05)
0.34 (0.05)
0.37 (0:01)
0.04 (0.02)
ComCNN
0.21 (0.00)
3.64 (0.04)
0.38 (0.02)
ConsGNP
0.01 (0:00)
3.94(0.11) (0.11)
0.49 (0.15)
0.31 (0.00)
0.39 (0.00)
(0.00)
2.25 (0.02)
(0.28)
-0.03 (0.10)
(0.01)
-0.57 (0.11)
-0.72 (0.08)
GNP
0.06 (0.00)
0.08 (0.01)
0.11 (0.01)
0.83 (0.04)
0.23 (0.13)
0.02 (0.05)
0.17 (0.01)
-0.17 (0.00)
-0.32 (0.00)
RCNP (sta)
0.24 (0.00)
0.28 (0.01)
0.31 (0.00)
3.04 (0.06)
0.85 (0.01)
0.44 (0.00)
0,20 (0.01)
-0.00 (0.00)
-0.31 (0.03)
RGNP (sta)
0.03 (0:00)
0.05 (0.01)
0.88 (0.00)
3.90 (0. 10)
1,09 (0.01)
1.13 (0.05)
0.34 (0.03)
0.37 (0.01)
0.04 (0.02)
O0ID
ConCNNP
0.21 (0.00)
3.64 (0.04)
0.38 (0.02)
CrevGNP
0.01 (0:00)
(0.08)
0.49 (0:15)
CNP
2,88 (0.91)
1.58 (0 (0.50)
2.20 (0.81)
F
40.37 (0.12)
-0.22 (0.03)
F
-2.55 (1.15)
-1.71 (0.55)
GNP
F 1.47 (0.27)
0.62 (0.04)
F
-0.67 (0.05)
-0.72 (0.03)
We further consider two GP tasks with isotropic EQ and Matérn- & kernels (invariant to rigid trans- formations). Within this set of experiments, we include the FullRCNP and FulIRGNP models, each equipped with the 'isotropic' distance comparison function. The results (Table 2) indicate that RCNPs and FulIRCNPs consistently outperform CNPs across both tasks. Additionally, we notice that FulIRCNPs exhibit better performance compared to RCNPs as the dimension increases. When dg = 2, the performance of our RGNPs is on par with that of ConvGNPs, and achieves the best results in terms of both INT and OOID when d₂ > 2, which again highlights the effectiveness of our models in handling high-dimensiooal tasks by leveraging existing equivariances.
5.2 Bayesian optimization
We explore the extent our proposed models can be used for higher-dimensinn meta-learning task, using Bayesian optimization (BayesOpt) as our application [12]. The neural processes, ours as well as the baselines, serve as surrogates to find the global minimum fmin f (xmin) of a black-box function. For this task, we train the models by generating random functions from a GP kernel sampled from a set of base kernels EQ, Matérn- ( 2. 2 2
1 5 as well as their sums and products- wiith randomly
sampled hyperparameters. By training on a large distribution over kernels, we aim to exploit the metalearning capabilities of neural processes. The trained CNP models are then used as surrogates to minimize the Hartmann function [40, p. 185] in three and six dimensions, common BayesOpt test function. We use the expected improvement acquisition function, which we can evaluate analytically. Specifics the experimental setup and further evaluations can be found in Appendix E.
7

---

Table 2: Comparison of the interpolation (INT) and cettof-innuuddisttribee (OOID) performance of our RCNP models with different CNP baselines on two GP synthetie regression tasks with isotropic kernels of varying input dimensions.
EQ KL divergence(()
Matémn- 3 divergence()) da 3
d,
d,
dg
4, 2
5
RCNP (sta)
0.26 (0.00)
0.40 (0.01)
0:45 (0.00)
0.30 (0:00)
0.39 (0.00)
0.35 (0.00)
RGNP (sta)
0.03 (0.00)
0.05 (0.00)
0.11 (0.00)
0.03 (0.00)
0.05 (0.00)
0.11 (0.00)
FullRCNP
0.26 (0.00)
0.31 (0.00)
0.35 (0.00)
0.30
0.32
0.29 (0.00)
RNIIKGNP (1so)
0.88 (0.00)
0.14 (0.00)
0.25 (0.00)
0.09 (0:00)
0.16 (0.00)
0.21 (0.00)
ConvCNP
0.22 (0.00)
0.26 (0.00)
ConyGNP
0.01
CNP
0.33 (0.00)
0.44 (0.00)
0.57 (0.00)
0.39 (0.00)
0.46 (0.00)
0.47 (0.00)
GNP
0.05 (0.00)
0.09 (0.01)
0.19 (0.00)
0.07 (0.00)
0.11 (0.00)
0.19 (0.00)
RCNP (sta)
(0.00)
(0.01)
0.45 (0.00)
0.30 (0.00)
0.39 (0.00)
0.35 (0.00)
RGNP (sta)
0.03 (0.00)
(0.00)
(0.00)
(0.00)
(0.00)
FullRCNP (iso)
0.26 (0.00)
0.31 (0.00)
0.35 (0.00)
0.30 (0.000
0.32 (0.00)
0.29 (0.00)
0OID
RuIIRGNP
0.08 (0:00)
(.14().00)
0.25
0.09
0.16
0.21 (0.00)
Com-CNP
0.22(0..00)
0.26 (0.00)
Con-CNP
(0.100
0.01 (0.00)
CNP
(1.7()
3.30 (1.55)
1.22 (0.09)
6.75 (2.72)
1,75 (0.42)
0.93 (0.02)
GNP
2.25 (0:61)
(1.44)
0.74 (0.02)
1.86 (0:26)
1.23 (0.17)
0.62 (0.02)
Hantmann 35
Hartmann Ed
3.0
2.5
CAP
2.5
2,3
2.0
KC
6:36
65
rumber quaiis
Figure 2: Bayesian Optimization. Error during optimization of a 3D/6D Hartmann function (lower is better). RCNP/RGNP improve upon the baselines, approaching the GP performance.
Results. As shown in Figure 2, RCNPs and RGNPs are able to learn from the random functions and come close to the performance of Gaussian Process (GP), the most common surrogate model for BayesOpt. Note that the GP is refit after every new observation, while the CNP models can condition on the new observation added to the context set, without any retraining. CNPs and GNPs struggle with the diversity provided by the random kernel function samples and fail at the task. In order to remain competitive, they need to be extended with an attention mechanism (ACNP, AGNP).
5.3 Lotka-Volterra model
Neural process models excel in the so-called sim-to-real task where a model is trained using simulated data and then applied to real-world contexts. Previous studies have demonstrated this capability by training neural process models with simulated data generated from the stochastic Aatka-Voltrra predator-prey equations and evaluating them with the famous hare-lynx dataset [32]. We run this benchmark evaluation using the simulator and experimental setup proposed by [3]. Here the CNP models are trained with simulated data and evaluated with both simulated and real data; the learning tasks represented in the training and evaluation data include interpolation, forecasting and reconstruction. The evaluation results presented in Table 3 indicate that the best model depends on the task type, but overall our proposed relational CNP models with translational equivariance perform comparably to their convolutional and attentive counterparts, showing that our simpler approach does not hamper performance on real data. Full results with baseline CNPs are provided in Appendix F.
5.4 Reaction-Diffuuc model
The Reaction-Diffusiee (RD) model is a large class of state-space models originating from chemistry [45] with several applications medieine [13) and biology [22]. We consider here a reduced model representing the evolution of cancerous cells [13], which interact with healthy cells through the
8

---

Table 3: Normalized log- log-likelihood scores in the Lotka-Volterra experiments (higher is better). The mean and (standard deviation) reported for each model are calculated based on 10 training outcomes evaluated with the same simulated (S) and real (R) learning tasks. The tasks include interpolation (INT), foreassinng (FOR), and reconstruction (REC). Statistically significantly (see Appendix B.1)
best results are bolded. RCNP and RGNP models perform on par with convolutional and attentive
baselines.
INT (S)
FOR (S)
REC (S)
INT (R)
FOR (R)
REC (R)
RCNP
-3.57 (0.02)
-485 (0.00)
4.20 (0.01)
-.24 (0.02)
4,83 (0.03)
-4.55 (0.05)
RGNP
-3,51 (0.01)
-4.77 (0.00)
-3.76 (0.00)
-4.31 (0.06)
-4,47 (0.03)
-4.39 (0.11)
ConvCNP
-3.47 (0.01)
-4855 (0.00)
-406 (0.00)
-4.21 (0.04)
-5.01 (0.02)
-4.75 (0.05)
ConvGNP
-3,46 (0.00)
4.30 (0.00)
-3.67 (0.01)
-4.19 (0.02)
4,61 (0.03)
-4.62 (0.11)
ACNP
4.04 (0.06)
-4877 (0.01)
-,36 (0.03)
-4.88 (0.05)
-,799 (0.03)
-4.48 (0.02)
AGNP
(0.17)
-4.35 (0.09)
-4.05 (0.19)
-4.33 (0.15)
-4,48 (0.06)
-4.29 (0.10)
production of acid. These three quantities (healthy cells, cancerous cells, acid concentration) are defined on discretized space-time grid (2+1 dimensions, d2 = 3). We assume we only observe the difference in number between healthy and cancerous cells, which makes this model hidden Markov model. Using realistic parameters inspired by [13], we simulated we subsample observations to generate training data for the models; another set of 10 trajectories is used for testing. More details can be found in Appendix G.
-00h full trajectories, from which
We propose two tasks: first, completion task, where target points at time are inferred through the context different spatial locations time 1, and t 1; secondly, a forecasting task, where the target at time is inferred from the context at t - 1 and { - 2. These tasks, along with the form of the equation describing the model, induce translation invariance in both space and time, which requires the models to incorporate translational equivariance for all three dimensions.
Results. We compare our translationallequiiaaaann RCNP models to ACNP, CNP, and their GNP variants in Table 4. Comparison with ConvCNP is not feasible, as this problem is tiree-dimensional. Our methods outperform the others on both the completion and forecasting tasks, showing the advantage of leveraging translational equivariance in this complex apatio-temponn modeling problem.
Table 4: Normalized log-likelihood scores in the Difffinniiffusion problem for both tasks (higher is better). Mean and (standard deviation) from 10 training runs evaluated on a separate test datasct.
RCNP
RGNP
ACNP
AGNP
CNP
GNP
Completion 0.22 (0.33)
1.38 (0.62)
0.17 (0.03)
0.20 (0.03)
0.10 (0.01)
0.13 (0.03)
Forecasting
0.07 (0.18)
-0.18 (0.58)
-0.65 (0.31)
-1.58 (1.05) -0.51 (0.20) -0.50 (0.30)
5.5 Additional experiments
As further empirical tests of our method, study technique in the context of autoregressive CNPs in Appendix H.I, present a proof proof-of-conee -of-cccceeptttt of incorporating rotational symmetry in Appendix H.2,
and examine the performance of RCNPs on image regression in Appendix H.3.
Related work
This work builds upon the foundation laid by CNPs [10] and other members of the CNP family,
covered at length in Section 2. A significant body of work has focused on incorporating equivariances into neural network models [33, 4, 23, 41]. The concept of equivariance has been explored in Convolutional Neural Networks (CNNs; [24]) with translational equivariance, and more generally in Group Equivariant CNNs [5], where rotations and reflections are also considered. Work on DeepSets laid out the conditions for permutation invariance and equivariance [48]. Set Transformers [25] extend

---

this approach with an attention mechanism to learn higher- oodder interaction terms among instances of a set. Our work focuses on incorporating equivariances into prediction maps, and specifically CNPs.
Prior work on incorporating equivariances into CNPs requires a regular discrete lattice of the input space for their convolutional operations [15, 30). EquivCNPs [19] build work by [7] which operates on irregular point clouds, but they still require a constructed lattice over the input space. SteerCNPs [18] generalize ConvCNPs other equivariances, but still suffer from the same scaling issues. These methods are therefore in practice limited to low-dimensional (one to two equivariant dimensions), whereas our proposal does not suffer from this constraint.
Our approach is also related to metric-based meta-learning, such as Prototypical Networks [36] and Relation Networks [37]. These methods learn an embedding space where classification can be performed by computing distances to prototype representations. While effective for few-shot classification tasks, they may not be suitable for more complex tasks those requiring uncertainty quantification. GSSM [46] aims to learn relational biases via graph structure on the context set, while we directly build exact equivariances into the CNP architecture.
Kernel methods and Gaussian processes (GPs) have long addressed issues of equivariance by cus- tomizing kernel designs to encode specific equivariances [14]. For instance, stationary kernels are used to capture globally consistent patterns [34], with applications in many areas, notably Bayesian Optimization [12]. However, despite recent computational advances [9]. kernel methods and GPs still struggle with high-dimensional, complex data, with open challenges in deep kernel learning [47] and amortized kernel learning (or metalearning) [26, 35], motivating our proposal of RCNPs.
Limitations
RCNPs crucially rely a comparison function g(x, x') to encode equivariances. The comparison functions we described (e.g., for isotropy and translational equivariance) already represent a large class of useful equivariances. Notably, key contributions the neural process literature focus only on translation equivariance (e.g.. ConvCNP [15], ConvGNP [30], FullConvGNP (2]). Extending our method to other equivariances will require the construction of new comparison functions.
The main limitation of the RCNP class is its increased computational complexity in terms of context and target set sizes (respectively, N and M). The FuiIRCNP model can be cumbersome, with O(N M) cost for training and deployment. However, we showed that the simple or 'diggonal' RCNP variant can fully implement translational invariance with O(NM) cost. Still, this cost is larger than O(N + M) of basic CNPs. Given the typical metalearning setting of small-data regime (small context sets), the increased complexity is often acceptable, outweighed by the large performance improvement obtained by leveraging available equivariances. This is shown in our empirical validation, in which RCNPs almost always outperformed their CNP counterparts.
Conclusion
In this paper, introduced Relational Conditional Neural Processes (RCNPs), member of the neural process family which incorporates equivariances through relational encoding of the context and target sets. Our method applies to equivariances that can be induced via appropriate comparison function; here we focused on translational equivariances (induced by the difference comparison) and equivariances to rigid transformations (induced by the distance comparison). How express other equivariances via our relational approach is an interesting direction for future work.
We demonstrated with both theoretical results and extensive empirical validation that method suc- cessfully introduces equivariances in the CNP model class, performing comparably to the translational- cquivariant ConvCNP models in low dimension, but with simpler construction that allows RCNPs to scale to larger equivariant input dimensions (dz > 2) and outperform other CNP models.
In summary, we showed that the RCNP model class provides simple and effective way to implement translational and other equivariances into the CNP model family. Exploiting equivariances intrinsic to problem can significantly improve performance. Open problems remain in extending the current approach to other equivariances which are not expressible via a comparison function, and making the existing relational approach more efficient and scalable to larger context datasets.
10

---

Acknowledgments and Disclosure of Funding
This work was supported by the Academy of Finland Flagship programme: Finnish Center for Artificial Intelligence FCAL Samuel Kaski was supported by the UKRI Turing AI World-Leading Researcher Fellowship, [EP/W002973/1] The authors wish thank the Finnish Computing Compe- tence Infrastructure (FCCI), Aalto Science-IT project, and CSC-IT Center for Science, Finland, for the computational and data storage resources provided.
References
[1] Wessel Bruinsma. Neural processes. httpp:://githuu.comm// http://hht..cceeeeeeee//aarrrpas
2023, Accessed: October 2023.
[2] Wessel Bruinsma, James Requeima, Andrew YK Foong, Jonathan Gordon, and Richard Turner. The Gaussian neural process. In 3rd Symposium on Advances in Approximate Bayesian Inference, 2021.
[3] Wessel P Bruinsma, Stratis Markou, James Requeima, Andrew YK Foong, Tom R Andersson, Anna Vaughan, Anthony Buonomo, J Scott Hosking, and Richard E Turner. Autoregressive conditional neural processes. In International Conference on Learning Representations, 2023,
[4] Benjamin Chidester. Tianming Zhou, Minh N Do, and Jian Ma. Rotation equivariant and invariant neural networks for microscopy image analysis. Bioinformatics, 35(14):530-i537, 2019.
[5] Taco Cohen and Max Welling. Group equivariant convolutional networks. In International Conference on Muchine Leurning, pages 2990-2999, 2016.
[6] Brian Drawert, Andreas Hellander, Ben Bales, Debjani Banerjee, Giovanni Bellesia, Bernie Daigle, Jr., Geoffrey Douglas, Mengyuan Gu, Anand Gupta, Stefan Hellander, Chris Horuk, Dibyendu Nath, Aviral Takkar, Sheng Wu, Per Lötstedt, Chandra Krintz, and Linda Petzold. Stochastic simulation service: Bridging the gap between the computational expert and the biologist. PLOS Computational Biology, 12012):e000522,, 2016.
[7] Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups arbitrary continuous data. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 3165-311 (6766, 2020.
[8] Andrew YK Foong, Wessel P Bruinsma, Jonathan Gordon, Yann Dubois, James Requeima, and Richard E Turner. Meta-learning stationary stochastic process prediction with convolutional neural processes. In Advances Neural Information Processing Systems, volume 33, pages 8284-8295, 2020.
[9] Jacob R Gardner, Geoff Pleiss, David Bindel, Kilian Q Weinberger, and Andrew G Wilson. GPyTorch: Blackbox matrix-matix Gaussian process inference with GPU acceleration. In Advances in Neural Information Processing Systems, volume 31, pages 7576-7586, 2018.
[10] Marta Garnelo, Dan Rosenbaum, Chris Maddison. Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. Conditional neural processes. In International Conference on Machine Learning, pages 1704-1713, 2018.
[11] Marta Gareel,, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Ali Eslami, and Yee Whye Teh. Neural processes. In ICML Workshop on Theoretical Foundations and Applications of Deep Generative Models, 2018.
[12] Roman Garnett, Bayesian Optimization. Cambridge University Press, 2023.
[13] Robert A Gatenby and Edward T Gawlinski. A reaction-diffusion model of cancer invasion. Cancer Research, 56(24):5745-577) 1996.
[14] Mare G Genton. Classes of kernels for machine learning: statisties perspective. Journal of Machine Learning Research, 29993)299-311, 2001.

---

[15] Jonathan Gordon, Wessel P Bruinsma, Andrew YK Foong, James Requeima, Yann Dubois, and Richard E Turner. Convolutional conditional neural processes. In International Conference on Learning Representations, 2020.
[16] Nikolaus Hansen and Andreas Ostermeier. Completely derandomized self-adaptation in evolu- tion strategies. Evolutionary Computation, 9(2):159-19) 195, 2001.
[17] Nikolaus Hansen, yoshihikoueno, ARFI, Gabriela Kadlecová, Kento Nozawa, Luca Rolshoven, Matthew Chan, Youhei Akimoto, brieglhostis, and Dimo Brockhoff. CMA-ES/pycma: r3.3.0, January 2023. URL
7573532.
[18) Peter Holderrieth, Michael Hutchinson, and Yee Whye Teh. Equivariant learning of stochas- tie fields: Gaussian processes and steerable conditional neural processes. In International Conference on Machine Learning, pages 4297-4307, 2021.
[19] Makoto Kawano, Wataru Kumagai, Akiyoshi Sannai, Yusuke Iwasawa, and Yutaka Matsuo. Group equivariant conditional neural processes. In International Conference on Learning Representations, 2021.
[20] Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. In International Conference on Learning Representations, 2019.
[21] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[22] Shigeru Kondo and Takashi Miura. Reaction-diffusion model as a framework for understanding biological pattern formation. Science, 329(5999):1616-166, 2010.
[23] Risi Kondor and Shubhendu Trivedi. On the generalization of equivariance and convolution neural networks to the action of compact groups. In International Conference on Machine Learning, pages 2747-2755, 2018.
[24] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(1 36(11):227-223) 2324, 1998.
[25] Juho Lee. Yoonho Lee, Jungtack Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention based permutation-ines invarrian neural networks. In Pro- ceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pages 3744-3753, 2019.
[26] Sulin Liu, Xingyuan Sun, Peter Ramadge, and Ryan P Adams. Task-agnostic amortized inference of Gaussian process hyperparameters. In Advances in Neural Information Processing Systems, volume 33, pages 1444-21452, 2020.
[27] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proc. International Conference on Computer Vision, 2015.
[28]
Alfred J Lotka. Elements of Physical Biology. Williams & Wilkins, 1925.
[29] Zhou Lu, Hongming Pu, Feicheng Wang, Zhiqiang Hu, and Liwei Wang. The expressive power of neural networks: view from the width. In Advances in Neural Information Processing Systems, volume 30, pages 6231-6239, 2017.
[30] Stratis Markou, James Requeima, Wessel P Bruinsma, Anna Vaughan, and Richard E Turner. Practical conditional neural processes via tractable dependent predictions. In International Conference Learning Representations, 2022.
[31] Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-wwar meta learning via sequence modeling. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 16569- 16594, 2022.
[32] Eugene P Odum and Gary W Barrett. Fundamentals of Ecology. Thomson Brooks/Cole, 5th edition, 2005.
12

---

[33] Chris Olah, Nick Cammarata, Chelsea Voss, Ludwig Schubert, and Gabriel Goh. Naturally occurring equivariance in neural networks. Distill, 5(12):c00024-004 2020.
[34] Carl Edward Rasmussen and Christopher KI Williams. Gaussian Processes for Machine Learning. MIT Press, 2006.
[35] Fergus Simpson, Ian Davies, Vidhi Lalchand, Alessandro Vullo, Nicolas Durrande, and Carl Ed- ward Rasmussen. Kernel identification through transformers. In Advances in Neural Information Processing Systems, volume 34, pages 10483-10495, 2021.
[36] Jake Snell, Kevin Swersky, and Richard Zemel. Prototypical networks for few-shot learning. In Advances in Neural Information Processing Systems, volume 30, pages 4077-4087, 2017.
[37] Flood Sung, Yongxin Yang. Li Zhang, Tao Xiang, Philip HS Torr. and Timothy M Hospedales. Learning to compare: Relation network for few-shot learning. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1199-1208, 2018.
[38] Abiy Tasissa and Rongjie Lai. Exact reconstruction of Euclidean distance geometry problem using low-rank matrix completion. IEEE Transactions on Information Theory, (5243331443114,, 2019.
[39] Warren Torgerson. Multidimensional scaling: I. Theory and method. Psychometrika, 17(4):
401-419, 1952.
[40] Aimo Törn and Antanas Zilinskas. Global optimization, volume 350. Springer, 1989.
[41] Clément Vignac, Andreas Loukas, and Pascal Frossard. Building powerful and equivariant graph neural networks with structural message-passing. Advances Neural Information Processing Systems, volume 33, pages 14143-14155, 2020.
[42] Vito Volterra. Fluctuations in the abundance of species considered mathematically. Nature, 118(2972):558-566 1926.
[43] Edward Wagstaff, Fabian B Fuchs, Martin Engelcke, Ingmar Posner, and Michael A Osborne. On the limitations of representing functions sets. In International Conference Machine Learning, pages 6487-6494, 2019.
[44] Edward Wagstaff, Fabian B Fuchs, Martin Engelcke, Michael A Osborne, and Ingmar Posner. Universal approximation of functions on sets. Journal of Machine Learning Research, 23(151): 1 1-56, 2022.
[45] Naoki Wakamiya, Kenji Leibnitz, and Masayuki Murata. A self-organizing architecture for scalable, adaptive, and robust networking. In Nazim Agoulmine, editor, Autonomic Network Management Principles: From Concepts to Applications, pages 119-140. Academic Press, 2011.
[46] Qi Wang and Herke Van Hoof. Model-based meta reinforcement learning using graph structured surrogate models and amortized policy search. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 23055-23077, 2022.
[47] Andrew Gordon Wilson, Zhiting Hu, Ruslan Salakhutdinov, and Eric Xing. Deep kernel learning. In International Conference on Artificial Intelligence and Statistics, pages 370-378, 2016.
[48] Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabás Póezos, Ruslan Salakhutdinov, and Alexander J Smola. Deep sets. In Advances in Neural Information Processing Systems, volume 30, pages 3391-3401, 2017.
13

---

Appendix
In this Appendix, we include our methodology. mathematical proofs, implementation details for reproducibility, additional results and extended explanations omitted from the main text.
Contents
A Theoretical proofs
15
A.I Definition of equivariance for prediction maps
15
A.2 Proof that RCNPs are equivariant
16
A.3 Proof that RCNPs are context-preserving
16
Methods
18
B.1
Training procedure, error bars and statistical significance
18
B.2 Experimental details and reproducibility
19
B.3 Details on Figure
19
C Computational cost analysis
19
C.1 Inference time analysis
19
C.2 Overall training time
20
D Details of synthetic regression experiments
20
D.1 Models
D.2 Data
21
D.3 Full results
22
Details of Bayesian optimization experiments
23
E.1
Models
23
E.2 Data
24
E,3 Bayesian optimization
25
E,4
Computation time
26
Details of Lotka-Volterra model experiments
26
F.1
Models
26
F2
Data
................................
28
F3 Full results
................................
28
F.4
Computation time...
29
G Details of Reaction-Diffussin model experiments
29
G.I
Models
29
G.2 Data
29
G.3 Additional results
31
G.4 Computation time
H Additional experiments
31
14

---

H.1
Autoregressive CNPs
31
H.1.1
Computation time
32
H.2
Rotation equivariance
32
H.2.1
Computation time
33
H.3
Image datasets...
34
H.3.1
Setup
................................
34
H.3.2
Experiment I: Centered images (no translation)
H.3.3
Experiment 2: Translated images
35
H.3.4
Computation time
38
A Theoretical proofs
In this section, we provide extended proofs statements and theorems the main text. First, we show that our definition of equivariance for the prediction map is equivalent to the common definition of equivalent maps (Section A. Then, we provide full proofs that RCNPs are equivariant (Section A.2) and context-preserving (Section A.3) prediction maps, as presented in the theorems in Section of the main text.
A.1 Definition of equivariance for prediction maps
We show that the definition of equivariance of prediction map, which is introduced in the main text based the invariance of the representation function (Eq. 3 in the main text), is equivalent to the common definition of equivariance for generic mapping.
First, need extend the notion of group action of group T into Rd, for d E N. We define two additional actions, (i) for т E T and f a function from Rd into R2, and (ii) for a couple yx, E Rd R Following [15]:
yx, (rx,y),
Tf(x) = f x).
These two definitions can be extended for n-uplets X and Y:
(X, Y) = (rX,Y)
Tx(X) = f(r fxx
In the literature, equivariance is primarily defined for map F Z E X' x Rd + Fz E F(X, Rd) by: + FFZ = Fiz-
(S1)
Conversely, in the main text (Section 2), defined equivariance for prediction map through the representation function: prediction map ग with representation function r is -equiivariant if and only if for all r E T,
(S2) Eqs. S1 and $2 differ in that they deal with distinct objects. The former corresponds to the formal definition of an equivariant map, while the latter uses the representation function, which simpler to work with for the scope of our paper.
r((X, Y), r((rX, rX,
To show the equivalence of the two definitions (Eqs. S1 and S2), we recall that a prediction map defined through its representation r:
Px(X*) = x(- Z, x*) = p(- (r((((( Y),")), (X,Y). Y). Then, starting from Eq. S1 applied the prediction map.
where Z Z=XX,x
VX*, Pxxx = P(X = vX*, p(- r((rX, Y), x²)) = pl r(X, Y), Xx))
- r((rX, Y), X-) -xx
Noting that 7 (x(X, (). X*) = r((X, Y). -X"), by applying -1 to each side of the last line of the equation above, and swapping sides, we get:
P,X,Y)= ,Px + VX r(X,, Y). x*) r((x ((, 7),xX),
that is, Eqs. S1 and S2 are equivalent.
15

---

A.2 Proof that RCNPs are equivariant
We recall from Eq. 3 in the main text that a prediction map т with representation function is T Tequivvariant with respect to group T of transformations, T A' + A', if and only if for all T E T:
r((X, X") - r((>X, rX,+
where 7x = (xx and 7X is the set obtained by applying T to all elements of X.
Also recall that a prediction map and its representation function relational with respect a comparison function g : X' x X - R.ome and only if r can be written solely through the set comparisons defined in Eq. 8 of the main text:
r((X, X*) r((((X, Y), (X, Y)), (((,))())))))))) x),x), gX"" (X, g x").
We can now prove the following Lemma from the main text.
Lemma (Lemmaa 4.4), Let ı be a prediction map. T a transformation group, and g a comparison function. If is relational with respect 9 and g is T Tinvvariant, then π is Tequivariant.
Proof. From Eq. 8 the main text, if g is -iinvariant, then for a target set X and T E T
g((X, Y), = (X(, yn)) {(((T,., xx yn)] = g((rX, y),xX*)
for all S n I N and 1 S m < M. A similar equality holds for all the comparison sets in the definition of relational prediction map (see above). Since is relational, and using the equality above, we can write its representation function r' as: r((X, Y) = r Y), (X, Y)). = =r(g((TX, Y), (rX, Y)). Y), x),xX)),((x3 (rX, 9)), =r((7X, rX*).
(X"" (X, 9x)
Thus, is Teqquivariant.
From the Lemma, we can prove the first main result of the paper, that RCNPs are equivariant prediction maps for transformation group T provided suitable comparison function g. Proposition (Proposition 4.5). Let g be the comparison function used in a RCNP, and T a group of transformations. If is -invariant, the RCNP is T -equivariant.
Proof. all RCNP models, elements the context and target sets are always processed via rela- tional encoding. By definition (Eqs. 5 and 7 in the main text), the relational encoding is written solely as a function of ((X,, Y). (X, Y)) and g((X, X)), so the same holds for the representation function r. Thus, any RCNP is relational prediction map with respect to g. and it follows that is also a -qquiaarian prediction map according to Lemma 4.4.
A.3 Proof that RCNPs are pontext-preserving
We prove here that RCNPs preserve information in the context set. We will do so by providing construction for function that reconstructs the context set given the relationally encoded vector Pm (up to transformations). As disclaimer, the proofs in this section should be taken as 'existence proofs' but the provided construction is not practical. There are likely better context-preserving encodings, as shown by the empirical performance of RCNPs with moderately sized networks. Future theoretical work should provide stronger bounds on the size of the relational encoding, as recently established for DeepSets [43, 44].
For convenience, we recall Eq. of the main text for the relational encoding provided a comparison function g:
Pm Piull (x
X. Y))
+ fo (g(Xn, x"n), R,nn) R = (g(Xn, ) Nnr), 10,t'==
where R = ((X, Y), Y)) is the relational matrix. First, we show that from the relational encoding Piul we can reconstruct the matrix R, modulo permutations of rows or columns as defined below.
16

---

Definition A.1. Let R. R' E R RNN, The two matrices are equal modulo permutations, denoted R = R', if there a permutation (y) such that R R (m)))(nn)) for all S n, n' - N.
In other words, R = R' R' is equal to R after an appropriate permutation of the indices of both rows and columns. Now proceed with the reconstruction proof.
Lemma A.2. Let g: X x X + Rdame be a comparison function, (X. Y) a context set, X a target set, R the relational matrix and Piull the relational encoding as per Eq. 5 in the main text. Then there is a reconstruction function h. for r (pr..... pMs), such that h(r) = R' where R' R.
Proof. For this proof, it is sufficient to show that there is a function h applied to local representation of any context point, Pm = Ptull (x (X, Y)), such that h(pm) = R', with R' R. In this proof, we also consider that all numbers in computer so with we denote the set of real numbers physically representable
represented up to numerical precision,
a chosen floating point representation. As final assumption, without loss of generality, we take yn * Yn' for * n'. If there are non-unique elements yn in the context set, in practice we can always disambiguate them by adding small jitter.
Let fr R + Rd be the relational encoder of Eq. 5 of the main text, a neural network parametrized by Or. where R = Rdwr 12d, Here, we also include in O, hyperparameters defining the network architecture, such as the number of layers and number of hidden units per layer. We pick ded large enough such that we can build a one-hot encoding of the elements of R, i.e., there is an injective mapping from R to 10.1)dd such that one and only one element of the out- put vector is 1. This injective mapping exists, since R is discrete due to finite numerical pre- cision. We select 0, such that fr approximates the one-hot encoding up to arbitrary precision, which is achievable thanks to universal approximation theorems for neural networks [29]. Thus, Pm = Dn,n' (x ) R.,n) ) is a (astronomically large) vector with ones at the elements that denote (9(X,, ),,,,, Raw') for 1 S n, N.
Finally, we build reconstruction function h that reads Pm, the of one-hot encoded vec- tors, mapping each 'one', based on its location, back to the corresponding input (g(x,, ) Rea'). We have then all the information build matrix R' indexed by y, that
Ryy = Run',
with y Yn,y' Yn': for some y, YY,15, S.N.
Since we assumed ya + Yne for n o n', R' is unique up to permutations of the rows and columns, and by construction R' R.
We now define the matrix
Q Qunnn' = (g(x,, x"n); (Xn,
= (,n),,"",n, R,n')
for 1 S m S. M and 1 1 nnn S. N. We recall that a comparison function is context- preserving with respect to a transformation group T if for any context set (X, Y) and target set X', for all 1 < m < M, there is submatrix Q' C Qm:: a reconstruction function 7, and transformation T such that > (Q') = (rX. Y). We showed in Section 4.2 of the main text that the distance comparison function Gdist is prrsexi-pr with respect to the group of rigid transformations: and that the difference comparison function gaif is context-preserving with respect to the translation group. Similarly, family of functions (XX, there exists E θ, dhop E N, a reconstruction function 7, and a transformation T E T such that (X,, Y), X*) (TX, Y).
Y), (X), Rd is contexi- prreext-preserving under T if
fuutitto
With the construction and definitions above, we can now show that RCNPs are in principle context- preserving. In particular, we will apply the definition of context- prreextin functions h to the class of representation functions of the RCNP models introduced in the paper, for which drep Mdat since r = (Pi..... Par). ). The proof applies both to the full RCNP and the simple (diagonal) RCNP, for an appropriate choice of the comparison function g.
Proposition (Proposition 4.8). Let T be transformation group and the comparison function used in FulIRCNP. If g is context preserrigg with respect to T, then the representation function r of the FulIRCNP is context-preserving with respect to T.
Proof. Given Lemma A.2, for any given m. we can reconstruct Qm: = (..., R) (modulo per- mutations) from the relational encoding. Since g is context- prrtext-nrsseriing with respect to T, we can
17

---

reconstruct (X, Y) from R modulo transformation T E T. Thus, Ptull is context- pressering with respect T and so is r.
Proposition (Proposition 4.9). Let T be the translation group and 9an the difference comparison function. The representation of the simple RCNP model with 9an is context-preserving with respect to T.
Proof. As shown in the main text, the difference comparison function is context-preserving for Q Qmazi n=1 (g( X,, xm). Run)n for a given m. other words, we can reconstruct (X. Y) given only the diagonal of the R matrix and the comparison between a single target input x and the elements of the context set, (a)) Following construction similar the one in Lemma A.2, we can map Pm back to Q Then since is context-preserving with respect to Q we can reconstruct (X, Y) up to translations.
N
A potentially surprising consequence of our context-preservation theorems is that in principle no information lost about the entire context set. Thus, principle RCNP with sufficiently large network is able to reconstruct any high-order interaction of the context set, even though the RCNP construction only builds upon interactions of pairs of points. However, higher-order interactions might be harder to encode in practice. Since the RCNP representation is built on two-point interactions, depending on the network size, may be harder for the network to effectively encode the simultaneous interaction of many context points.
B
Methods
B.1
Training procedure, error bars and statistical significance
This section describes the training and evaluation procedures used in our regression experiments. The training procedure follows [3, Appendix GJ. Stochastic gradient descent is performed via the Adam algorithm [21] with a learning rate specified in each experiment. After each epoch during model training. each model undergoes validation using pre-generated validation set. The validation score a confidence bound based on the likelloodo values. Specifically, the mean (Hval) and standard deviation (oval) of the likelikel values over the entire validation set used to calculate the validation score as Pval v. /N.ai where Nval is the validation dataset size. The validation score compared to the previous best score observed within the training run, and if the current validation score is higher, the current model parameters (e.g., weights and biases for neural networks) saved. The models are trained for set number of epochs in each experiment, and when the training is over, the model is returned with the parameters that resulted in the highest validation score.
To ensure reliable statistical analysis and fair comparisons, we repeat model training using the above procedure ten times in each experiment, with each run utilizing different seeds for training and validation dataset generation and model training. In practice, generated ten seed sets that cach training run with the same model utilizes different seeds, but we maintain consistency by using the same seeds across different models in the same experiment. Each model is then represented with ten training outcomes each experiment, and we evaluate all training outcomes with a fixed evaluation set evaluation sets.
We evaluate the methods based on two performance measures: (i) the log-likelihood, normalized with respect to the number of target points, and (ii) the Kullback-Leibler (KL) divergence between the posterior prediction map π and a target πτ, defined as
π (xx) x² (x)
(L(T||**)>
((x) log dr,
assuming that π is absolutely continuous to x with respect to a measure À (usually, the Lebesgue measure). We report the KL normalized with respect to the number of target points.
Throughout our experiments, we rely on the KL divergence to evaluate the experiments with Gaussian process (GP) generated functions (" ("Gaussian'), for which we can compute the ground-truth target m (the multivariate normal prediction of the posterior GP). We use the normalized log- likelihood for all the other
)) cases. For each model and evaluation set, report both the average score
18

---

calculated across training runs and the standard deviation between scores observed in each training run.
We compare models based on the average scores and also run pairwise statistical tests to identify models that resulted in comparable performance. The results from pairwise comparisons used throughout the paper in the results tables to highlight in bold the models that are considered best in each experiment. Specifically, we use the following method determine the best models:
First, highlight in bold the model (A) that has the best empirical mean metric.
Then, for each alternative model (B), we run a one-sided Student's 's (-lest for paired samples, with the alternative hypothesis that model (B) has higher (better) mean than model (A), and the null hypothesis that the two models have the same mean, Samples are paired since they share random seed (i.c., the same training data).
The models that do not lead statistically significant (p- < 0.05) rejection of the null hypothesis are highlighted in bold as well.
B.2 Experimental details and reproducibility
The experiments carried out in this work used the open-source neural processes package released with previous work [3]. The package is distributed under the MIT license and available at https: //gi thub. hom/wessebb/ weell//nnuuaaapp http//euuallppseseeeeeeeee [1]. We extended the package with the proposed relational model architectures and the new experiments considered in this work. Our implementation weabb/eelationnn
of RCNP is available https:/ / giithii thub. com/
neeural-poocesees
For the purpose of open science and reproducibility, all training and evaluation details are provided in the following sections of this Appendix as well as in the Github repository linked above.
B.3 Details on Figure
In Figure in the main text, compare CNP and RCNP models. The RCNP model used in this example utilizes a stationary kernel to encode translation equivariance. The encoding in the CNP model and the relational encoding the RCNP model are 256-dimensional, and both models used an encoder network with three hidden layers and 256 hidden units per layer. In addition, both models used a decoder network with six hidden layers and 256 hidden units per layer. We used the CNP architecture from previous work [3] and modeled the RCNP architecture on the CNP architecture. The encoder and decoder networks use ReLU activation functions.
We trained the neural process models with datasets generated based on random functions sampled from a noiseless Gaussian process prior with Matérn- 3 covariance function with lengthscale f = 0.5. The datasets were generated by sampling 1-30 context points and 50 target points from interval 2- 2]. The neural process models were trained for 20 epochs with 244 datasets in each epoch and learning rate 3- 10 The datasets used training the CNP and RCNP models were generated with the same seed.
Figure shows model predictions in two regression tasks. First, we visualized the predicted mean and credible bounds calculated 1.96 times the predicted standard deviation in the range [-2, 2] given 10 context points. This corresponds to the standard interpolation (INT) task. Then we shifted the context points and visualized the predicted mean and credible bounds in the range [2, 6]. This corresponds to an out-of-input- coa--innuuddistributtonn (OOID) task. The Gaussian process model used as a reference had Matérn- 5 covariance function with lengthscale e 0.5, the same process used train the CNP and RCNP models.
C Computational cost analysis
C.1 Inference time analysis
In this section, we provide a quantitative analysis of the computational cost associated with running various models, reporting the inference time under different input dimensions d, and varying context set sizes N. We compare our simple and full RCNP models with the CNP and ConvCNP families.
19

---

All results are calculated on an Intel Core 17- CPU, under the assumption that these models can be deployed on devices or local machines without GPU access.
Table S1 shows the results with different input dimensions dz- We set the size of both context and target sets 20. Firstly, we can the runtime of convolutional models increase significantly as they involve costly convolutional operations, and cannot be applied in practice to dz > 2. Conversely, the runtime cost remains approximately constant across de for the other models. Secondly, our RCNP models have inference time close to that of CNP models, while FulIRCNP models are slower (although still constant in terms of dz). Since context set size is the main factor affecting the inference speed of our RCNP models, we delve further into the computational costs associated with varying context set sizes N, keeping the input dimension dæ 1 and target set size M 20 constant (Table S2), Given that RCNP has cost of O(NM), , we observe small, steady increase in runtime. Conversely, the Fu models, with their O(N2M) cost, show an approximately quadratic surge in runtime as N increases, expected.
Table S1: Comparison of computational costs across different models under different input dimen- sions dg- The size of both context sets (N) and target sets (M) is set to 20. We report the mean value and (standard deviation) derived from 50 forward passes exccuted on various randomly generated datasets.
Runtime (x 10 3s)
d,
dg 2
d,
d,
RCNP
2.35 (0.39)
2.29 (0.28)
2.28 (0.32)
2.25 (0.27)
RGNP
2.31 (0.30)
2.41 (0.34)
2.31 (0.26)
2.50 (0.40)
FuIlRCNP
8.24 (0.46)
9.09 (1.78)
8.27 (0.56)
9.08 (1.72)
FulIRGNP
8,40 (0.56)
(1.85)
8.37 (0.62)
9.20 (1.76)
ConvCNI
4,46 (0.38)
31.19 (2.43)
ConvGNP
4.51 (0.48)
38.04 (2.31)
CNP
1,86 (0.30)
2.10 (0.38)
2.14 (0.30)
2.09 (0.29)
GNP 2.01 (0.27)
2.09 (0.26)
2.20 (0.26)
2.10 (0.32)
Table S2: Comparison of computational costs across different models under different context set sizes N. We maintain the input dimension at dg = 1, and a fixed target set size of M = 20 We report the mean value and (standard deviation) derived from 50 forward passes executed on various randomly generated datasets.
Runtime (x 10 3x)
N = 10
N = 20
N = 50
N = 100
RCNP
1,90 (0.26)
2.34 (0.48)
2.40 (0.30)
2,90 (0.46)
RGNP
2.10 (0.26)
2.53 (0.34)
2.55 (0.22)
2.93 (0.37)
FNIIRCNP
2.83 (0.33)
8.80 (1.59)
38.03 (4.71)
154.17 (16.04)
FuIIRGNP
3,08 (0.34)
8.68 (1.02)
38.33 (5.88)
161.07 (19.54)
CNnVCNP
4.47 (0.43)
4.47 (0.51)
5.10 (0.62)
5.83 (0.88)
ConvGNP
5.10 (0.54)
4.55 (0.35)
5.26 (0.50)
4.70 (0.46)
CNP
1.95 (0.27)
1.86 (0.28)
1.87 (0.21)
1.97 (0.19)
GNP
2.48 (0.70)
2.03 (0.19)
2.17 (0.48)
2.11 (0.23)
C.2 Overall training time
For the entire paper, we conducted all experiments, including baseline model computations and preliminary experiments not included in the paper, a GPU cluster consisting of mix of Tesla P100, Tesla V100, and Tesla A100 GPUs. We roughly estimate the total computational consumption to be around 25 000 GPU hours. A more detailed evaluation of computing for each set of experiments is reported when available in the following sections.
D Details of synthetic regression experiments
We report details and additional results for the synthetic regression experiments from Section 1 of the main text. The experiments compare neural process models trained on data sampled from both Gaussian and non Gaussian synthetic functions, where 'Gaussian' refers to functions sampled from
20

---

Gaussian process (GP). Our procedure follows closely the synthetic regression experiments presented in [3].
D.1 Models
We compare RCNP models with the CNP and ConvCNP model families in regression tasks with input dimensions dz = (1, 2, 3, 5, 10). The CNP and GNP models used in this experiment encode the context sets as 256-dimensional vectors when dæ < 5 and as 128-dimensional vectors when d 2 5. Similarly, all relational models including RCNP, RGNP, FuIIRCNP, and FuIlRGNP produce relational encodings with dimension 256 when dg < 5 and dimension 128 when d, 2 5. The encoding network used in both model families to produce the encoding or relational encoding three-layer MLP, featuring 256 hidden units per layer for de 5 and 128 for dg > 5. We also maintain the same setting across all CNP and RCNP models in terms of the decoder network architecture, using a six-layer MLP with 256 hidden units per layer for d₂
and 128 for de 2 The encoder and decoder networks use ReLU activation functions. Finally, the convolutional models ConvCNP and ConyGNP, which are included in experiments where dN {1, 2}. are employed with the configuration detailed in [3, Appendix FJ, and GNP, RGNP, FalIRGNP, and ConvGNP models all use linear covariance with 64 basis functions.
Neural process models are trained with datasets representing regression task with context and target features. The datasets used in these experiments were generated based on random functions sampled from Gaussian and non-Gaussian stochastic processes. The models were trained for 100 epochs with 21 datasets in each epoch and learning rate 3 10-4, The validation sets used in training included 2 datasets and the evaluation sets used to compare the models in interpolation (INT) and caaf-iiuudddiirbuutte (OOID) tasks included 2 datasets each.
D.2 Data
We used the following functions to generate synthetic Gaussian and non- Gaussan data:
Exponentiated quadratic (EQ). We sample data from a GP with an EQ covariance function:
f GP (0, exp
[|lx
x11 ))
(S3)
262
where is the lengthscale. Matérn- We sample data from a GP with a Matérn- & covariance function:
5 )00) exp (-v5r)) f gp (o, 1+ Vor (1+vo++
(S4)
where ||x and is the lengthscale.
Weakly periodic. We sample data from GP with weakly periodic covariance function:
f g~ (0,
1 26
2 sin [xxxxx) (xxx)))]l
(0, exp
x
x11
(S5)
where la is the lengthscale that decides how quickly the similarity between points in the output of the function decays as their inputs move apart, lp determines the lengthscale of periodic variations, and p denotes the period.
Sawtooth. We sample data from sawtooth function characterized by stochastic frequency, orientation, and phase as presented by:
f ~~ (x, u)2 + o mod 1,
(S6)
where a ~ (S2) is the frequency of the sawtooth wave, the direction of the wave is given as ~ ((( E Rdd, ||x||2 = 1)), and & ~ ((0 1) determines the phase.
Mixture. We sample data randomly from either one of the three GPs or the sawtooth process, with each having an equal probability to be chosen.
21

---

In this set of experiments, we evaluate the models using varying input dimensions d, (1, 3, 5, 10). To maintain a roughly equal level of difficulty across data with varying input di- mensions dz. the hyperparameters for the above data generation processes are selected in accordance with dz:
1
1
{ Vz.
la= = vdz. = dg. p p= v= Ω
(S7)
2//2
vdz
For the EQ, Matérn- 3. and weakly periodic functions, we additionally add independent Gaussian noise with x² = 0.05.
The datasets representing regression tasks were generated by sampling context and target points from the synthetic data as follows. The number of context points sampled from each EQ, Matérn- 5 or weakly periodic function varied uniformly between and 30 dz., and the number of target points was fixed to 50 dz. Since the datasets sampled from the sawtooth and mixture functions represent more difficult regression problems, we sampled these functions for 1-30 context points when dæ = 1 and 1-50- d₂ context points when dg > 1, and we also fixed the number of target points to 100 dg.
2:
All training and validation datasets were sampled from the range X' = |1 while evaluation sets were sampled in two ways. To evaluate the models in an interpolation (INT) task, we generated evaluation datasets by sampling context and target points from the same range that was used during training, Mest = [-2, 2] To evaluate the models in an out-of- of--nnptt--aae anffiipptddssiiibbet distribution (OOID) task, we generated evaluation datasets with context and target points sampled from range that is outside the boundaries of the training range, specifically Xeu d.
D.3
Full results
The results presented in the main text (Section 5. compared the proposed RCNP and RGNP models to baseline and convolutional CNP and GNP models in selected synthetic regression tasks. The full results encompassing a wider selection of tasks and an extended set of models are presented in Table S3, S4, S5, S6, S7. We constrained the experiments with EQ and Matérn- - tasks (Table S3 and S4) to input dimensions de = {1, 2, 3, 5] owing to the prohibitive computational memory costs associated with the FuP model included in these experiments. The other experiments (Table S5- S7) were repeated with input dimensions d₂ = {1, 3. 5. 103. Our RCNP models demonstrate consistently competitive performance when compared with ConvCNP models in scenarios with low dimensional data, and significantly outperform the CNP family of models when dealing with data of higher dimensions.
Table S3: Synthetic (EQ). Comparison of the interpolation (INT) and cwwrroofinuttdistriboo (OOID) performance of our RCNP models with different CNP baselines on EQ function with varying input dimensions. We use normalized Kullback- Leibler divergences our metric and show mean and (standard deviation) obtained from 10 runs with different seeds. Trivial refers to a model that predicts a Gaussian distribution utilizing the empirical mean and standard deviation derived from the context outputs. "F" denotes failed attempts that yielded very bad results. Missing entries could not be run. Statistically significantly best results are bolded.
d.
d,
d,
d,
INT
OOID
INT
OOID
INT
OOID
INT
OOID
RCNP (sta)
0.22 (0.00)
0.22 (0.00)
0.26 (0.00)
0.26 (0.00)
0.40 (0.01)
0.40 (0:01)
0.45 (0.00)
0.45 (0:00)
RGNP (sta)
0.01 (0.00)
0.01 (0.00)
0.03 (0.00)
0.03 (0.00)
0.05 (0.00)
0.05 (0.00)
0.11 (0.00)
0.11 (0.00)
FullRCNP (iso)
0.22 (0.00)
0.22 (0.00)
0.26 (0.00)
0.26 (0.00)
0.31 (0.00)
0.31 (0.00)
0.35 (0.00)
0.35 (0.00)
FuIRGNP (isn)
0.03 (0.00)
0.03 (0.00)
0.08 (0.00)
0.08 (0.00)
0.14 (0.00)
0.14 (0.00)
0.25 (0.00)
0.25 (0.00)
ConvCNP
0.21 (0.00)
0.21 (0.00)
0.22 (0.00)
0.22 (0.00)
ConvGNP
0.00 (0.00)
0.00 (0.00)
0.01 (0.00)
0.01 (0.00)
CNP
0.25 (0.00)
2.21 (0.70)
0.33 (0.00)
4.54 (1.76)
0.44 (0.00)
3.30 (1.55)
0.57 (0.00)
1.22 (0.09)
GNP
0.02 (0.00)
0.05 (0.00)
2.25 (0.61)
0.09 (0.01)
2.54 (1.44)
0.19 (0.00)
0.74 (0.02)
Trivial
1.03 (0.00)
1.03 (0.00)
1.13 (0.00)
1.13 (0.00)
1.12 (0.00)
1.12 (0.00)
1.03 (0.00)
1.03 (0.00)
We also conducted an additional experiment to explore all the possible combinations of comparison functions with the RCNP and FullRCNP models on EQ and Matérn- tasks. Table S8 provides a comprehensive view of the results for EQ tasks across three input dimensions. The empiri- cal results support our previous demonstration that RCNP models are capable of incorporating translation-cquivvrrincee with no loss of information: the stationary versions of the RCNP model
22

---

Table S4: Synthetic (Matérn-3). Comparison of the interpolation (INT) and out-ef-inputtdistributt (OOID) performance of our RCNP models with different CNP baselines on Matérn- function with varying input dimensions. We use normalized Kullback-Leibler divergences as metric and show mean and (standard deviation) obtained from 10 runs with different seeds. Trivial refers to a model that predicts Gaussian distribution utilizing the empirical mean and standard deviation derived from the context outputs. "F" denotes failed attempts that yielded very bad results. Missing entries could not be run. Statistically significantly best results are bolded.
d, 3
INT
OOID
INT
OOID
INT
OOID
INT
OOID
RCNP (sta)
0,25 (0.00)
0.25 (0.00)
0.30 (0.00)
0.30 (0.00)
0.39 (0.00)
0.39 (0.00)
0.35 (0.00)
0.35 (0.00)
RGNP (sta)
0.01 (0:00)
0.01 (0.00)
0.03 (0.00)
0.03 (0.00)
0.05 (0.00)
0.08 (0.00)
0.11 (0.00)
0.11 (0.00)
FullRCNP riso)
0.25
0.25 (0.00)
0.30 (0.00)
0.30 (0:00)
0.32 (0.00)
0.32 (0.00)
0.29 (0.00)
0.29 (0.00)
FulRRGNP (iso)
0.03 (0.00)
0.03 (0.00)
0.09 (0.00)
0.09 (0.00)
0.16 (0.00)
0.16 (0.00)
0.21 (0.00)
0.21 (0.00)
ConvCNP
0.24 (0:00)
(0.00)
0.26 (0.00)
0,26 (0.00)
ConyGNP
0.00 (0.00)
0,00 (0.00)
0.01 (0:00)
0.01 (0.00)
CNP
0.29 (0.00)
F
0.39 (0.00)
6.75 (2.72)
0.46 (0.00)
1.75 (0.42)
0.47 (0.00)
0.93 (0.02)
GNP
0.02 (0.00)
2.96 (1.77)
0.07 (0.00)
1,86 (0.26)
0.11 (0.00)
1.23 (0.17)
0.19 (0.00)
0.62 (0.02)
Trivial
1.04 (0.00)
1.04 (0.00)
1.06 (0.00)
1.06 (0.00)
0.98 (0.00)
0.98 (0.00)
0.79 (0.00)
0.79 (0.00)
Table S5: Synthetic (weakly-periodic). Comparison of the interpolation (INT) and out- in-itt distribution (OOID) performance of our RCNP models with different CNP baselines on weakly- periodic function with varying input dimensions. We use normalized Kullback-Leibler divergences as our metric and show mean and (standard deviation) obtained from 10 runs with different seeds. Trivial refers to model that predicts Gaussian distribution utilizing the empirical mean and standard deviation derived from the context outputs. "F"" denotes failed attempts that yielded very bad results. Missing entries could not be run. Statistically significantly best results are bolded.
d_
2
d_ 3
4.
10
ODID
INT
0DID
OOID
INT
ODID
RCNP (sta)
0.24 (0:00)
(0.00)
0.24 (0.90)
(0.00)
0.28 (0.00)
0.33 (0.00)
0.31 (0.00)
(0.00)
(0.00)
RGNP (sta)
0:03 (0.00)
0.00 (0.00)
0.05 (0.00)
0.05 (0.00)
0.05 (0.00)
0.05 (0.00)
0.08 (0.00)
0.08 (0.00)
0.11 (0.00)
0.11 (0.00)
CNP
0.21 (0.00)
0.21 (0.00)
0.20 (0.00)
0.20.0..0 (0.00)
CoEYGNP
0.01 (0.00)
0.01 (0.00)
0.02 (0.00)
0,02 (0.00)
CNP
0.31 (0.00)
2.88 (0.91)
0.39 (0.00)
1.81 (0.43)
0,39 (0.00)
1.58 (0.50)
0.42
2.20 (0.81)
0.75 (0.00)
1.00 (0.11)
GNP
0.06
F
(0.00)
(0.76)
0.08
1.47 (0:27)
0.11
0.62
(0.01)
(0.05)
Twivaa
0.78 (0:00)
0.78 (0.00)
0.81 (0.00)
(0.00)
0.80 (0.00)
0.80 (0.00)
0.37 (0.00)
0.77 (0.00)
(0.00)
(0.00)
deliver performance comparable to the stationary FullRCNP models. Conversely, since the RCNP model is not context- preexxe-rreeriing for rigid transformations, the isotropic versions of the RCNP model exhibit inferior performance in comparison to the isotropic FulIRCNPs.
E Details of Bayesian optimization experiments
We report here details and additional results for the Bayesian optimization experiments from Sec- tion 5.2 of the main text. We evaluated our proposed approach on a common synthetic optimization test function, the Hartmann function in its three and six dimensional versions [40, 185) Each of two test functions is evaluated on a d di-diennioonl hypercube X' = [0, 1]<=.
E.1 Models
We compared our proposals, RCNP and RGNP, against CNP and GNP as well as their attentive variants ACNP and AGNP. A GP with Matérn- 5 kernel additionally served as the "gold standard" to be compared against. We mostly followed the architectural choices discussed in Section D. However, we kept the number of hidden units per layer fixed at 256 and the relational encoding dimension at 128, irrespective of the dimensionality dz. The models were trained with a learning rate of 10 over up to 500 epochs with 2 datasets in each epoch. A validation set with 2 datasets was used to track the training progress and early stopping was performed if the validation score had not improved for 150 epochs.
'See httpp://www.cct ca/
html for their precise definitions and minima.
23

---

Table S6: Synthetic (sawtooth). Comparison of the interpolation (INT) and out-offiinuutddistibbiiii (OOID) performance of our RCNP models with different CNP baselines on sawtooth function with varying input dimensions. We use normalized logllikelshoods as our metric and show mean and (standard deviation) obtained from 10 runs with different seeds. Trivial refers to a model that predicts Gaussian distribution utilizing the empirical mean and standard deviation derived from the context outputs. "F" denotes failed attempts that yielded very bad results. Missing entries could not run. Statistically significantly best results are bolded.
d,
d: 2
de 3
d, 5
de 10
INT
00ID
INT
OOID
INT
OOID
INT
ODID
INT
OOID
RCNP
3,5 (0.06)
3,04 (0.06)
1,73 (0.03)
1.34
0.85 (0.01)
0.85
0.44 (0000)
0.44 (0.00)
0.35 (0.00)
0.75 (0,00)
RGINP
3.90
2.33
(0.01)
1,09
1,13 (0.05)
1.32 (0.03)
Com-CNP
(0,04)
3.64
3,66
CNI
(0:11)
3.97 (0.08)
(0.08)
3.99 (0.11)
CNP
2:25 (0.07)
-3.27
0.36 (0.28)
-0.37 (0.12)
-0.03 (0. 10)
4.22 (0.03)
0.27 (0.00)
-2.29 (0.67)
GNP
0.83 (0.04)
P
1.04 (0.09)
0.23 (0.13)
0.02 (0.05)
0.00 (0.04)
Triut
-0.27 (0.00)
-0.27 (0.00)
.0.99 (0.00)
-0.19 (0.00)
-0.19 (0.00)
-0.39 (0.00)
20.18 (0.00)
-0.18 (0.00)
-0.14 (0.00)
-0.14 (0.00)
Table S7: Synthetic (mixture). Comparison of the interpolation (INT) and out-of coff--nput-ttiee distribution (OOID) performance of our RCNP models with different CNP baselines on mixture function with varying input dimensions. We use normalized log- ligeliiiiissos as metric and show mean and (standard deviation) obtained from 10 runs with different seeds. Trivial refers to a model that predicts Gaussian distribution utilizing the empirical mean and standard deviation derived from the context outputs. "F" denotes failed attempts that yielded very had results. Missing entries could not be run. Statistically significantly best results are bolded.
d
d, 3
d, 5
d, 10
INT
00ID
INT
INT
INT
ODID
INT
OOID
RCNP
(0.01)
(0.00)
0.37 (0:00)
0,10 (0.00)
0.10
(0.03)
0.31
0.32
0..2
RGNP (sm)
(0.03)
0.34 (0.08)
B:46 (0.02)
0.46 (0.02)
(0:01)
0.37 (0:01)
(0.02)
0.04 (0.02)
.0.11 (0.02)
.0.11 (0.02)
CNP
0.38 (0.02)
0.38 (0.02)
(0.0))))))
0.63 (0.01)
ComGNN
0.49 (0.15)
0.49 (0.15)
0.87 (0.03)
0.57 (0005)
CNP
0.01 (0.01)
F
.0.22 (0.01)
(3.66)
-0.57 (0 11)
-2.55 (1.15)
.0.72 (0.08)
-1.71 (0.55)
-0.88 (0.00)
-1.15 (0
CNP
0.17 (0.01)
P
0.01 (0.01)
-0.17 (0.00)
-0.67 (0.05)
0.32 (0.00)
0.72 (0.00)
0,38 (0.10)
-0.75 (0.09)
-0.8 (0.00)
.0.7% (0.00)
10.11 (0.00)
42.8 (0.00)
-1.84 (0.00)
-0.84 (0.00)
-0.6 (0:00)
-0890 (1.00)
-0.87 (0.00)
-0.87 (0.00)
E,2 Data
The datasets used to train the neural process models were generated based on synthetic functions sampled from Gaussian process model with kernel k, with context and target set sizes as described in Section D.2. However, while the datasets used in the synthetic regression experiments were generated using fixed kernel setups, the current experiment explored sampling from a set of base kernels. To explore four different training regimes for each du = {3, 6). we changed how the GP kernel k<z) = (k)(z/E) with an output scale and a lengthscale & was chosen as follows.
(i) Matérn-fixed. The kernel k is a Matérn 2 kernel with fixed é = vd./4 and 0 = 1.
(ii) Matérn-sampled. The kernel k is Matérn- 3 kernel, with
(~ LN (Io /d. 0.5*). A A~ LN(, where LN is log-Normal distribution, i.e., log(E) follows standard-Normal distribution. (iil) Kernel-single. The kernel k is sampled as Ko U ((EQ, Matérn-,, Matérn- Matérn- ()) E ~ CN 0 0~ LN 1),
where UI is uniform distribution over the set. (iv) melnel- The kernel k is sampled as k k2 ka, K4 ~ l4 (f (ANA, EQ, Matérn 1, Matérn 3, Matérn-
f, ~ LV (1og(v log(v/o d=54),0.5)),
i (1,.....4) 4),
0, ~ LC((,, 1), jE (1,2),
A: = (e,) k2(l2) + (2kx((z) (l4)),
24

---

Table S8: An ablation study conducted on RCNPs and FullRCNPs using synthetic regression data. We evaluated both the 'stationary" (sta) version, which employs the difference comparison function, and the 'isotropic' (iso) version, utilizing the distance comparison function, across both RCNP and Fu models. The table provides metrics for both the interpolation (INT) and out-of- of--o----utt-- distribution (OOID) performance. We show mean and (standard deviation) obtained from 10 runs with different seeds. Statistically significantly best results are bolded.
EQ KL divergenoc())
Matém- 3 KL divergenee())
d2
d2
d =
d2
d₂ =
dg
RCNP (sta)
0.26 (0.00)
0.40 (0.01)
0.45 (0.00)
0.30 (0.00)
0.39 (0.00)
0.35 (0.00)
RGNP (sta)
0.03 (0.00)
0.05 (0.00)
0.11 (0.00)
0.03 (0.00)
0.05 (0.00)
0.11 (0.00)
RCNP (iso)
0.32 (0.00)
0.41 (0.00)
0.46 (0.00)
0.34 (0.00)
0.39 (0.00)
0.36 (0.00)
RGNP (iso)
0.12 (0.00)
0.17 (0.00)
(0.00)
0.12 (0.00)
0.17 (0.00)
0.24 (0.00)
FullRCNP (sta)
0.25
0..0
0.36 (0.00)
0.30 (0.00)
0.31 (0.00)
0.29 (0.00)
FalIRGNP (sta)
0.03
0.04
0.09 (0.00)
0.03 (0.00)
0.04 (0.00)
0.10 (0.00)
FullRCNP (iso)
0.26 (0.00)
0.31 (0.00)
0.35 (0.00)
0.30 (0.00)
0.32 (0.00)
0.29 (0.00)
FulIRGNP (iso)
0.08 (0.00)
0.14 (0.00)
0.25 (0.00)
0.09 (0.00)
0.16 (0.00)
0.21 (0.00)
RCNP (sta)
0.26 (0.00)
0.40 (0.01)
0.45 (0.00)
0.30 (0.00)
0.39 (0.00)
0.35 (0.00)
RGNP (sta)
0.03 (0.00)
0.05 (0.00)
0.11 (0.00)
0.03 (0.00)
0.05 (0.00)
0.11 (0.00)
RCNP (iso)
0.32 (0.00)
0.41 (0.00)
0.46 (0.00)
0.34 (0.00)
0.39 (0.00)
0.36 (0.00)
OOID
RGNP (iso)
(0.00)
0.17 (0.00)
0.30 (0.00)
0.12 (0.00)
0.17 (0.00)
0.24 (0.00)
FulIRCNP (sta)
0.25 (0.00)
0.30 (0.00)
0.36 (0.00)
0.30 (0.00)
0.31 (0.00)
0.29 (0.00)
FulIRGNP (sta)
0.03 (0.00)
0.04 (0.00)
0.09 (0.00)
0.03 (0.00)
0.04 (0.00)
0,10 (0.00)
FulIRCNP (iso)
0,26 (0.00)
0,31 (0.00)
0.35 (0.00)
0.30 (0.00)
0.32 (0.00)
0.29 (0.00)
FulIRGNP (iso)
0.08 (0.00)
0.14 (0.00)
0.25 (0.00)
0.09 (0.00)
0.16 (0.00)
0.21 (0.00)
where NA indicates that no kernel is chosen, i.e., term in the sum can consist of a single kernel or be missing completely. This setup is based on the kernel learning setup in [35].
See Section D.2 for the precise definition of each kernel. Note that each setting is strictly more general than the one before and entails its predecessor as a special case. Setting (iv) was reported in the main text.
E.3 Bayesian optimization
After training, the neural process models were evaluated in a Bayesian optimization task. Starting from common set of five random initial observations, X E u1(0, 1)d,, each optimization run performed fifty steps of requesting queries by maximizing the expected improvement acquisition function:
El (x) = F. [max([(x)) - Shess 0)]
where f(-) is the negative target function (as we wanted to minimize f), fbest the current observed optimum, and the expectation is calculated based on the neural process or GP model predictions. The acquisition function was optimized via PyCMA [17], the Python implementation of Covariance Matrix Adaptation Evolution Sirategy (CMA-ES) [16]. Each CMA-ES run was initialized from the best acquisition function value estimated over a random subset of 100 locations from the hypercube. After each step, the newly queried point and its function evaluation were used to update the surrogate model. For neural process models, this means that the point was added to the context set, whereas for the the GP model, this means that the model was also reinitialized and its hyperparameters refit via type-II maximum likelihood estimation. The neural process model parameters were kept fixed without retraining throughout the Bayesian optimization task.
The results reported in Figure 2 in the main paper were computed using a single pretrained neural process model and ten random Bayesian optimization restarts for each model. For the full results presented here, each neural process model was pretrained three times and used as a surrogate model in ten optimization runs initialized with different sets of five observations, giving us a total of thirty optimization runs with each model. See Figure S1 and Figure S2 for the results and their respective discussion.
25

---

E.4
Computation time
Each model was trained on single GPU, requiring about 90-120 min per run. The subsequent CMA-ES optimization and analysis took about 60-90 min running CPU with eight cores per experiment. The total runtime was approximately 280 of GPU time and approximately 7 on the CPUs.
Matém-fixed
Malérn-sarpaa
2.5
2.0
amor mihn (ix)) (ix))-far) Tear
1.5
imin
0.0
0.0
10
20
30
40
10
20
30
40
number of queries
number of queries
Kernel single
Kernel multiple
2.5
2.5
CNP
GNP
faiel
2.0
error Imin (tx) 1- fmal
2.0
AGNP
ACMP
RCMP
RGNP
1.5
1.5
1.0
1.0
0.5
0.5
0.0
0
20 number queries
querles
Figure S1: Ablations on Hartmann 3d. GNP struggled throughout all four pretraining scenarios. while its attentive and relational variants (AGNP, RGNP) were able to benefit from the additional flexibility, greatly improving their performance. CNP struggled to learn the flexibility in the context sets increases, while its attentive and relational counterparts adapted themselves and were able to converge optimal performance even in the most restrictive pretraining scenario.
F Details of Lotka-Volteree model experiments
We report here details and additional results for the Lotka-Volterra model from Section 5.3 of the main text.
F1
Models
The RCNP architecture used in this experiment is modeled on the multi-output CNP architecture proposed in previous work [3, Appendix FJ. The multi- CNP encodes the context set associated with each output as 256-dimensional vector and concatenates the vectors. The encoder network used in this model is an MLP with three hidden layers and 256 hidden units per layer, and the decoder network an MIP with six hidden layers and 512 hidden units per layer. The RCNP model replicates this architecture by encoding the targets with respect to the context set associated with each output as 256-dimensional vectors, concatenating the vectors, and using encoder and decoder networks that the same the encoder and decoder networks used in the CNP model. We note that the encoder
26

---

Matérn-
seremm--ampped
3.0
error Imin fonl
2.0
2.0
1.5
0.5
0.5
0.00
0.0g
number queries
number queries
Kernel-single
kemel-multtple
GP
CNP
2,5
2.5
ACNP
2.0
RCNP
RGNP
amrr Iml
1.0
1.0
0.0
0.0
20
40
60
80
20
40
60
80
number of queries
number of queries
Figure S2: Ablations on Hartmann 6d. As in the 3d case (see Figure S1), CNP and GNP struggled to learn all. Similarly, the advanced attentive and relational CNP models converged their optimum even with the most restrictive prior kernel set. As the flexibility increased, RGNP was able to benefit the most and improved both upon AGNP as well as their conditional variants (ACNP, RCNP).
network dimensions the same used in the synthetic regression experiments (Section D.I) while the decoder network has more hidden units per layer in the multi-output architectures. The encoder and decoder networks use ReLU activation functions.
The models evaluated in these experiments include RCNP, CNP, ACNP, and ConvCNP as well as RGNP, GNP, AGNP, and ConvGNP, RCNP and RGNP models share the architecture described above while the other models are evaluated with the architectures proposed in previous work [3]. However, in contrast to previous work, where CNP, GNP, ACNP, and AGNP models encoded the context associated with each output with separate encoder, here we used shared encoder RCNP, RGNP, CNP, GNP, ACNP, and AGNP. We made this choice based on preliminary experiments carried out with synthetic multi-output data, where we observed that using a shared encoder improved both RCNP and CNP results. We note that all models extended with nonlinear transformation to take into account that the population counts modeled in this experiment are positive [3].
The neural process models were trained with simulated data and evaluated with simulated and real data. The models were trained for 200 epochs with 2 datasets each epoch and learning rate 10-4 The validation sets used in training included 22 datasets also generated based on the simulated data. The evaluation sets generated based on simulated data included 22 datasets each and the evaluation sets generated based on real data 2 dataset each, and both simulated and real data were used to generate three separate evaluation sets to represent the three task types considered in this experiment: interpolation, forecasting, and reconstruction.
27

---

F.2 Data
The datasets used in this task included simulated and real datasets. The simulated datasets were constructed based on simulated predator and prey time series generated using stochastic Lotka- Volterra equations. The Lotka-Voltenn equations introduced independently by Lotka [28] and Volterra [42] model the joint evolution between prey and predator populations as
{ û cnt Swn ? = Sux YYU
where denotes the prey, denotes the predators, is the birth rate of the prey, /3 is the rate at which predators kill prey, ó the rate at which predators reproduce when killing prey, and is the death rate of the predators; the dot denotes a temporal derivative, i.e. au. .. In this work, we generated simulated data with a stochastic version proposed by [3].
{ du au df Buv dt + / dd
dv duv df yu dt + 02/6 a dW
where W and W are two independent Brownian motions that add noise to the population trajectories and Ø > 0 controls the noise magnitude.
We simulated trajectories from the above equation using the code released with the previous work [3]. The parameters used in each simulation were sampled from: a ~ ((0.2, 3 ~ (0.044 0.08), 5 ~ ((0.8, 1.2), 7 (0.04, 0.08), and a ~ U(.55 10). The initial population sizes were sampled from u, v ~ (5, 100), and the simulated trajectories were also scaled by a factor r (11 5). To construct a dataset, 150-250 observations were sampled from the simulated prey time series and 150-250 observations from the simulated predator series. The selected observations were then divided into context and target observations based on the selected task type which was fixed in the evaluation sets and selected at random in the training and validation data.
The datasets representing interpolation, forecasting, or reconstruction tasks were created based on the selected simulated observations as follows. To create an interpolation task, 100 predator and 100 prey observations, selected at random, were included in the target set, and the rest were included in the context set. To create a forecasting task, a time point was selected at random between 25 and 75 years, and all observations prior to the selected time point were included in the context set and the rest were included in the target set. To create reconstruction task, a random choice was made between prey and predators such that the observations from the selected population were all included in the context set while the observations from the other population were divided between the context and target sets by choosing at random time point between 25 and 75 years and including all observations prior to the selected time point in the context set and the rest in the target set.
The real datasets used to evaluate the neural process models in the sim-to-rea tasks were generated based on the famous hare-lynx data [32]. The hare and lynx time series include 100 observations at regular intervals. The observations were divided into context and target sets based on selected task types, and the same observations were used to generate separate evaluation sets for each task. The forecasting and reconstruction task was created with the same approach that was used with simulated data, while the interpolation task was created by choosing at random 1-20 predator observations and 1-20 prey observations that were included in the target set and including the rest in the context set.
F33 Full results
We evaluate and compare the neural process models based on normalized log-likelihood scores calculated based on each evaluation set and across training runs. We compare RCNP, CNP, ACNP, and ConvCNP models in Table S9 (a) and RGNP, GNP, AGNP, and ConvGNP models with 32 basis functions in Table 59 (b) and with 64 basis functions in Table 59 (c). We observe that the best model depends on the task, and comparison across the tasks indicates that RCNP or RGNP are not consistently better or worse than their convolutional and attentive counterparts. This is not the case with CNP and GNP, and we observe that the other models including RCNP and RGNP are almost always substantially better than CNP or GNP. The models with 64 basis functions performed on average better than the models with 32 basis functions, but we observed that some RGNP, GNP, and AGNP training runs did not converge and most GNP runs did not complete 200 epochs due to numerical errors when 64 basis functions were used. We therefore reported results with 32 basis functions in the main text.
28

---

Table S9: Normalized log-likelihood scores in the volka-Volterr experiments (higher is better). We compared models that predict (a) mean and variance and models that also predict a low-rank covariance matrix with (b) 32 basis functions or (c) 64 basis functions. The mean and (standard deviation) reported for each model are calculated based on 10 training outcomes evaluated with the same simulated (S) and real (R) learning tasks. The tasks include interpolation forecasting (FOR), and reconstruction (REC). Statistically significantly best results are bolded. RCNP and RGNP models perform par with their convolutional and attentive counterparts, and nearly always substantially better than the standard CNPs and GNPs, respectively.
INT (S)
FOR (S)
REC (S)
INT (R)
FOR (R)
REC (R)
RCNP
-3.57 (0.02)
-4.85 (0.00)
-4.20 (0.01)
4.24 (0.02)
-4.83 (0.03)
-4.55 (0.05)
(a)
ConvCNP
-3.47 (0.01)
-4.06 (0.00)
4.21 (0.04)
-5.01 (0.02)
-4.5 (0.05)
ACNP
-4.04 (0.06)
-4.87 (0.01)
-4.36 (0.03)
-4.18 (0.05)
-4.79 (0.03)
-4.48 (0.02)
CNP
4477 (0.00)
-4.88 (0.00)
-4.66 (0.00)
4.74 (0.01)
-4.81 (0.02)
-4.70 (0.01)
RGNP
-3.51 (0.01)
-.27 (0.00)
-3.76 (0.00)
-4.31 (0.06)
-4.77 (0.03)
-4.99 (0.11)
ConvGNP
-3.46 (0.00)
4.30 (0.00)
-3.67 (0.01)
4.19 (0.02)
4.61 (0.03)
-4.22 (0.11)
(b)
AGNP
(0.17)
4.35 (0.09)
-4.05 (0.19)
4.33 (0.15)
-448 (0.06)
-4,29 (0.10)
4.62
4.38 (0.04)
-4336 (0.07)
4.79 (0.03)
4.72 (0.04)
-4722 (0.04)
RGNP
(0.06)
4.15 (0.10)
-3.75 (0.10)
4.30 (0.04)
4.44 (0.02)
-4.38 (0.10)
ConvGNP
-3,46
4.15 (0.01)
-3.66
4.20
4.57 (0.05)
(0.22
AGNP
(0.10)
(0.04)
(011I)
(.11
4.41 (0.06)
-4.30 (0.08)
GNP
-4.63 (0.03)
4.36 (0.05)
(.005
4.82
4.75 (0.04)
-4.75 (0.03)
E4
Computation time
When the models were trained on a single GPU, each training run with RCNP and RGNP took around 6.5-7 hours and each training run with the reference models around 2-3 hours depending on the model. The results presented in this supplement and the main text required 120 training runs, and 10-15 training runs had been carried out earlier to confirm as a sanity check our ability reproduce the results presented in previous work. The multi-output models used in this experiment were additionally studied preliminary experiments using synthetic regression data. The computation time used in all the experiments is included in the estimated total reported in Section C.2.
G Details of Reaction-Diffusinn model experiments
We report here details and additional results for the Reaction-Diffusion model from Section 5.4 of the main text.
G.1 Models
We compared RCNP, RGNP, ACNP, AGNP, CNP, and GNP models trained using simulated data with input dimensions d., = {3, 4]. We mostly used the architectures described in Section D. but since less training data was used in this experiment, we reduced the number of hidden layers in the decoder network to 4. The models were trained for 100 epochs with 220 datasets in each epoch and learning rate 10-4 The validation sets used in training included 2 datasets and the evaluation sets used to compare the models in completion and forecasting tasks included 2 datasets each.
G.2 Data
We generated simulated data using a simplified equation inspired by the cancer evolution equa- tion from Gatenby and Gawlinski [13]. The equation involves three quantities: the healthy cells (z, f) + u(z,t), t), the cancerous cells (z, + vzz t), and acidity measure (z, + w(z,t) t), with
29

---

E [0, and z E C R². Their temporal and spatial dynamics are described by:
= Tu u11 - k//ky) du w r 1/ l/) D, V. ú ww dw w Dw V
(S8)
where Tu: Te, and Tw are apparition rates, ka and kp control the maximum for the number of healthy and cancerous cells, d, and d, are death rates relative other species, and D, and D., are dispersion rates.
We simulate trajectories from Eq. S8 discretized space-time grid using the SpatialPy simulator [6]. For the purpose of our experiment, we selected realistic parameter ranges that would also produce quick evolution. The selected parameter ranges are therefore close [13], but for example the diffusion rates were increased. The parameters for the simulations are sampled according to the following distributions: kn 10, hig 100, Tu
0.0033), To (0.0227, 0.033). Tw ~ (1.88, 2.2), D ~ U(000009, 0.0011), Dw- (000009, 0.011). du = O, d, ~ 4(0.45, 0.55). number of cancerous cells, with fast decay in healthy cells due to acid spread, as depicted in Figure S3. The process starts with the healthy cells count u constant across the whole space, cancerous cells count v zero except for one cancerous cell at uniformly random position, and acid count w zero across the whole space.
These parameters lead to relatively
Sane
Cancer cels count
[-0
t-2
count
4.0
25
10
model at times t 0, 1,2. Each row represent a different quantity in the model (from top to bottom, healthy cells, cancerous cells, and acid). At time the population of healthy cells is 10 at each point of space, and single point occupied by additional cancerous cells. After this time, cancerous cells start spreading and slowly multiplicating, while acid spreads faster, killing healthy cells along its way.
Figure S3: Example simulation of the Reaction-Diffii
30

---

We generated a total 5 10 trajectories on a 11 X 11 grid at times t 0,.... 5. We used the first 4 - 10 trajectories to generate datasets for model training and the other 10 generate datasets for evaluation. The datasets were generated by sampling from the trajectories as follows. First, we extracted the population counts at three consecutive time points from a trajectory and time point selected at random. We then sampled 50-200 context points and 100-200 target points from this subsample to create a dataset representing either a completion or forecasting task. For the completion task, the context and target points were sampled from the subsample at random, and for the forecasting task, all context points were sampled from the first two time points, and all target points from the third time point were included in the subsample. All datasets included in the training and validation data represented the completion task.
Finally, the input and output variables at each context and target point were determined as follows. In the experiments reported in the main text, we considered the case of partial observations, where only the difference thus intractable. Predicting future states of the quantity u v without running inference and without having access to is challenging problem, similar to several application cases of interest. In this appendix, we additionally consider a scenario where we include the acid count in the input variables and take the healthy cells count as the output variable.
v is observed. This makes the simulator model a hidden Markov model and
G.3 Additional results
A promising scientific application of neural processes is inferring latent variables given other measured quantities. To explore this capability, we attempt to infer the quantity of healthy cells (u(z, t)) given as input the acidity concentration (w(z, f)) for some z, t. For this study, we augmented the inputs to (z, f, wzz t)). Notably, the regression problem is not equivariant in the last input dimension, so we encoded w 1) in a non-relational manner also in the RCNP and RGNP models. The results are presented in Table S10. The difference between the completion and forecasting tasks is small in this experiment, and the best results are observed with RGNP both tasks. AGNP performance is surprisingly poor and will require further investigation.
Table S10: Normalized log- aigliike scores in the additional experiments of the Reaction- problem for both tasks (higher is better). Mean and (standard deviation) from 10 training runs evaluated on a separate test dataset.
RCNP
RGNP
ACNP
AGNP
CNP
GNP
Completion
0.35 (0.12)
0.74 (0.18) 0.24 (0.25)
-2.20 (3.67)
0.44 (0.17)
0.32 (0.26)
Forecasting 0.29 (0.11)
0.63 (0.15) 0.38 (0.21)
-1.7 (3.0)
0.46 (0.13) 0.41 (0.20)
G.4 Computation time
On GPU, for all methods the training runs of this experiment were particularly fast, each one lasting no more than minutes. We carried out a total of 10 training runs for each method, with total of 120 training runs. We also performed around 60 exploratory training runs (with longer training time, larger networks, etc.). The total training time comes less than 15 hours.
H Additional experiments
In addition to the experimental evaluation in the main paper and their extended discussion in the sections above, explored several further experimental setups to demonstrate the wide applicability of our proposal. These are (i) an application of the autoregressive structure proposed by [3] to our proposed approach (Section H.)), (ūi) a demonstration of how incorporate rotational equivariance into the model (Section H.2), and (ili) an exploration of the method's performance on two popular image data regression tasks (Section H.3).
H.1 Autoregressive CNPs
To further demonstrate the effectiveness of our RCNP architectures, we ran set of experiments adopting autoregressive approach to both our RCNP and other CNP models. Autoregressive
31

---

CNPs (AR-CNPs) are a new family of CNP models which were recently proposed by [3]. Unlike traditional CNPs, which generate predictions independently for each target point, or GNPs, which produce a joint multivariate normal distribution over the target set, AR-CNPs define a joint predictive distribution autoregressively, leveraging the chain rule of probability. Without changing the CNP architectures and the training procedure, AR-CNPs feed earlier output predictions back into the model autoregressively to predict new points. This adjustment enables the AR-CNPs to capture highly dependent, non-Gaussian predictive distributions, resulting in better performance compared to traditional CNPs.
In this experiment, we compare AR-RCNE AR-CNP, and AR-ConvCNP on synthetic regression data. The AR approach is applied at the evaluation stage to generate the target predictions autoregressively. When applying AR, the training method for the models remains unchanged. Thus, for this experiment, we use the same models that were trained with synthetic data. Details about the data generation and model configurations can be found in Section D. We chose the translational-equiarriio version of RCNP for each task.
The full results are reported in Table S11. As we maintain consistent experimental settings with previous synthetic experiments (Section D), we can directly compare the performance of models without the AR approach (Table S3-S7) to those enhanced with the AR approach. We observe that the AR approach improves RCNP performance the Gaussian experiments in all dimensions and in the non -a--aussiian sawtooth and mixture experiments when d₂ > 1. When de = 1. AR-RCNP results are comparable to RCNP results in the mixture task and worse in the sawtooth task. We also observe that AR-CNP results are worse than CNP results in most sawtooth and mixture experiments. Overall we can see that AR-RCNPs are better than AR-CNPs across all dimensions and tasks under both INT and OOID conditions, and achieve close performance compared with IR-ConvCNPs under dN = 1 in the Gaussian experiments.
Table S11: Comparison of the interpolation (INT) and out-of-iinpu-ddistiibutton (OOID) performance of AR-RCNP model with other AR baselines on various synthetic regression tasks with varying input dimensions. We show the mean and (standard deviation) obtained from 10 runs with different seeds. "F" denotes failed attempts that yielded very bad results. Our A RCNP method performs competitively in low dimension, and scales to higher input dimensions (d 2), where AR-ConvCNP is not applicable.
Malin-
Weakly-peridiic
Sawsooth
Misture
(is
(ivergmnss(
(ivaggeoss()
Ing-Iikclihnosl(11
lgg-ildihhoddt))
INT
O0ID
INT
O0ID
OOID
INT
COID
INT
OOID
RCNP
(6:00)
(0.00)
(0.00)
0:02 (0.00)
(0.00)
0.04 (0:00)
1,93 (0.05)
1.97 (0.04)
0.20 (0.27)
(.20
AR ConyUNP
0.01 (0.00)
0.01 (0.00)
(0.00)
0.00 (0.00)
(0.00)
0.01 (0:00)
(0.09)
(0.00)
0.83 (0:00)
(0000)
AR-CNP
0.09 (0.01)
4.33 (3.44)
0.12 (0.01)
0.21 (0.01)
3,05 (4.52)
-177 (1.20)
.0.2 (0.19)
AR-RCPP
0.14
0.14 (0.01)
0,13
0.13 (0.00)
0.11 (0.00)
0.11 (0:00)
1.55 (0.01)
1.55 (0.0))
0.20 (0:00)
0.20 (0.00
AR-CNH
0.29 (0.00)
3,05 (1.26)
0.32 (0.02)
2.24 (1.20)
0.33 (0.00)
2,00 (0:75)
0.25 (0.33)
0.26 (0.09)
0.39 (0.01)
-372 (1 68)
AR RCNP
0,23 (0.00)
0.23 (0.00)
0.15 (0.00)
0.15 (0.00)
0.16 (0.00)
0.16 (0.00)
0.87 (0.01)
0.87 (0.01)
0.10 (0.00)
.0.00 (0.00)
AR-CNP
0.49 (0.00)
1.23 (0.05)
0.59 (0.00)
0.97 (0.05)
0.38 (0.00)
1.47 (0.35)
-0.00 (1 13)
-0.65 0..
-3.51 (2.92)
H.1.1
Computation time
The AR procedure is not used in model training, so the training procedure aligns with that of standard CNP models. The experiments reported in this section used AR with models trained for the synthetic regression experiments in Section 5. the models had been trained from scratch, we estimate that the total computational cost for these experiments would have been approximately 2000 GPU-hours. AR is applied in the evaluation phase, where we predict target points autoregressivelly This process is considerably lengthier than standard evaluation, with the duration influenced by both the target set size and data dimension. Nonetheless, since only a single pass is required during evaluation, the overall computational time remains comparable to standard CNP models.
H.2 Rotation equivariance
The experiments presented in this section use a two-dimensional toy example to investigate modeling datasets that are rotation equivariant but not translation equivariant. We created regression tasks based
32

---

on synthetic data sampled from GP model with an exponentiated quadratic covariance function and mean function m [x) = | where A =iag(a) a fixed diagonal matrix with unequal entries and R is random rotation matrix. We considered an isotropic-kernel model version with the standard EQ covariance function (Equation S3) and an anisotropic- covariance function k(x, x') exp(-||BRa - BRx'|] (1)), where B diag(b) is fixed diagonal matrix with unequal entries and R is the same random rotation matrix as in the mean function. Both models use the anisotropic GP mean function defined previously.
model version with the
We generated the datasets representing regression tasks by sampling context and target points from the synthetic data generated with petropic-kernd or anisotropic-kernel model version as follows. The number of context points sampled varied uniformly between 1 and 70, while the number of target points was fixed at 70. All training and validation datasets were sampled from the range X' = [-4,0, 0) [-4,0] 0]. To evaluate the models in an interpolation (INT) task, we generated evaluation datasets by sampling context and target points from the training range. To evaluate the models in an out-of-input distribution (OOID) task, we generated evaluation datasets with context and target points sampled in the range [0, 4] x [0, 4].
We chose a comparison function g that is rotation invariant but not translation invariant defined as²
g(x, x') ([xx x x'|22, ||x||2, 12)
Since 'simple' RCNPs are only context preserving for translation translation-oquivariinnt equivariance and not for other equivariances (Proposition 4.9), we combine the ootation-equivarinn comparison function with a FuIIRCNP model.
We compared FulIRCNP with the rotation-equivariant comparison function, FuiRRCNP (rot), with CNP and ConyCNP models in both the isotropic-kemee and anisotropic-kernel test condition. We also used isotropic-kenee
data to train and evaluate a translation-eqiivariant ("stationary') ') RCNP version, RCNP (sta), that uses the difference comparison function. The model architectures and training details were the same as in the two-dimensional synthetic regression experiments (Section D.I).
The results reported in Table S12 show that, as expected, the isotropic-kernel test condition is easier than the amisotropic-kent Kerpel test condition where the random rotations change the covariance structure. The best results in both conditions interpolation task and with FuIIRCNP (rot) when the models are evaluated in the OOID task. Since the input range in the OOID task is rotated compared the training range, FulIRCNP (rot) results on the same level in both tasks while the other models are not able to make reasonable predictions in the OOID task. We also observe that that the FulIRCNP (rot) results in the interpolation task are better than CNP or RCNP (sta) results in the isotropic-kernel test condition and similar to CNP in the anisotropic-kerne test condition.
observed with ConvCNP when the models are evaluated in the
Table S12: Normalized log- likelihhoood scores in the experiments using synthetic data generated with isotropic and anisotropic GP kernels and an anisotropic GP mean function. The mean and (standard deviation) reported for each model are calculated based 10 training outcomes evaluated with separate datasets representing interpolation (INT) and intt-ff-inpuudddsiibuttoo (OOID) tasks. "F" denotes failed attempts with likeliill scores below 60.
Isotropic GP kernel
Anisotropic GP kemel
INT
OOID
INT
OOID
FuIIRCNP (rot)
-0.89 (0.04)
-0.90 (0.04)
-6.37 (0.34)
-6.38 (0.31)
RCNP (sta)
-1.16 (0.03)
F
ConvCNP
-0.60 (0.02)
F
-2.33 (0.12)
F
CNP
-1.42 (0.03)
-6.43 (1.13)
H.2.1 Computation time
The models were trained on a single GPU, and each training run took around 6 hours with FulRRCNP 2 hours with CNP, 4 hours with ConyCNP and 3 hours with RCNP The results provided in this
"This comparison function is invariant to proper and improper rotations around the origin (i.e., including mirroring).
33

---

appendix required 70 training runs and consumed around 250 GPU hours. Additional experiments related to this example consumed around 2500 GPU hours.
H.3 Image datasets
The experiments reported in the main paper demonstrate how RCNP can leverage prior information about equivariances in regression tasks including tasks with more than 1-2 input dimensions. In this section, we provide details and results from additional experiments that used MNIST [24] and CelebA [27] image data to create regression tasks. The neural process models were used in these experiments to predict the mean and variance over pixel values across the two-dimensional image given some pixels as observed context. The experiments investigate how RCNP compares other conditional neural process models when the assumption about task equivariance is incorrect.
H.3.1 Setup
The experiments reported in this section compared RCNP, CNP, and ConvCNP models using in- terpolation tasks generated based on 16 x 16 pixel images. The pixels were viewed as data points with d₂ 2 input variables that indicate the pixel location on image canvas and dy output variables that encode the pixel value. The output dimension depended on the image type. Pixels in grayscale images were viewed as data points with one output variable (dy = that takes values between 0 and I while pixels in RGB images were viewed as data points with three output variables (dy = 3) that each take values between and encode one RGB channel. We evaluated the selected models using standard MNIST and CelebA images (Section H.3.2) and examined modeling translation cquivariance with toroidal MNIST images (Section H.3.3).
The model architectures used in the image experiments match the architectures used in the synthetic regression experiments (Section D.I) extended with an output transformation to bound the predicted mean values between 0 and 1. While we used images with multiple output channels, there was no need use multi-output architecture such as discussed in Section F.I. This is because we assumed that all output channels are observed when a pixel is included in the context set and unobserved otherwise, meaning that the output channels in image data are not associated with separate context sets. RCNP models were used with the standard stationary comparison function unless otherwise mentioned.
The models were trained for 200 epochs with 2 datasets in each epoch and learning rate 3- 10 1-04 and tested with validation and evaluation sets that included 2 datasets each. The log- aigliikllhoodd score that used as training objective is usually calculated based on the predicted mean and variance, but when the models were trained on MNIST data, started training with the predicted variance fixed to small value. This was done to prevent the neural process models from learning to predict the constant background values with low variance and ignoring the image details. We trained with the fixed variance between epochs 1-150 and with the predicted variance as usual between epochs 151-200.
The datasets used train and evaluate the neural process models dataset was generated based on an image that was sampled from the selected image set and downscaled to 16 x 16 pixels. The downscaled images were either used as such transformed with random translations as explained in Section H.3.3. The context sets generated based on the image data included 2- 128 pixels sampled at random while the target sets included all 256 pixels. The evaluation sets used each experiment were generated with the context set size fixed to 2, 20, 40, and 100 pixels.
generated as follows. Each
H.3.2 Experiment Centered images (no translation)
The experiments reported in this section used MNIST and CelebA images scaled to the size 16 x 16 (Figure S4). The normalized log-likelihood scores calculated based on each evaluation set and across training runs for each neural process model are reported in Table $13. The best results are observed with ConvCNP, while the order between RCNP and CNP depends the context set size. CNP results are better than RCNP results when predictions are generated based on N context points while RCNP results are better when more context data is available.
We believe CNP works better than RCNP when N 2 because the absolute pixel locations are generally informative about the possible pixel values in the MNIST and CelebA images. This means
34

---

6 3
9
7
(a)
(b)
Figure S4: Example (a) MNIST and (b) CelebA 16 x 16 images.
Table S13: Normalized log-like oggliiklihood scores in the image experiments using (a) 16 x 16 MNIST images or (b) 16 16 CelebA images (higher is better). The mean and (standard deviation) reported for each model are calculated based 10 training outcomes evaluated in interpolation task with context sizes 2, 10, 20, and 100. RCNP results are the lowest when the context set is small but improve when the context size increases.
2
20
40
100
RCNP
5.52 (0.30)
7.13 (0.10)
8.08 (0.13)
9.35 (0.15)
(a)
ConvCNP
6.43 (0.02)
7.35 (0.04)
8.27 (0.04)
9.87 (0.06)
CNP
6.29 (0.12)
7.04 (0.03)
7.47 (0.03)
8.00 (0.05)
RCNP
0.18 (0.00)
1.19 (0.01)
1.79 (0.01)
3.01 (0.03)
(b) ConvCNP
0.33 (0.00)
1.36 (0.01)
2.03 (0.01)
3.37 (0.03)
CNP 0.31 (0.00)
0.90 (0.01)
1.09 (0.01)
1.25 (0.01)
that the regression tasks generated based on these images are not translation equivariant and the assumption encoded in the RCNP comparison function is incorrect. Since RCNP does not encode the absolute target location, it needs context data to make predictions that CNP and ConvCNP can make based on the target location alone (Figure S5).
While RCNP needs context data to work around the incorrect assumption about translation equivari- ance, RCNP results are better than CNP results when more context data is available. We believe this because the context set representation CNP does not depend on the current target location while RCNP can learn to preserve context information that is relevant to the current target. For example, while CNP looses information about the exact observed values when the context set is encoded, RCNP generally learns reproduce the observed context. RCNP may also learn to emphasize context points that are close to the current target and may learn local features that are translation equivariant.
H.3.3 Experiment 2: Translated images
The interpolation tasks in the previous image experiments do not exhibit translation equivariance since the input range fixed and the image content centered on the canvas. To run an image experiment with translation equivariance, we converted the 16 x 16 MNIST images into toruses and applied random translations to move image content around. The translations do not change the relative differences between image pixels on the torus, but when the torus is viewed an image, both the absolute pixel locations and relative differences between pixel locations are transformed. This means that to capture the translation equivariance, needed to calculate differences between pixel locations on the torus.
The experiments reported in this section compare models trained using either centered image data (Figure $4 (a)) or image data with random translations on the torus (Figure S6). We ran experiments with both RCNP (sta) that calculates the difference between pixel locations on the image canvas and RCNP (tor) that calculates the difference between pixel locations on the torus. The models trained with centered image data were evaluated with both centered and translated data while the models
35

---

N-2
N-20
N-40
N-100
Context
3
a
RCNP
9
9
ConvCNP
3
7
CNP
N=100
Context
RCNP
ConVCNP
CNP
Figure S5: Example context sets and predicted mean values in the image experiment using (a) MNIST and (b) Celeba 16 16 data.
36

---

377
4
3
Figure S6: MNIST 16 x 16 example images with random translations on a torus.
trained with translated data were evaluated with translated data. The normalized log- logllikeihhod scores are reported in Table S14.
Table S14: Normalized log- likelihood scores in the image experiments using 16 x 16 MNIST images with random translations on torus (higher is better). We compare (a) models trained on centered image data evaluated on centered image data, (b) models trained on centered image data evaluated on image data with random translations, and (c) models trained on image data with random translations evaluated on image data with random translations. The mean and (standard deviation) reported for each model are calculated based on 10 training outcomes evaluated in interpolation task with context sizes 2. 10, 20, and 100. "F" denotes very bad results with negative log- logllikelhhood scores. RCNP (tor) trained on centered image data is able to make predictions about images with random translations because the translation mechanism is modeled in the comparison function.
2
20
100
RCNP (sta)
5.52 (0.30)
7.13 (0.10)
8.08 (0.13)
9.35 (0.15)
(a)
RCNP (tor)
5.51 (0.01)
6.63 (0.04)
7.59 (0.03)
9.15 (0.09)
ConvCNP
6.43 (0.02)
7.35 (0.04)
8.27 (0.04)
9.87 (0.06)
CNP
6.29 (0.12)
7.04 (0.03)
7.47 (0.03)
8.00 (0.05)
RCNP (sta)
F
F
F
F
RCNP (tor)
5.62 (0.01)
6.68 (0.04)
7.61 (0.06)
9.13 (0.09)
(b)
ConvCNP
F
F
CNP
F
F
F
RCNP (sta)
5.51 (0.21)
6.66 (0.11)
7.67 (0.09)
9.35 (0.11)
RCNP (tor)
5.56 (0.10)
6.72 (0.06)
7.70 (0.09)
9.31 (0.14)
(c)
ConvCNP
5.61 (0.01)
6.72 (0.03)
7.78 (0.04)
9.60 (0.08)
CNP
5.55 (0.03)
6.29 (0.03)
6.80 (0.02)
7.52 (0.09)
The results Table S14 (a) extend the results discussed in the previous section (Table S13 (a)). The results indicate that when the models are trained and tested with interpolation tasks generated based on centered image data, RCNP (sta) works better than RCNP (tor). We believe this is because the relative locations calculated based on pixel coordinates on the image canvas are more informative about the absolute location than relative locations calculated on the torus. In other words, while RCNP (sta) and RCNP (tor) do not encode the absolute target location on the image canvas, RCNP (sta) can learn to derive equivalent information based on regularities between context sets.
The image interpolation task becomes translation equivariant on the torus when the models are trained or tested using the image data with random translations. The results reported in Table $14 (b) indicate that using differences calculated on the torus allows RCNP (tor) to generalize between centered and translated images while the other models are not able to make sensible predictions in this OOID task. All models work well when trained with matched data (Table S (c)), meaning that training on image data with random translations allows the other models to learn the equivariance which encoded in RCNP (tor). The best results in Table S14 (c) are observed with ConvCNP, while RCNP (tor) and RCNP (sta) are similar CNP when = 2 and better than CNP when more context data available.
37

---

H.3,4 Computation time
The models were trained on single GPU, and each training run with RCNP took around 3-4 hours, each training run with CNP around 1-2 hours, and each training run with ConvCNP around 4.5-5.5 hours depending on the training data. The results presented in this supplement required 1 11 training runs with total runtime around 350 hours. In addition other experiments carried out with image data took around 2000 hours.
38


Reviews:
[{"review_id": "BvdjSyovgc", "reviewer": "Reviewer_C6E4", "summary": "The authors propose a class of neural processes that can be constructed to enforce invariance to particular properties like translation and rotation.", "questions": "See weaknesses", "limitations": "Limitations are adequately addressed", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- As far as I know, the proposed architecture and technique for enforcing invariances in Section 3 is novel.\n- Improvements are shown over standard neural processes on a range of tasks\n- The authors show both theoretically and empirically that the proposed architecture does indeed enforce the invariances described.", "weaknesses": "- Choosing $g$ in Eq 5 to enforce a particular invariance seems difficult. The authors provide choices that enforce isotropy and translation invariance, but it would not be obvious how to enforce a different type of invariance.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review and comments, and we are glad to see that you found our paper interesting.\n\nRegarding your concern about enforcing particular invariances, we would like to comment that it is indeed true that our method only applies to equivariances that can be enforced via a comparison function $g$. This is explicitly mentioned within the paper, such as the Introduction:\n\n> Our proposed method works for equivariances that can be expressed relationally via comparison between pairs of points (e.g., their difference or distance); [...]\n\nand the Discussion (just before the *Limitations* paragraph):\n\n> Our method applies to equivariances that can be induced via an appropriate comparison function; here we focused on translational equivariances (induced by the difference comparison) and equivariances to rigid transformations (induced by the distance comparison).\n\nTo further address your point and provide an additional example, we ran a new set of experiments with a FullRCNP which implements equivariance to proper and improper rotations (i.e., rotations + mirroring). As per Proposition 4.5 of our paper, we need to select an adequate comparison function, which we define as follows:\n$$g_\\text{rot}(\\mathbf{x},\\mathbf{x}') = (\\Vert \\mathbf{x} - \\mathbf{x}' \\Vert_2, \\Vert \\\\mathbf{x} \\Vert_2, \\Vert \\mathbf{x}' \\Vert_2)$$\nThis comparison function is based on the distance between points as well as their distance from the origin (the center of rotation). This representation makes the comparison function invariant to rotations and mirroring, but not translations. Thus, it will induce a FullRCNP which is equivariant to these transformations.\n\nAs mentioned in the global rebuttal, we tested this new RCNP variant in a new set of synthetic experiments, namely a regression task involving a Gaussian Process whose mean and covariance functions introduce rotational and mirror symmetry, but not translational symmetry. Such a task is representative for example of a physical model set in a potential well (e.g., a point charge or mass in the origin). Due to the limited time available for the rebuttal, we present preliminary results in Table 2 of the rebuttal PDF. Full results will be included in the revised paper and Supplementary Material. Briefly, our results show that the FullRCNP is able to leverage the equivariances intrinsic to the task to outperform a standard CNP in both low and high dimensions.\n\nIn conclusion, our relational neural process framework is suitable for isotropy, translational equivariance, and the newly implemented (proper and improper) rotational equivariance. Whether other equivariances can be expressed via our relational approach is an interesting direction for future work. Nonetheless, we believe the equivariances addressed in the paper represent a large class of useful equivariances. This is demonstrated by the fact that several key papers and models proposed in the neural process literature (as well as the broader machine learning literature) focus *only* on translational equivariance (e.g., ConvCNP, Gordon et al. 2020; ConvGNP, Markou et al., 2022; FullConvGNP, Bruinsma et al., 2020), with the limitations we discussed in the paper and that our proposed method overcomes.\n\nFor clarity, we will include an additional explanation based on the paragraph above to the *Limitations* section of the paper to make the contribution and limitations of the paper clearer.\n\n### References\n\n- Bruinsma et al. (2020). The Gaussian Neural Process. *AABI*.\n- Gordon et al. (2020). Convolutional Conditional Neural Processes. *ICLR*.\n- Markou et al. (2022). Practical Conditional Neural Processes via Tractable Dependent Predictions. *ICLR*.\n"}, {"review_id": "y4vd5h2ecB", "reviewer": "Reviewer_N64F", "summary": "The paper presents a novel approach for incorporating equivariance into conditional neural processes (CNPs) which can scale to high dimensions. Modelling equivariance is essential to improve the performance of CNPs. Unlike previous approaches that use convolution and become impractical with increased input dimensions, this work uses relational information and discards absolute information. As a result, this simple method can handle high-dimensional inputs. The authors also prove that their approach is context preserving which means they do not lose other information. Their empirical results demonstrate that the proposed method is comparable to convolutional CNPs, GNPs etc. on a diversified range of tasks.\n", "questions": "**Questions**\n\n1. I wonder if there is a middle ground between the full RCNP and the diagonal RCNP. In the matrix terminology, it is a bit similar to a low-rank approximation of the full comparison set. I think this would allow people to balance computational cost and expressivity. Also by only comparing only to a subset of K \"important\" context points, we can lower the cost to O(NK) rather than the quadratic cost of the full RCNP.\n2. By only using pairwise comparison in the model construction, I wonder whether we would lose high-order interaction which involves more than two context points.", "limitations": "The authors have discussed the limitations of their work and I have no concerns of any potential negative societal impact.", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "**Strengths**\n\n1. This paper addresses the challenge of incorporating equivariance into Conditional Neural Processes (CNPs) for high-dimensional problems. The proposed models are shown to be translation-equivariant, allowing them to scale to higher dimensions and are comparable/outperform existing CNP and GNP models on a wide range of tasks.\n2. The paper is well-structured and clearly presented, with a strong motivation for the research, precise technical statements, and comprehensive background. \n3. The paper provides robust theoretical results, all of which are supported by proof. The empirical investigation is extensive, covering a diversified range of tasks. The authors demonstrate the effectiveness of their models through experiments on synthetic Gaussian and non-Gaussian regression tasks, Bayesian optimization, Lotka-Volterra models and reaction-diffusion models. ", "weaknesses": "**Weakness**\n\nGiven the extensive use of RGNP in the experiments and its significant role in the paper's findings, it would still be beneficial for the authors to provide a concise description or a mathematical formulation of the RGNP in the main text. Even though it is not difficult to extend the mathematical description of RCNP to RGNP.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your very positive and insightful comments, we are glad to see that you find our paper particularly well-suited for NeurIPS. In the following we address your questions and points raised.\n\n### Weaknesses\n\n> *Given the extensive use of RGNP in the experiments and its significant role in the paper's findings, it would still be beneficial for the authors to provide a concise description or a mathematical formulation of the RGNP in the main text. Even though it is not difficult to extend the mathematical description of RCNP to RGNP.*\n\nThanks for pointing this out. In the revised paper we will include a concise mathematical description of the RGNP model in the main text.\n\n### Questions\n\n> *1. I wonder if there is a middle ground between the full RCNP and the diagonal RCNP [...] similar to a low-rank approximation of the full comparison set. [...] Also by only comparing only to a subset of K \"important\" context points [...].* \n\nThis is a very interesting and insightful suggestion to lower the asymptotic cost of full RCNPs, which would be related to other work in the Gaussian process literature (“sparse” GPs and inducing points), and also to the concept of coresets. Indeed, this is a potential direction of future work aimed at reducing the computational complexity. \n\nA naive baseline to consider in this direction would be to only encode a target point based on the $K$ closest context points, thus reducing the cost of FullRCNP to $O(K^2 M + N M)$, where the first term is the relational encoding and the second term is the naive cost of computing the distances between $M$ targets and $N$ context points (although the latter can be sped up via smarter nearest-neighbor search algorithms; see e.g. Hyvönen et al. 2022). While $K$ is nominally a constant, to guarantee a reasonable performance in practice we envision it could be chosen based on the (maximum) number of context points in the task. Still, by choosing a square-root scaling ($K \\propto \\sqrt{N}$) we can recover a manageable asymptotic cost of $O(N M)$.\n\n> *2. By only using pairwise comparison in the model construction, I wonder whether we would lose high-order interaction which involves more than two context points.* \n\nThis is an excellent point and the perhaps surprising theoretical answer is no, in our construction. Our context-preservation theorems demonstrate that no information is lost about the entire context set, so in principle a RCNP with a sufficiently large network is able to reconstruct any high-order interaction of the context set (see Section 4.2 in the main text and the full proof in Appendix B.2).\n\nHowever, in practice the answer can be somewhere in the middle, in that the chosen representation is built on two-point interactions, so depending on the network size it may be harder for the network to effectively encode the simultaneous interaction of many context points.\n\nWe will add these informative remarks in the theoretical section of the Supplement.\n\n### References\n\n- Hyvönen et al. (2022). A Multilabel Classification Framework for Approximate Nearest Neighbor Search. *NeurIPS*.\n"}, {"review_id": "wnsZB00lUY", "reviewer": "Reviewer_jSJf", "summary": "This work introduces a new member of the neural process model family that is designed for biasing the model towards representing equivariances in the data. It does this by including relational information among the context set, and between the predicted and context inputs in the encoder for a new input. Also, only the relational information is used for encoding and the absolute information is discarded. Experiments are given for a several types of equivariances/applications to show the method successfully models equivariances in the data.", "questions": "Rather than inputting the full unweighted relational matrix to the embedding, would it be possible to learn which pairs of context set input/outputs and the value of their comparison function to attend to? In addition to weighting pairs of data points, some variant of attention could help break the quadratic scaling.\n\nWhat is the cause of the jagged mean line for RCNP in Figure 1 (b), (e) relative to CNP and GPs?", "limitations": "An honest assessment of the limitations is given in Section 7, including the point that the full RCNP model has quadratic scaling.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Simple but effective technical idea with convincing experiments. For instance, I found Figure 1 very illuminating. The work shows promise for scaling GP-like models (or rather, emulating GP kernels) for distributions over functions to large data/dimensions.", "weaknesses": "The full RCNP variant has tractability issues, although it is shown that the simpler RCNP using the diagonal elements of the relational matrix performs satisfactorily.\n\nEquivariances other than translations and rigid transformations, and other variants of CNP are left to future work. Would it not be too much effort to examine a few of them in this paper?\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your very positive comments and useful remarks about our work. We address your remarks and questions below.\n\n### Weaknesses\n\n> *The full RCNP variant has tractability issues, although it is shown that the simpler RCNP using the diagonal elements of the relational matrix performs satisfactorily.*\n\nThis is a fair point. As counterpoints to consider, we would like to remark that CNPs are commonly applied in the low-data setting (e.g., up to several hundred data points). Thus, while the quadratic scaling of full RCNPs is admittedly improvable, it is not as limiting as it might seem from its asymptotic analysis. \nMoreover, as you mentioned, for the very common application of translational invariance, “diagonal” RCNPs are enough. Importantly, for the case of translational equivariance, the diagonal RCNP is not just an approximation of the full RCNP, but an exact solution. This is proved in our theorems (see Proposition 4.9 in the main text and its proof in Appendix B.2), and confirmed by our empirical results.\n\n> *Equivariances other than translations and rigid transformations, and other variants of CNP are left to future work. Would it not be too much effort to examine a few of them in this paper?*\n\nGood point. In terms of examining other variants of (R)CNP, we implemented and ran new experiments with the autoregressive RCNP (AR-RCNP), where the AR-CNP is a recent model introduced in Bruinsma et al. (2023). The AR-RCNP demonstrates the application of our technique to another model of the CNP family. Results are presented in Table 1 of the rebuttal PDF, and will be included in the revised paper.\n\nRegarding other equivariances: we now also implemented equivariance to proper and improper rotations. Please refer to our answer to reviewer C6E4 for a detailed description of the newly implemented model, experiments and related discussion. We report a summary below.\n\nThe newly implemented equivariance is based on the comparison function:\n$$g_\\text{rot}(\\mathbf{x},\\mathbf{x}') = (\\Vert \\mathbf{x} - \\mathbf{x}' \\Vert_2, \\Vert \\\\mathbf{x} \\Vert_2, \\Vert \\mathbf{x}' \\Vert_2),$$\nwhich is based on the distance between points as well as their distance from the origin (the center of rotation). This makes the comparison function invariant to rotations and mirroring, but not translations. Thus, it will induce a FullRCNP which is equivariant to these transformations, according to our Proposition 4.5.\n\nWe tested this new RCNP variant in a new set of synthetic experiments which incorporate rotational and mirror symmetry. Preliminary results are presented in Table 2 of the rebuttal PDF, and full results will be included in the revised paper and Supplementary Material. Our results show that the FullRCNP is able to leverage the equivariances intrinsic to the task to outperform a standard CNP in both low and high dimensions.\n\nIn conclusion, our framework is suitable for isotropy, translational equivariance, and the newly implemented (proper and improper) rotational equivariance. Whether other equivariances can be expressed via our relational approach is an interesting direction for future work. Nonetheless, we believe the equivariances addressed in the paper represent a large class of useful equivariances. This is demonstrated by the fact that several key papers and models proposed in the neural process literature (as well as the broader machine learning literature) focus *only* on translational equivariance (e.g., ConvCNP, Gordon et al. 2020; ConvGNP, Markou et al., 2022; FullConvGNP, Bruinsma et al., 2020), with the limitations we discussed in the paper and that our proposed method overcomes.\n\n\n### Questions\n\n> *Rather than inputting the full unweighted relational matrix to the embedding, would it be possible to learn which pairs of context set input/outputs and the value of their comparison function to attend to? In addition to weighting pairs of data points, some variant of attention could help break the quadratic scaling.*\n\nThis is an interesting point and related to work on Transformer neural processes (Nguyen & Grover, 2022). However, we are unsure how this would address the scaling issue per se, since the attention mechanism is also notoriously quadratic in the size of the context set. \n\nA somewhat similar idea, along these lines of finding a middle ground between the full RCNP and the diagonal RCNP, would consist of something like a low-rank approximation of the full comparison set, or by only comparing to a chosen subset of K \"important\" context points. This would relate to work in the Gaussian process literature (“sparse” GPs and inducing points), and also to the concept of coresets. Indeed, this is a potential direction of future work that we have been considering. Please see our response to reviewer N64F for additional remarks on this point.\n\n> *What is the cause of the jagged mean line for RCNP in Figure 1 (b), (e) relative to CNP and GPs?*\n\n- The jagged mean is a byproduct of the common ReLU activation function used in the (R)CNP architecture. Due to the ReLUs, the output of the network will naturally be a piecewise linear function (i.e., with discontinuous derivatives, which makes it look jagged).\n- Standard CNPs here underfit the data, so their mean *appears* smoother. However, note that it is still jagged after zooming in – the discontinuous derivative is just less prominent due to the underfitting. The outputs of CNPs commonly look jagged (e.g., see Figure 2 of Garnelo et al. 2018).\n- GPs are kernel methods, here with a Matérn 5/2 kernel which is twice differentiable, so the posterior mean function will also be twice differentiable (i.e., it will look reasonably smooth).\n\n### References\n\n- Bruinsma et al. (2023). Autoregressive Conditional Neural Processes. *ICLR*.\n- Garnelo et al. (2018). Conditional Neural Processes. *ICML*.\n- Nguyen & Grover (2022). Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. *ICML*.\n"}, {"review_id": "lYe01KaI9Q", "reviewer": "Reviewer_bQ2F", "summary": "This paper is an extension of the conditional neural process. The primary contribution is to augment the conditional neural process with the relational structure. The developed model was examined in synthetic regression, Bayesian optimization, Lotka-Volterra simulation, and others.", "questions": "No.", "limitations": "Limitations are highlighted in Section 7. Here it seems that the runtime complexity of predictive distributions is higher than that of Gaussian processes, which is quadratic w.r.t. the number of data points. Remember that one of the motivations in CNPs or NP-like models is the reduction of runtime complexity.\nOther limitations are included in the weakness section.\nThe definition of equivariance is incorrect.\nSome missing citations [1-5].", "rating": 2, "confidence": 4, "soundness": 1, "presentation": 2, "contribution": 1, "strengths": "1. This paper is complete in organization and easy to understand.\n2. The relational inductive biases mentioned in this work are crucial in some scenarios but were already studied in previous NP related work (See weakness).", "weaknesses": "**1. Unclear Motivation.** The motivation for using relational inductive biases is not well clarified in the Introduction section, e.g., in which scenarios or datasets the relational information is required. Meanwhile, I am not sure of the definition of relational in this work, and it seems the concept is similar to that in Set Transformer .\n\n**2. Lack of Novelty.** The novelty of this work seems relatively limited. There have been extensive works to incorporate equivariance into neural process models, such as work in [1-4]. As for the use of relational inductive bias, especially when the input/output are high dimensional, this has appeared in work [5].\n\n**3. Incomplete Analysis.** In Table (1)/(2), the kl divergences are reported, but I did not find what it means in the results. Notably, most empirical analysis claims the proposed model outperforms others, but the reasons behind these observations, e.g., what kind of equivariance matters, are unclear.\n\n**4. Missing benchmarks and baselines.** Images naturally hold translation equivariance/invariance properties and are more appropriate for evaluation. This work fails to compare with other equivariant (C)NPs [1-4].\n\n \n**5.** The contribution and the organization of the paper are mixed in Line 49-Line 56.\n\n**References:**\n\n[1] Kawano, M., Kumagai, W., Sannai, A., Iwasawa, Y., & Matsuo, Y. (2021). Group equivariant conditional neural processes. arXiv preprint arXiv:2102.08759.\n\n[2] Holderrieth, P., Hutchinson, M. J., & Teh, Y. W. (2021, July). Equivariant learning of stochastic fields: Gaussian processes and steerable conditional neural processes. In International Conference on Machine Learning (pp. 4297-4307). PMLR.\n\n[3] Markou, S., Requeima, J., Bruinsma, W., Vaughan, A., & Turner, R. E. (2021, October). Practical Conditional Neural Process Via Tractable Dependent Predictions. In International Conference on Learning Representations.\n\n[4] Foong, A., Bruinsma, W., Gordon, J., Dubois, Y., Requeima, J., & Turner, R. (2020). Meta-learning stationary stochastic process prediction with convolutional neural processes. Advances in Neural Information Processing Systems, 33, 8284-8295.\n\n[5] Wang, Q., & Van Hoof, H. (2022, June). Model-based meta reinforcement learning using graph structured surrogate models and amortized policy search. In International Conference on Machine Learning (pp. 23055-23077). PMLR.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and effort spent reading our paper. We would like to clarify what we perceive as misunderstandings of the nature and contribution of our paper. Our core contribution is in providing a simple, effective way to build exact equivariances (with a focus on translational equivariance) directly in the CNP architecture, *in a way that easily scales to higher input dimensions*. This is something that is currently not addressed in existing work, including the references you provided. Here we would like to clarify the perceived weaknesses of our paper, views which we do not share and explain further below.\n\n### Weaknesses\n\n> *1. Unclear Motivation. The motivation for using relational inductive biases is not well clarified in the Introduction section, e.g., in which scenarios or datasets the relational information is required.*\n\nWe clarified our motivation in the Introduction: proposing a simple method to implement some exact equivariances into CNPs (particularly, translational equivariance), in a way that scales efficiently with input dimension. The relational approach is how we achieve our goal.\n\n> *Meanwhile, I am not sure of the definition of relational in this work, and it seems the concept is similar to that in Set Transformer.*\n\nWe use Set Transformers to achieve invariance to permutations of the context set (see line 128). Unlike Set Transformers, we encode data points via the comparison function $g$. We call our encoding (and method) “relational” because data points are encoded based on how they compare or relate to each other.\n\n> *2. Lack of Novelty. The novelty of this work seems relatively limited. There have been extensive works to incorporate equivariance into neural process models, such as work in [1-4]. As for the use of relational inductive bias, especially when the input/output are high dimensional, this has appeared in work [5].*\n\nWe firmly disagree with this claim, which misses our main contribution. All existing methods for incorporating exact equivariances in CNPs (ConvCNP and variants [3,4], SteerCNP [2], EquivCNP [1]) require at some point a regular discretization (lattice) of the input space and convolution operations. This limits in practice the number of input dimensions these methods are applied to (often up to 1 or 2 equivariant dimensions), as lattice-based methods scale exponentially with the input dimension.\n\nThank you for the references. Note we already had [3,4] in our paper. **We will include a more extensive discussion of [1-5] in the Related Works.** Briefly:\n\n- [2] introduces SteerCNP, which directly generalizes ConvCNP to other equivariances and suffers from its exact same scaling issues (e.g., Section 8 of [2], “*One practical limitation of this method is the necessity to discretise the continuous RKHS Embedding, which can be costly*”). [2] mentions the possibility of using alternative architectures to bypass this issue, but this is not developed in the paper.\n- [1] introduces EquivCNP, based on LieConv (Finzi et al., *ICML* 2020). While LieConv per se can operate on irregular point clouds, EquivCNPs require constructing a lattice over the input space (see Algorithm 1 and Section 4.3 of [1]), thus reinstating the scalability problem.\n- [5] aims to learn arbitrary relational inductive biases in a NP model using a relational graph structure on the context set. While interesting, this differs from our goal of directly building *exact equivariances* in the CNP architecture.\n\n> *3. Incomplete Analysis. In Table (1)/(2), the kl divergences are reported, but I did not find what it means in the results. Notably, most empirical analysis claims the proposed model outperforms others, but the reasons [...] are unclear.*\n\nOur RCNP models implement specific equivariances (translational or isotropy), as specified in each experiment. If our method works better, it is because encoding that equivariance is useful for the task; we will add an explanation.\n\n> *4. Missing benchmarks and baselines. Images naturally hold translation equivariance/invariance properties and are more appropriate for evaluation.*\n\nWe added new results on MNIST and CelebA (see Table 3 and Fig 1 in the rebuttal PDF). In our paper, we focused on higher-dimensional tasks to showcase the capabilities of our model.\n\n> *This work fails to compare with other equivariant (C)NPs [1-4].*\n\nThis is not correct. We compared against ConvCNP and ConvGNP [3], where applicable. Due to the mentioned scalability issues, we could not test these methods above input $d_x=2$. ConvNPs [4] have the same scalability issues as ConvCNP and are generally outperformed by ConvGNP (see [3]). The other equivariant CNPs [1,2] have the same scalability issues. Moreover, SteerCNP [2] is a strict generalization of ConvCNP, so there is no need to run separate tests in the translational equivariant case.\n\n> 5. The contribution and the organization of the paper are mixed.\n\nWe rewrote a separate paragraph, distinct from the paper organization, with the contributions described at the beginning of this response.\n\n### Limitations\n\n> *Here it seems that the runtime complexity of predictive distributions is higher than that of Gaussian processes [...].* \n\nThis is not correct. Unlike CNPs which are amortized, GPs still require training (fitting the kernel parameters), which is generally *cubic* in the number of data points. There are quadratic techniques but they require tuning for good results (Maddox et al., 2021). Regarding scalability of FullRCNP, see also our answer to reviewer N64F.\n\n> *The definition of equivariance is incorrect.* \n\nOur definition of equivariance is correct and equivalent to the standard definition for general mappings, $\\tau f(X) = f(\\tau X)$. We added a section in the Supplement showing that our definition of equivariance of a prediction map based on the invariance of its representation (Eq. 3) is equivalent to the standard definition, as found e.g. in [2] or in the ConvCNP paper (Gordon et al., 2020)."}]

Author rebuttal:
We thank the anonymous reviewers for their comments and suggestions for improving our paper. We are glad to see that the majority of reviewers found the paper interesting and of impact. We provide clarifications and detailed answers to perceived weaknesses and raised questions in our individual responses.

In particular, we added a number of requested experiments and comparisons to baselines in different settings to the paper. Tables and Figures presenting these results can be found in the attached pdf. Please find below a short overview of the implemented changes:

### New experiments and results

- As asked by reviewer jSJf, we implemented and ran new experiments with another variant of RCNP – the autoregressive RCNP (AR-RCNP), where the AR-CNP is a recent model introduced in Bruinsma et al. (2023). The AR-RCNP demonstrates the application of our technique to another model of the CNP family. Results are presented in Table 1 of the rebuttal PDF, and will be included in the revised paper.
- As suggested by reviewers C6E4 and jSJf, we provided the RCNP implementation of a new class of equivariances, namely proper and improper rotations. We tested this new RCNP variant in a new set of synthetic experiments, a regression task involving a Gaussian Process whose mean and covariance function introduce rotational and mirror symmetry, but not translation. Such a task is representative for example of a physical model with a potential well. Due to the limited time available for the rebuttal, we present preliminary results in Table 2 of the rebuttal PDF. Full results will be included in the revised paper and Supplementary Material. Please also see our response to reviewer C6E4 for further details.
- As requested by reviewer bQ2F, we ran our RCNP model on an image completion task using classic image datasets such as MNIST and CelebA, to compare our performance to other members of the CNP family. Due to the limited time available for the rebuttal, we could only complete preliminary experiments with downscaled images. Results are presented in Table 3 of the rebuttal PDF, with examples in Figure 1 of the PDF, and will be included in the revised Supplementary Material.

### Other changes

Besides the new results explained above, we implemented a number of changes to the main text and Supplementary Material to further improve the clarity of the paper. 

- We clarified our motivation and contributions (see our response to reviewer bQ2F for details). 
- We expanded the Related Works section with new references and explained more in detail how our work differs from the existing literature on equivariant (C)NPs (please see our answer to reviewer bQ2F). Importantly, we highlighted how all existing works require a lattice (regular grid) construction in the input space, often followed by convolutions, which strongly limits scalability of the methods to higher input dimensions; a limitation which is not shared by our distinct “relational” approach.
- We expanded our limitations section, to better highlight that our relational approach hinges on finding an appropriate comparison function $g$. While our paper covers several important equivariances of large practical interest (translational, isotropy, now rotations), whether other equivariances can be expressed via our relational approach is left for future work. For details, see the response to reviewer C6E4.
- In addressing a remark by reviewer bQ2F, we clarified our definition of equivariance of prediction maps, based on the invariance of their representation (Eq. 3 in the paper). Our definition, used for its convenience in our paper, is mathematically equivalent to the more common definition of equivariance for general mappings; $\tau f(X) = f(\tau X)$, but the comment made us realize it is not immediately apparent. We added a section in the Supplementary Material explicitly showing this equivalence, starting from standard definitions of equivariance and group action (Gordon et al., 2020; Holderrieth et al., 2021).

### References

- Bruinsma et al. (2023). Autoregressive Conditional Neural Processes. *ICLR*.
- Gordon et al. (2020). Convolutional Conditional Neural Processes. *ICLR*.
- Holderrieth et al. (2021) Equivariant Learning of Stochastic Fields: Gaussian Processes and Steerable Conditional Neural Processes. *ICML*.


Meta-review: This work proposes the relational conditional neural process, a new neural process that biases the model towards representing exact equivariances in the data. Unlike previous approaches that use convolution, this approach uses relational information instead of absolute information and so scales well to high-dimensional inputs.

Strengths
- Good presentation and writing.
- Simple and effective idea.
- Convincing experiments compared to existing CNP/GNP models and robust theoretical results.

Weaknesses
- It's unclear how to enforce other types of invariance with a comparison function g.
- The full RCNP has tractability issues though the simpler one still performs well.
- There's no concise description of RGNP in the main text.
Decision: Accept (poster)