Approximately Equivariant Neural Processes
Authors: Matthew Ashman
Cristiana Diaconu
Adrian Weller
Wessel P Bruinsma
Richard E. Turner
Venue: NeurIPS 2024 poster (2024)
Keywords: equivariance
neural processes
meta learning
deep learning
probabilistic methods
Primary area: probabilistic_methods
TL;DR: A general approach to constructing approximately equivariant architectures with applications to neural processes

Abstract: Equivariant deep learning architectures exploit symmetries in learning problems to improve the sample efficiency of neural-network-based models and their ability to generalise. However, when modelling real-world data, learning problems are often not *exactly* equivariant, but only approximately. For example, when estimating the global temperature field from weather station observations, local topographical features like mountains break translation equivariance. In these scenarios, it is desirable to construct architectures that can flexibly depart from exact equivariance in a data-driven way. Current approaches to achieving this cannot usually be applied out-of-the-box to any architecture and symmetry group. In this paper, we develop a general approach to achieving this using existing equivariant architectures. Our approach is agnostic to both the choice of symmetry group and model architecture, making it widely applicable. We consider the use of approximately equivariant architectures in neural processes (NPs), a popular family of meta-learning models. We demonstrate the effectiveness of our approach on a number of synthetic and real-world regression experiments, showing that approximately equivariant NP models can outperform both their non-equivariant and strictly equivariant counterparts.

Approximately Equivariant Neural Processes
Matthew Ashman" University of Cambridge mca390cam.c au.uk
Cristiana Diaconu" University Cambridge cdd43@cam. .c.uk
Adrian Weller University of Cambridge The Alan Turing Institute aw665@cam. ac uk
Wessel Bruinsma Microsoft Research Al for Science wessel. p. bruinsma@gmall. com
Richard E. Turner
University of Cambridge
Microsoft Research AI for Science
The Alan Turing Institute ret23@cam ac uk
Abstract
Equivariant deep learning architectures exploit symmetries in learning problems to improve the sample efficiency of neural-network-bassd models and their ability generalise. However, when modelling real-world data, learning problems are often not exactly equivariant, but only approximately. For example, when estimating the global temperature field from weather station observations, local topographical features like mountains break translation equivariance. these scenarios, is desirable to construct architectures that can flexibly depart from exact equivariance in a data-driven way. Current approaches to achieving this cannot usually be applied out-of the-box to any architecture and symmetry group. In this paper, we develop a general approach to achieving this using existing equivariant architectures. Our approach is agnostic to both the choice of symmetry group and model architecture, making it widely applicable. We consider the use of approximately equivariant architectures in neural processes (NPs), popular family of meta- meta-learning models. We demonstrate the effectiveness of our approach on a number of synthetic and real- world regression experiments, showing that approximately equivariant NP models can outperform both their non-equivariant and strictly equivariant counterparts.
Introduction
The development of equivariant deep learning architectures has spearheaded many advancements in machine learning, including CNNs [LeCun et al., 1989], group equivariant CNNs [Cohen and Welling, 2016a, Finzi et al., 2020], transformers [Vaswani et al., 2017], DeepSets [Zaheer et al., 2017], and GNNs [Scarselli et al., 2008]. When appropriate, equivariances provide useful inductive biases that can drastically improve sample complexity [Mei et al., 2021] and generalisation capabilities [Elesedy and Zaidi, 2021, Bulusu et al., 2021, Petrache and Trivedi, 2023] by exploiting symmetries present in the data. Yet, real-world data are seldom strictly equivariant. As an example, consider modelling daily precipitation across space and time. Such data may be close to translation equivariant and therefore translation cquivariance serves as a useful inductive howeve, it is clear that not strictly translation equivariant due geographical and seasonal variations. More generally, whilst there aspects of modelling problems which are universal and exhibit symmetries such as equivariance (e.g. atmospheric physics) there often unknown local factors (e.g. precise local topography) which break these symmetries and which the models do not have access to. It is thus desirable for the
"Equal contribution.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

model to be able to depart from strict equivariance when necessary; that is, to develop approximately
equivariant models.
A family of models that have benefited greatly from the use of equivariant deep learning architectures neural processes [NPs; Garnelo et al., 201 Sa,b]. However, many of the problem domains that NPs are applied to exhibit only approximate equivariance. It is therefore desirable to build approximately equivariant NPs through the use of approximately equivariant deep learning architectures. Whilst there exist several approaches to constructing approximately equivariant architectures, they are limited to CNN- and MLP-based models [Wang al., 2022a, van der Ouderaa et al., 2022, Finzi et al., 2021, Romero and Lohit, 2022]-which whhcc restricts their applicability to certain symmetry groups and often require modifications to the loss function used train the models [Finzi et al., 2021, Kim al., 2023a]. As there exist NPs that utilise variety of architectures, such as transformers and GNNs [Carr and Wingate, 2019, Nguyen and Grover, 2022, Kim et al., 2019, Feng et al., 2022], these approaches are not sufficient. We address this shortcoming through the development of an approach to constructing approximately equivariant models that is both architecture and symmetry-gne group agnostic. Importantly our approach can be realised without modifying the core architecture of existing equivariant models, enabling use in a variety of equivariant NPs such as the convolutional conditional NP [ConvCNP: Gordon et al., 2019], the equivariant CNP [EquivCNP Kawano et al., 2021]. the steerable CNP [SteerCNP; Holderrieth et al., 2021], the translation equivariant transformer NP TNP; Ashman et al., 2024], and the relational CNP [RCNP; Huang et al., 2023].
We outline our core contributions follows:
1. We demonstrate that under certain technical conditions, such as Hölder regularity and compactness, any non-equivariant mapping between function spaces can be approximated by equivariant mapping with fixed additional inputs (Theorems 2 and 3). This provides insight into how to con- struct approximately equivariant models, and generalises several existing approaches to relaxing equivariant constraints.
2.
We apply this result and the insights it provides to construct approximately equivariant versions of several popular equivariant NPs. The modifications required are very simple, yet effective: we demonstrate improved performance relative to both strictly equivariant and non-equivariant models on number of spatio-temporal regression tasks.
2 Background
We consider the supervised learning setting, where X. y denote the input and output spaces, and (X' & x V) be a collection of all finite y) E X' x y denotes an input-output pair. Let S = U ON
data sets, which includes the empty set Ø. Let SM (X y) be the collection of M input-output pairs and SCM = U Sm be the collection of at most M pairs. We denote a context and target set with De. D. E S, where 'Del = No. [D,]= = N₁. Let X E (2: Y. (2)) be the inputs and corresponding outputs of De, and let X, (X) N. Y, (D)) Mi be defined analogously. We denote single task as { = (De: Dz) = ((Xe, Y.), (X,, Y.)). Let P(X) denote the collection of V-valued stochastic processes on X. Let e denote the parameter space for some family of probability densities, e.g. means and variances for the space of all Gaussian densities.
2.1 Neural Processes
NPs [Garnelo et al., 2018a,b] can be viewed as neural-network- parametrisations of prediction
maps म: S + P(X)) from data sets S to predictions P()), where predictions are represented by y-valued stochastic processes on X' [Foong et al., 2020, Bruinsma, 2022]. Throughout, we denote the density of the finite-dimensional distribution of (D) at inputs X by p( X, D). In this work, we restrict our attention to conditional NPs [CNPs; Garnelo et al., 2018a], which only target marginal predictive distributions by assuming that the predictive densities factorise: p( X D)= II, p(y, /x,, D). We denote all parameters of CNP by w. CNPs are trained in meta-learning fashion, in which the expected predictive log-probability is maximised:
WM. arg max, (w)(W) where (Mi,(w) Eyce) E>> log Pa (Yt,n Kt. D.). (1)
Here, the expectation taken with respect to the distribution over tasks p(E). In practice, we only have access to finite number of tasks for training, so the expectation is approximated with average

---

over tasks. The global maximum is achieved if and only if the model recovers the ground-truth marginals [Proposition 3.26 by Bruinsma, 2022].
2.2 Group Equivariance
We consider equivariance with respect to transformations in some group G. For example, G can be the group of translations the group of rotations. Mathematically, group G is set endowed with binary operation G x G + G (denoted as multiplication) such that (i) (fg)h = f(gh) for all f.g. h E G: (ii) there exists an identity element e E G such that eg = ge g for all g E G; and (iii) every element g E G has an inverse g E G. A G-space is a space X for which there exists a function G x X - X called a group action (again denoted by multiplication) such that (i) ex x for all x E X; and (ii) f(zz) = ( for all f, E G and z E X. The notion of group equivariance is used to describe mappings for which, when the input to the mapping is transformed by some g E G, the output of the mapping is transformed equivalently. This is formalised in the following definition.
Definition 1 (G-equivariance), Let X and Y be G-spaces. Call a mapping X + Y G-equivariant if p(g=)= = gp(x) for all g E G and x X.
2.3 Group-Equivariant Conditional Neural Processes
The property of Geequivariance particularly useful in NPs when modelling ground-trute -ruth stochastic processes that exhibit G-stationarity:
Definition 2 (G-stationary stochastic process). Let X be a G-space. We say that stochastic process P P(X) G-stationary if. for J ~ and all g E G, f(()) ) is equal in distribution to f'g
Let P E P(X) be a ground-truth stochastic process. For dataset D S, define (D()) by integrating P againstt some density x (D), such that dTp = p dP. TD((D) is the Radon- Nikodym derivative of the posterior stochastic process with respect to the prior; intuitively, (f) is proportional to the likelihood, (f) = p(D f)/p(D), so it determines how the data is observed (e.g. under which noise?). Assume that (2) x that (2) P. We say that T'p is G-invariant if T (gD) o = (D)
Proposition (G-stationarity and G-equivariance). The ground-truth stochastic process P is G- stationary and T'p is G-invariant if. and only if. πP is G-equivariant.
Proof. This is Thm 2.1 by Ashman et al. [2024] with translations replaced by applications of g E G.
Thus, for data generated from a stochastic process that is approximately G-stationary, incorporating G-equivariance into NP approximations of the corresponding prediction map can serve as a useful inductive bias that can help generalisation and improve parameter efficiency. Intuitively, requiring NP to be equivariant effectively ties parameters together, which significantly reduces the search space during optimisation, enabling better solutions to be found with fewer data. The improved generalisation capabilities of G-equivariant models formalised by Elesedy and Zaidi [2021], Bulusu et al. [2021], Petrache and Trivedi [2023].
For CNP, assume that every marginal p( |x, D) is in some fixed parametric family with parameters (x: D)= E O. For example, could consist of the mean and variance for a Gaussian distribution. Let C(X, θ) denote the set of continuous functions X' + θ For CNP, define the associated parameter map क: S - C(²' θ) by (D))xx) = {x:: D). Intuitively, the parameter map maps a dataset 'D to function & (DD) giving the parameters (xD)(x) for every input x. Assume that X' is G-space. We turn S and C(X' θ) into G-spaces by applying to the inputs: gD E S consists of the input-output pairs (gx,, n), and g0(x; D) C(X, θ) is defined by g0(x; D) = 00g D)-
Definition 3 (G-equivariant CNP). A CNP is G-equivariant if the associated parameter map & is G-equivariant: ((gD) = (D) for all datasets D.
To parametrise G-equivariant CNP, we must parametrise the associated parameter map $. Here present a general construction by Kawano et al. [2021].
Theorem 1 (Representation of G-equivariant CNPs, Theorem by Kawano et al. [2021]). Let )) C _ be compact. Consider appropriate collection S'Sas C SCM. Then a function 4: S S.M
3

---

C(X, 0) is continuous, permutation-invaricn and G-equivariant if and only if is of the form
(DD p(Senc()) where
(Xmm (xm>ym)) = - ,x))
(2)
for some continuous Φ: y + RDD, an appropriate G-invariant positive-defiite kernel ub: X² (ie. 6xx,, ((,)) = (g 9Xm) for all g E G), and some conliruous and G-equivariant p: H
R
C(X, Q) with HH an appropriate G-invariant space of functions (i.e. closed under g E &G).
Theorem naturally gives rise to architectures which can be deconstructed into two components: an encoder and a decoder. The encoder, Sene : S + 2, maps datasets to some embedding space 2. The decoder, p: 2 - C(X' θ), takes this representation and maps to function that gives the parameters of the CNP's predictive distributions: p(y]x, x. D) = p(ylp(fem (D) (x)) where ([)) e O. For simplicity, we often view the decoder as function Z x X + ① and more simply write (D)exeD)), x). In Theorem I, both the encoder fene and decoder p are G-equivariant. Many neural process architectures are of this form, including the ConvCNP [Gordon et al., 2019], the EquivCNP [Kawano et al., 2021], J, the RCNP [Huang et al., 2023], and the TE-TNP [Ashman et al., 2024]. 2 We discuss the specifics of each of these architectures in Appendix B.
3 Equivariant Decomposition of Non-Equivariant Functions
In this section, we demonstrate that, subject to regularity conditions, any non-equivariant operator between function spaces can be constructed as, approximated by, an equivariant mapping with additional, fixed functions as input. These results motivate simple construction for approximately equivariant models, which we use construct approximately equívariant NPs in Section 3.1. We first illustrate the proof technique by proving the result for linear operators (Theorem 2) and then extend the result to nonlinear operators (Theorem 3).
Setup. Let (HI, (.,.) .)) be a Hilbert space of functions X' + R. Let G be a group acting linearly on H from the left. For every g E G and f E H., applying g f gives another function: of H. By linearity, for fi. f2 E H. g(11 + f2) = gfi + 9f2. Assume that H is separable and that Hi is G-invariant, meaning that, for all fi, f2 E HH and g G, (gff. 9 = (fi, f2). If = (2) and X' is a separable metric space, then IH is also separable [Brezis, 2011, Theorem 4.13]. Let B be the collection of bounded linear operators on H.. We say that an operator T B is of finite rank if the dimensionality of the range of T is finite. Moreover, we say that an operator T E B is compact if the image of every bounded set is relatively compact. Let (e:)121 be an orthonormal basis for H. Define Pa =Ei e (e,, .). The following proposition is well known:
-Finite-rank approximation of compact operators; e.g., Corollary by Brezis [2011].). Let T E B. Then T is compact if and only if there exists a sequence of operators of finite rank (Tn)x2 S B such that | - Tn + 0. In particular, one may take T,, = P,T, so T is a compact operator if and only if T P, | + O.
Proposition
Intuitively, compactness of an operator implics that its inputs and outputs can be well approximated with inite-dimensional basis. The following new result shows that every compact T B can be approximated with a G-equivariant function with additional fixed inputs: T = E,(: t1,..., t2a) where E,, is G-equivariant and t1. ,2. are the additional fixed inputs.
Theorem 2 (Approximation of non con-equivariant linear operators.). Let T E B. Assume that T is compact. Then there exists sequence of continuous nonlinear operators E,; Hl42n - H. n 2 1, and a sequence of functions (tm))21 S HH such that every E,, is G-equivariant,
Em Sosi,g2,.......
= (J (21.J2..... J2....,,2nn1. for all g E G and Ji, S2011 E HH, (3)
and IT - F...., 11,))l -
Proof. Observe that PaTS = Ein es (T ei, f) = ei f), with T, = e E H.. Define the TI,
continuous nonlinear operators F., H2 - H as F,, TI, G1, ..,....... Tn, en) = E (T By G-invariance of IH, these operators are G-equivariant:
En (gf, gT,, ge..ge...... = Σ. I.""""ge..rra (gT,, 9f) = 9 Db (, f) = gf
(4)
We note that the form of the embedded dataset fen: (D) diffeer slightly in the RCP and TE-TNN TE-TNP. However, in both cases the form in Equation can be recovered special cases. We discuss this more in Appendix B.

---

Since T is compact, ||- P, + 0 (Proposition 2), so |IT - En(: : which proves the result.
Ti... o Taa e1 ea.e + 0.
As example, consider G' to be the group of translations such that E,, is translation equivariant. Translation-eqiiariann mappings between function spaces can be approximated with a CNN [Kuma- gai and Sannai, 2020]. Therefore, Theorem 2 gives that T CNN( ft....,ttu tt.) where f1,....,4. . are additional fixed inputs that are given as additional channels to the CNN and can be treated as model parameters to be optimised. These additional inputs in the CNN break translation equivariance, because they are not translated whenever the original input is translated. The number & of such inputs roughly determines to what extent equivariance is broken: the larger k. the more non-equivariant the approximation becomes. This example holds true for any CNN architecture provided it can approximate any translation-equivaraann operator.
Theorem is not exactly what we set out to achieve: we wish to construct approximately equivariant mappings, not use equivariant mappings to approximate non con-vauivariant models. However, instead of applying Theorem directly to T, consider the decomposition T = Tequiv + T Tnon-euuv where Tequiy is in some sense the best "equivariant approximation" of T and T non-eqquiv i is the residual. Then, if we approximate Toon equiv with Enoon-qqivv using Theorem 2, and approximate Tequiv with some G-equivariant mapping Fequiv directly, we have T R: Frequiv + Fmon equiv (', 9,1,.-.. -.)) (.).. If k O, this approximation roughly recovers Fequir. the best equivariant approximation of 7: and, if & is increased, the equivariance starts break and the approximation starts wholly approximate T. Specifically, in Theorem 2, k = 2n determines the dimensionality of the range of the approximation Enn Therefore, a small k means that that one would deviate from T Teguiv in only a few degrees of freedom. The idea that k can be used to control the degree to which our approximation is G-equivariant inspires our training procedure in Section 3.1, where, with some probability, we set the additional inputs to zero. This forces the model to learn the 'best equivariant approximation', so that the number of additional fixed inputs has the desired influence of controlling only the flexibility of the non-equivariant component.
We now extend the result to any continuous, possibly nonlinear operator T: H - H. For T E B, it is true that 7 is compact if and only if 7|7- P, - (Proposition 3 in Appendix A). In generalising Theorem 2 to nonlinear operators, we shall use this equivalent condition. Roughly speaking, [ P,, TrPTl - 0 says that both the domain and range of T' admit finite-dimensional approximation, and the proof then proceeds by discretising these finite-dimensional approximations.
Theorem 3 (Approximation of non-equivariant operators.). Te::: H + H be a continuous, possibly nonlinear operator. Assume that || P,TPP, - O, and that T is (c,a for c, a in
the following sense:
(5) Moreover, assume that the orthonormal basis (ct)121 is chosen such that, for every n N and E G. span (ei....,ex) span (gei,.... ye.), meaning that subspaces spanned by finitely many basis elements (km)n21 C N, a sequence of continuous nonlinear operators F,,: Hl+ka + H, n 2 1, and a sequence of functions (tn)n21 HH such that every En is G-equivariant
[[T(u) T(p) ul for all u, v H..
invariant under the group action (*). Let M >> O. Then there exists sequence
(.fff....,99144.. =f (f1tk......tkk) for all g Gand fa...., fitk. H.
(6)
and
sup
[[T(u) En(u))Enu,,, .i...., thx)) +
M
If assumption (*) does not hold, then the conclusion holds with En(.,ti....... tt....,,tt... replaced by E, (t.,,,t,...,, ,t..).,tt..). We provide a proof in Appendix A.
Condition (+) says that the orthonormal basis (es)a21 must be chosen way such that applying the group action does not "introduce higher basis elements". Note that only one such basis needs to exist for the result to hold. An important example where this condition holds is HH = (S where 81 = R/E is the one-dimensional torus (00, 1] with endpoints identified); G the oneddimensional translation group: and (c:);21 the Fourier basis: (x( = c c22=kz Then (z - = eaakr es (z) x (x), so span (e*( - T)} span {ex). ). This example shows that the result holds for translation equivariance, which is the symmetry that we will primarily consider in the experiments. Importantly,
Theorem requires ||T- P.Il - 0 to hold. For this example, this condition roughly means that the dependence of T on high-freqqency basis elements goes to zero the frequency increases.
5

---

In Theorem 3, the sequence (ha)a21 determines how many additional fixed inputs are required to obtain a good approximation. For linear T (Theorem 2). kn grows linearly in n. In the nonlinear case, kn may grow faster than linear in n, so more basis functions may be required. We leave it to future work to more accurately identify how the growth of kn in depends on T. promising idea is to consider [Tgs fll as a measure of how equivariant a mapping is [Wang et al., 2022a].
3.1 Approximately Equivariant Neural Processes
Theorem provides a general construction of non-equivariant functions from equivariant functions, lending insight into how approximately equivariant neural processes can be constructed. Consider a G-equivariant CNP with G-equivariant encoder fenc and decoder p. The construction that we consider in this paper is to insert additional fixed inputs into the decoder p:
P PY(X, = II, ply,.(yneeeen. (x.. D) = IIn P(yn|>a,
tB))
(8)
where the decoder p: H C(X, 9) now takes in the dataset embedding fene well as B additional fixed inputs (to) S H. These become additional model parameters and break G- equivariance of the decoder. 'he more are included, the more non-equivariant the decoder becomes. Crucially, Theorem shows that including sufficiently many additional inputs eventually recovers any non-equivariant decoder. Conversely. by only including few of them, the decoder deviates from G-equivariance in only a few degrees of freedom. Hence, the number of additional fixed inputs determines which extent the decoder can become non-equivariant.
There are a myriad of ways in which (l)) can be incorporated into existing architectures for P. If p is a CNN a G-equivariant CNN, the additional inputs become additional channels, which can be either concatenated or added to the dataset embedding. We use the latter approach in implementation. For more general p. such as that used in the TE-TNP, we can employ effective alternative choices, as discussed in Appendix B. Theorem 3 also intimately connects to positional embeddings in wansformer-based L.LMs [Vaswani et al., 2017] and vision transformers [VIT; Dosovitskiy et al., 2021]. These architectures consist of stacked multi-head self-attention (MHSA) layers, which are permutation equivariant with respect to the input tokens. However, after tokenising individual words or pixel values into tokens, the underlying transformer which processes these tokens should not be equivariant with respect to the permutations of words pixels, as their position is important to the overall structure. Following Theorem 3, we can add word- or pixel- sesition-specifee inputs to break this equivariance. This corresponds exactly to the usual positional embeddings that are regularly used. This connection between breaking equivariance and positional encodings was also noted by Lim et al. [2023], where they interpret several popular positional encodings as group representations that can help incorporate approximate equivariance.
Recovering equivariance -nttof-distributionn We stress that equivariance is crucial for models to generalise beyond the training distribution-thie is the 'shared' component that is inherent to the system we are modelling. Whilst the non-equivariant component is able to learn local symmetry- breaking features that are revealed with sufficient data, these features do not reveal themselves outside the training domain. To obtain optimal generalisation performance, we desire that the model ignores the non-equ nonivvuiva component outside of the training domain and instead reverts to equivariant predictions. We achieve this in the following way: (i) During training, set t, 0 entirely with some fixed probability. This allows the model to learn and produce predictions for the equivariant component of the underlying system wherever the fixed additional inputs are zero. (ii) Forcefully set the additional fixed inputs to zero outside the training domain: (x() = 0 for x not close to any training data.² 3 This forces the model to revert to equivariant predictions outside the training domain, which should substantially improve the model', 's ability to generalise.
Our approach recovering equivariance out-of-distribution also connects Wang et al.'s (2022b) notion of e-approximate equivariance. Let aaennp be the learned neural process, and let NENN be the neural process obtained by setting 4 O in AEE-NP. If the predictions of AEE-N are not too different from those of TE-NP. because the fixed additional inputs only affect the predictions in a limited way, then TeE-NP is e-approximately equivariant in the sense of Wang et al. [2022b]. See Appendix D.
For example, the rectangle covering the maximum longitude and latitude of geographical locations seen during training of an environmental model.

---

4 Related Work
Group equivariance and equivariant neural processes. Group equivariant deep learning architec- tures are ubiquitous in machine learning, including CNNs [LeCun et al., 1989], group equivariant CNNs [Cohen and Welling, 20]6 transformers [Vaswani et al., 2017, Lee et al., 2019] and GNNs [Scarselli et al., 2008]. This is testament to the usefulness of incorporating inductive biases into models, with significant body of research demonstrating improved sample complexity and generali- sation capabilities [Mei et al., 2021, Flesedy and Zaidi, 2021, Bulusu et al., 2021, Zhu et al., 2021]. There exist number of NP models which realise these benefits. Notably, the ConvCNP [Gordon et al., 2019], RCNP [Huang et al., 2023], and TE-TNP [Ashman et al., 2024] all build translation equivariance into NPs through different architecture choices. More general equivariances have been considered by both the EquivCNP [Kawano et al., 2021] and SteerCNPs [Holderrieth et al., 2021], both of which consider architectures similar to the ConvCNP.
Approximate group equivariance. Two methods similar to ours are those of Wang et al. [2022a] and van der Ouderaa et al. [2022], who develop approximately equivariant architectures for group equivariant and steerable CNNs [Cohen and Welling, 20]6a, We demonstrate in Appendix C that both approaches are specific examples of our more general approach. Finzi et al. [2021] and Kim et al. [2023a] obtain approximate equivariance through a modification to the loss function such that the equivariant subspace of linear layers is favoured. These methods are less flexible than our approach, which can be applied to any equivariant architecture. Romero and Lohit [2022] only enforce strict equivariance for specific elements of the symmetry group considered. Their approach hinges on the ability to construct and sample from probability distributions over elements of the groups which both restricts their applicability and complicates their implementation. An orthogonal approach to imbuing models with approximate equivariance is through data augmentation. Data augmentation is trivial to implement; however, previous work [Wang al., 2022h] has demonstrated that both equivariant and approximately equivariant models achieve better generalisation bounds than data augmentation. Kim et al. [2023b] propose similar approach data augmentation, replacing uniform average over group transformations with an expectation over an input-dependent probability distribution. Implementation of this distribution is involved, and must be hand-crafted for each symmetry group.
5 Experiments
In this section, we evaluate the performance of number of approximately equivariant NPs derived from existing strictly equivariant NPs in modelling both synthetic and real-world data. We provide detailed descriptions of the architectures and datasets used in Appendix E. Throughout, we postfix to the name of each model the group it is equivariant with respect to, with the postfix G' denoting approximate equivariance with respect to group G. We shall also omit reference to the dimension when denoting a symmetry group (e.g. ((n) becomes 7). In all experiments, we compare the performance of three TNP-based models [Ashman et al., 2024]: non-equivariant, T, and T; three ConvCNP- based models [Gordon et al., 2019]: T, T using the approach described in Section 3. 1, and T using the relaxed CNN approach of Wang et al. [2022a]; and two EquivCNP- -bsee models [Kawano et al., 2021]: E' and E. For experiments involving large number of datapoints, we replace the Tas-based models with pseudo-token TNP-based models [PT-TNP; Ashman et al. 2024].
5.1 Synthetic 1-D Regression With the Gibbs Kernel
We begin with synthetic 1-D regression task with datasets drawn from a Gaussian process (GP) with the Gibbs kernel [Gibbs, 1998]. The Gibbs kernel similar to the squared exponential kerne,, except the lengthscale {(x) is function of position x. The non-stationarity of the kernel implies that the predictive map is not translation equivariant, hence we expect an improvement of approximately equivariant NPs with respect to their equívariant counterparts. We construct each dataset by first sampling change point, either side of which the GP lengthscale is either small (f(x) = 0.1) or large (f(x)= 4.0). The range from which the context and target points are sampled is itself randomly sampled, so that the change point is not always present in the data. We sample N, ~ (111, 64) and the number of target points is set as N,- = 128. See Appendix for complete description.
An implementation of our models be found at cembridge-mlg/aeng
7

---

Table 1: Average test log- ligeliieiisoos (1) for the synthetic 1-D GP and 2-D smoke experiments. For the 1-D dataset we used the regular TNP, while for the 2-D experiment we used the PT-TNP. Results are grouped together by model class. Best in-class result is bolded.
1-D GP
2-D Smoke
Model
ID Log-lik. (1)
OOD Log-lik. (1)
Log-lik. (t)
TNP
0.406 0.004
-1.3734
0.002
4.299 + 0.008
TNP (T)
-0.500 + 0.001
-0.430 + 0.007
4.181 + 0.011
TNP (T)
0.406 0.004
-0.444 0.007
4.715 0.010
ConvCNP (7)
-0.499
0.004
-0.442
0.006
3,637 = 0.041
ConvCNP (T)
-0.330 0.004
-0.42
0.006
3.827 0.011
RelaxedConvCNP (T)
-0.419 0.004
-0.405 + 0.007
4.006 0.010
Equ (E)
-0.504 0.004
-0.443
0.007
4.194 0.015
EquivCNP (E)
0.435 + 0.004
0.413 + 0.007
4.233 0.012
Context
Model
TE Model
Ground truth
(a) TNP.
(b) ConvCNP (T).
(c) EquivCNP (E).
(d) TNP (T).
(c) ConvCNP (T).
RelaxedConvCNP (T).
(g) EquivCNP (E).
(h) TNP (T).
Figure 1: A comparison between the predictive distributions on a single synthetic 1-D regression dataset of the TNP-, ConvCNP-, and EquivCNP-basee models. For the approximately equivariant models, we plot both the model's predictive distribution (blue), well as the predictive distributions obtained without using the fixed inputs (red). The dotted black lines indicate the target range.
We evaluate the log- lieleikelioood on both the in-distribution (ID) training domain and on out-of- distribution (OOD) setting in which the test domain is far away from the change point. Table presents the results. We observe that the approximate equivariant models are able to: 1) recover the performance of the non-equivariant TNP within the non-equivariant well as the equivariant models when tested OOD. We provide illustrative comparison of the predictive distributions for each model in Figure 1. More examples can be found in Appendix When transitioning from the low-lengthscale to the high-lengthscale region, the equivariant predictive distributions behave as though they are in the low-lengthscale region. This is due to the ambiguity as to whether the high-lengthseeae region has been entered. In contrast, the approximately equivariant models are able to learn that a change point always exists at x 0, resolving this ambiguity.
regime; and 2) generalise
We show in Appendix E.1 that the approximately equivariant models deviate only slightly from strict equivariance, thus being equivariant the approximate sense as we describe in Section 3.1. We also analyse the performance of the models with increasing number of basis functions in Appendix E. The biggest improvement is obtained when going from to 1 basis function, highlighting the importance of relaxing strict equivariance. We generally find that few basis functions suffice to capture the coumetry-breakiinn features, but using more fixed inputs than necessary does not hurt performance, thus justifying in practice the use of 'sufficiently high number' of fixed inputs.
5.2 Smoke Plumes
There are inherent connections between symmetries and dynamical systems, yet it also true that real world dynamics rarely exhibit perfect symmetry. Motivated by this, investigate the utility of

---

approximate equivariance in the context of modelling caanetrricalll-pprreeeee simulations generated from partial differential equations. We consider setup similar to Wang et al. [2022a]-we construct a dataset of 128 x 128 2-1) smoke simulations, computing the air flow in a closed box with smoke source. Besides the closed boundary, we also introduce a fixed spherical obstacle through which smoke cannot pass, and we sample the position of the spherical smoke inflow out of three possible locations. These three components break the symmetry of the system. We consider 25,000 different initial conditions generated through PhiFlow [Holl et al., 2020]. We use 22,500 for training and the remaining for testing. We randomly sample the smoke sphere radius r ~ [5, 30) and the buoyancy coefficient B ~ 1.0... For each initial condition, we run the simulation for a fixed number of steps and only keep the last state. We sub-sample a 32 x 32 patch from each state to construct dataset. We sample the number of context points according to Ne 4{10, 250} and set the remaining datapoints as target points. Table compares the average test log-likelihood of the models. As in the 1-D regression experiment, the approximately equivariant versions of each model outperform both the non-equivariant and equivariant versions, demonstrating the effectiveness of our approach in modelling complex symmetry- -erenker--bariigggg... features. We provide illustrations of the predictive means for each model in Appendix Figure 6.
5.3 Environmental Data
As remarked in Section 1, aspects of modelling climate systems adhere to symmetries such as equivariance. However, there are also unknown local effects which may be revealed by sufficient data. We explore this empirically by considering a real-world dataset derived from ERA5 [Copernicus Climate Change Service, 2020], consisting of surface air temperatures for the years 2018 and 2019. Measurements resolution of an hour. We also have access to the surface elevation at each coordinate, resulting in a 4-D input (x,, E I). We consider measurements collected from Europe and from central S. We train each model on Europe's ' 2018 data, and test on both Europe's 's and central US' 2019 data. Because the LNN-based architectures have insufficient support for 4-D convolutions, we first consider a 2-D experiment in which the inputs consist of longitudes and latitudes, followed by a 4-D experiment consisting of all four inputs upon which the transformer- based architectures are evaluated.
collected at latitudinal and longitudinal resolution of 0.5%, and temporal
2-D Spatial Regression. We sample datasets spanning 16 across each axis. Each dataset consists of maximum of N 1024 datapoints, from which the number of context points are sampled according to N₂ ~ [[[NN110 [N/3]], J. The remaining datapoints are set as target points. Table 2 presents the average test likeelolooodd on the two regions for each model. We observe that approximately equivariant models outperform their equivariant counterparts when tested on the same geographical region as the training data. As the central US data falls outside the geographical region of the training data, we zero-out the fixed inputs, so that the predictions for the approximately equivariant models depend solely on their equivariant component. Surprisiggly, they also outperform their equivariant counterparts. This suggests that incorporating approximate equivariance acts to regularise the equivariant component of the model, improving generalisation in finite-data settings such as this. In Figure 2, we compare the predictions of the PT-TNP and EquivCNP- test dataset with the ground-truth data and the equivariant predictions made by the same models with the fixed inputs zeroed out. The predictive means of both models are almost indistinguishable from ground-trutt We can gain valuable insight into the effectiveness of our approach by comparing a plot of the difference between the approximately equivariant and the equivariant predictions (Figures 2 and 2i) to that of the elevation map for this region (Figure 2c). As elevation is not provided as an input, yet is crucial in predicting surface air temperature, the approximately equivariant models can infer the effect of local topographical features on air temperature through their non -quivirivar component. As with the 1-D GP experiment, analyse the degree of approximate equivariance and the performance with an increasing number of fixed inputs in Appendix E.3, drawing similar conclusions.
models for a single
4-D Spatio-Tempp
Regression. In this experiment, we sample datasets across 4 days with measurements every day and spanning 8 across each axis. Each dataset consists of maximum of N 1024 datapoints, from which the number of context points are sampled according to [N/3]]. The remaining datapoints are set as target points. We provide results
No ~ in Table 2. Similar to the 2-D experiment, we observe that the approximately equivariant PT-TNP outperforms both the equivariant PT-TNP and non-equivariant PT-TNP on both testing regions.
A longitude/ latitude range of [35", 60%] [10", 45"] and [- 1200, -80" [30%, 50%). respectively.

---

Table 2: Average test log- lielliiilihodss (1) for the 2-D and 4-D environmental regression experiment.
Results are grouped together by model class. Best in-class result is bolded.
2-D Regression
4-D Regression
Model
Europe (1)
US (1)
Europe (1)
US (t)
PT-TNP
1.14 0.01
<-10
0.94 0.01
< <-10
PT-NNP (2)
1.06
1.14 0.01
0.73 0.01
PT-TNP (7)
1.22
0.55 = 0.01
1.21 + 0.01
0.76 0.01
ConvCNP (7)
1.11
0.12 0.02
ConvCNP (7)
1.18 118+0000
0.15 0.02
RelaxedConCCNP (7)
1.20 0.01
0.22 0.02
EquivCNP (L)
0.64 0.02
EquivCNP (E)
1.36 + 0.01
0.69 0.01
300
315
24
700
315
30
4.5%%
4.5%
4.5%
(a) Context.
(b) Ground-truth
(c) Elevation.
20
40.5%
40.5%
PT-TNP: T.
(e) PT- T.
(1 PT-TNP: T -TT.
4
4
4PN
0.9%%
0..%%
0..5%
1..9
3PB 23 wn 10.0
20: 3011 30.91
35.9%
3PR 20 2011
10.0
(g) EquivCNP: E.
(h) EquivCNP: E.
(i) EquivCNP: E
Figure 2: A comparison between the predictive distributions of the equivariant (left column) and approximately equivariant (middle column) components of the PT-TNP (I') and EquivCNP (E) models on single (cropped) test dataset from the 2-D environmental data experiment.
6 Conclusion
The contributions of this paper are two-fold. First, we develop novel theoretical results that provide insights into the general construction of approximately equivariant operators which is agnostic to the choice of symmetry group and choice of model architecture. Second, we use these insights to construct approximately equivariant NPs, demonstrating their improved performance relative to non- equivariant and strictly equivariant counterparts on number of synthetic and real-world regression problems. We consider this work to be an important step towards understanding and developing approximately equivariant models. However, more must be done to rigorously quantify and control the degree to which these models depart from strict equivariance. Further, we only considered simple approaches to incorporating approximate equivariance into equivariant architectures, and provided empirical results for relatively small-scale experiments. We look forward to addressing each of these limitations in future work.
10

---

Acknowledgements
CD is supported by the Cambridge Trust Scholarship. AW acknowledges support from a Turing AI fellowship under grant EP/VO25279/1 and the Leverhulme Trust via CFI. RET is supported by The Alan Turing Institute, Google, Amazon, ARM, Improbable, EPSRC grant EP/T005386/1, and the EPSRC Probabilistic AI Hub (ProbAL, EP/Y028783/1).

---

References
Matthew Ashman, Cristiana Diaconu, Junhyuck Kim, Lakee Sivaraya, Stratis Markou, James Re- queima, Wessel P Bruinsma, and Richard E Turner. Translation-eoo Translation-qqiivarianl transformer neural processes, In International conference on machine learning. PMLR, 2024.
Haim Brezis. Functional analysis, Sobolev spaces and partial differential equations, volume Springer, 2011.
Wessel P. Bruinsma. Convolutional Conditional Neural Processes. PhD thesis, Department of Engineering, University of Cambridge, 2022. URL http::::/wwwwwwpwoeee cam. handle/1810/354533
cam ac uk/
Srinath Bulusu, Matteo Favoni, Andreas lpp. David I Müller, and Daniel Schuh. Generalization capabilities of translationally equivariant neural networks. Physical Review D, 104(7):07450,, 2021.
Andrew Carr and David Wingate. Graph neural processes: Towards Bayesian graph neural networks. arXiv preprint arXiv: /902. 10042, 2019.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference machine learning, pages 2990-2999. PMLR, 201
Taco S Cohen and Max Welling. Steerable emns. arXiv preprint arXiv:1612.08499, 2016b.
Copernieus Climate Change Service. Near surface meteorological variables from 1979 2019 derived from bias-corrected reanalysis, 2020. URL http:///d....mmaa..oooooo atimate..cpernicuu eu/
cdsapp#! /home.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition scale. In International Conference on Learning Representations, 2021. URL https://se /ppnreviir. net./forumm.d--iicbaarrr
Bryn Elesedy and Sheheryar Zaidi. Provably strict generalisation benefit for equivariant models, In International conference on machine learning, pages 2959-2969. PMLR, 2021.
Leo Feng, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Latent bottle- necked attentive neural processes. arXiv preprint arXiv:2211.0845: 2022.
Mare Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In International Conference on Machine Learning, pages 3165-3176 PMLR, 2020.
Mare Finzi, Gregory Benton, and Andrew Wilson. Residual pathway priors for soft equivariance constraints. Advances in Neural Information Processing Systems, 34:30037-30049 2021.
Andrew Foong, Wessel Bruinsma, Jonathan Gordon, Yann Dubois, James Requeima, and Richard Turner. Meta-learning stationary stochastic process prediction with convolutional neural processes. Advances in Neural Information Processing Systems, 33:8284-8229 2020.
Marta Garnelo, Dan Rosenbaum, Christopher Maddison, Tiago Ramalho, David Saxton, Murray Shanahan, Yee Whye Teh, Danilo Rezende, and SM Ali Eslami. Conditional neural processes, In International conference machine learning, pages 1704-1713. PMLR, 2018a.
Marta Garelo,, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv:1807.00162) 2018b.
Mark Gibbs. Bayesian Gaussian processes for regression and classification. PhD thesis, Citeseer, 1998.
Jonathan Gordon, Wessel P Bruinsma, Andrew YK Foong, James Requeima, Yann Dubois, and Richard Turner. Convolutional conditional neural processes. arXiv preprint arXiy: 1910. 13556, 2019.
12

---

Peter Holderrieth, Michael Hutchinson, and Yee Whye Teh. Equivariant learning of stochastic fields: Gaussian processes and steerable conditional neural processes. In International Conference on Machine Learning, pages 4297-4307. PMLR, 2021.
Philipp Holl, Nils Thuerey, and Vladlen Koltun. Learning to control pdes with differentiable physics. In International Conference on Learning Representations, 2020. URL https: https://opentt //tt::ee//ppeeevirr. net/ SoerrumiidHyySSinnee
Daolang Huang, Manuel Haussmann, Ulpu Remes, ST John, Grégoire Clarté, Kevin Sebastian Luck, Samuel Kaski, and Luigi Acerbi. Practical equivariances via relational conditional neural processes. arXiv preprini arXiv:2306. 10011 2023.
Andrew Jaegle, Felix Gimeno, Andy Brock, Oriol Vinyals, Andrew Zisserman, and Joao Carreira. Perceiver: General perception with iterative attention. In International conference on machine learning, pages 4651 4664. PMLR, 2021.
Makoto Kawano, Wataru Kumagai, Akiyoshi Sannai, Yusuke Iwasawa, and Yutaka Matsuo. Group equivariant conditional neural processes. arXiv preprint w:Xiv:2102.0875: 2021.
Hyunjik Kim, Andriy Mnih, Jonathan Schwarz, Marta Garnelo, Ali Eslami, Dan Rosenbaum, Oriol Vinyals, and Yee Whye Teh. Attentive neural processes. arXiv preprint arXiv: arXivv/901 /ete://0105550 .05761, 2019.
Hyunsu Kim, Hyungi Lee, Hongseok Yang, and Juho Lee. Regularizing towards soft equivariance under mixed symmetries. In International Conference on Machine Learning, pages 16712-16727. PMLR, 2023a.
Jinwoo Kim, Dat Nguyen, Ayhan Suleymanzade, Hyeokjun An, and Seunghoon Hong. Learning prob- abilistic symmetrization for architecture agnostic equivariance. Advances in Neural Information Processing Systems, 36: 18582-186113 2023b.
Wataru Kumagai and Akiyoshi Sannai. Universal approximation theorem for equivariant maps by group CNNs. arXiv preprint arXiv:2012.13882, 2020.
Yann LeCun, Bernhard Boser, John Denker, Donnie Henderson, Richard Howard, Wayne Hubbard, and Lawrence Jackel. Handwritten digit recognition with propaaatonn network. Advances in neural information processing systems, 2, 1989.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: framework for attention- based permutation-invarian neural networks. In International conference on machine learning, pages 3744-3753. PMLR, 2019,
Derek Lim, Hannah Lawrence, Ningyuan Teresa Huang, and Erik Henning Thiede. Positional encodings as group representations: A unified framework. 2023.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint 17110: 2017.
Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory, pages 3351-3411 PMLR, 2021.
Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. arXiv preprint arXiv:2207.041179 2022.
Mircea Petrache and Shubhendu Trivedi. approximatioo-geeeaaliaaaiii trade-offs under (approxi- mate) group equivariance. Advances in Neural Information Processing Systems, 36:61936-61959, 2023.
David W Romero and Suhas Lohit, Learning partial equivariances from data. Advances in Neural Information Processing Systems, 35:36466-36478 2022.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical
image segmentation. In Medical Image Computing and Computer-Assister Intervention-MICCAA 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pages 234-241. Springer, 2015.
13

---

Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE transactions on neural networks, 20-1):61-80, 2008.
Tycho van der Ouderaa, David W Romero, and Mark van der Wilk. Relaxing equivariance constraints with non-stationary continuous filters. Advances in Neural Information Processing Systems, 35: 18818-33830, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Rui Wang, Robin Walters, and Rose Yu. Approximately equivariant networks for imperfectly symmetric dynamics. In International Conference on Machine Learning, pages 23078-23091. PMLR, 2022a.
Rui Wang, Robin Walters, and Rose Yu. Data augmentation vs. equivariant networks: A theory of generalization on dynamics forecasting. arXiv preprint arXiv:2206. 09450, 2022b.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
Sicheng Zhu, Bang An, and Furong Huang. Understanding the generalization benefit of model invariance from a data perspective. Advances in Neural Information Processing Systems, 34: 4328 4341, 2021.
14

---

A Proof of Theorem 3
Proposition 3. Let I E B. Then T is compact if and only if - P.TPA + 0.
Proof. Assume that - P,TPP, - Then T is the limit of sequence of finite-rank operators. so T is compact (Proposition 2).
On the other hand, assume that T is compact. Then - - inequality to bound
(Proposition 2). Use the triangle
|T - P. IT P,T + P.T P,TPA
T P, PP T - TP,
(9)
We already know that ||T - P.T + O, and it is true that | = Finally, since compact, the adjoint 7 is also compact. Hence, again by Proposition 2,
|| - P,, [I TP,))
- +
(10)
which proves that ||T - 1.TP, +
Theorem 3 (Approximation of non- non-equivariant operators.). Let : H + He a continuous, possibly nonlinear operator. Assume that ||T P.TPP, - O, and that T is (c,a)-Hölder for e, or in the following sense:
(5) Moreover, assume that the orthonormal basis (ez)i21 is chosen such that, for every n. E N and g E G. span (ei...., Pn] span .. meaning that subspaces spanned by finitely many basis elements (kn)n21 S N, a sequence of continuous nonlinear operators En: H2+ka + H., n 2 1, and sequence of functions (nn)n21 Hil such that every En is G-equivariant,
(T(u) T(b) cllu w11a for all u. H.
invariant under the group action (*). Let M >> O. Then there exists sequence
F.igfh....99tth... = gE. ;.11..). for all g and f1...., f1tk, H.
(6)
and
sup DellIIas
[T((u) Eu, th +.
(7)
If assumption (*) does not hold, then the conclusion holds with En( fall,,i.... ,t....., .t.........))) tx..) replaced by En st.....,ty,) We provide a proof in Appendix A.
Proof. Let E 0. Choose N such that | P,TPal < f,, and choose h > such that a = Je. Set in = h/v2n. Consider the following collection of vectors:
A = (jhnet : +1 ..... +1... h(M x ... {jhnen : 0, +1, +1, [M/h..] (11)
By construction of A. for every e H. such that |u||| s M, there exists an E A such that | all2 < nh 1 < 122
Let h: IR + 00, xc) be continuous function with support equal to (-h, Set H H + [0, oc).
k(u,v) k( lu el]) Consider the following map:
Σ k(u, ai)4 E kuu,, a;)
H MI H. Elu, 01 4 9A 44,4)A) =
(12)
where 0/0 defined This map E is continuous: the numerator is continuous H-valued function, and the denominator is continuous Ik-valued function which is non-zero wherever the numerator non-zero. Moreover, by G-invariance of & G-invariant, so E is G-equivariant:
Al k(gu, ga. ga:)gt;
| a(u,ai). gE(u, an, (1,4)....) (13)
E(gu, ga,, gli,g.,,.....
E k(gu, ga:)
Σ 1 k(u, az)
15

---

Now set tq = (a,TTP))a.) Consider some E HI such that ||u|| S M. By construction of A, Σ k(Pnu, a;) > O. Therefore, by (c, (H.lder continuity of
[IPPT - E(P,u, an, (,,,..,),....))) .1,...)
E k(P, u. a, (P,TT,))))
EA k(P u( (P.TP..a))
(14)
EA (
Σ ((,u
E k(Pn u a)) - (P,TT,,((,))) (PP
(15)
((Puu()
k(Pau u, a;) (P,TT,)))) - (P.TP,))aa)
(16)
Σ E k(Pnu,
1=1 k(PRu, 0z (a)ccl P a:ll (i))
ch 2
(17)
where (i) uses the triangle inequality; (ii) (P,TT))(u) (P.TP))(PP.)) and (c,a)-Hölder conti- nuity of P.TPa and (iii) that [P.u a:ll implies that k(P, a:) 0. Therefore, for all u E such that ||u||| _ M,
(T(u) - F(P,u, )II (T (P.TT.) + [IP E(P,.,.....
(18)
Finally, if condition (x) holds, then Pa is G-equivariant, which show now. Let u H and consider E G. Since span fei,...,enI = span fgei..... ]en) and ge,....: forms an orthonormal basis for span, (ei,...,C,) C, we can also write P, Ei=1 ge; ge; ).
Pa(gu) =
Σ ge, (ge, gu) = Σ ee u) = ((Pu))
(19)
Therefore, in the case that condition (x) holds, Pn can be absorbed into the definition En.
From the proof, we find that k, = 2|A|| = 2(1 + 2[22m//1 I 2(2 2/2+2/2 and cha = fe, which implies that / = (e/ (//a so k, S 2(2 + 22 (E/2c)))-/ This growth estimate is faster than exponential. Like the linear case shows (Theorem 2), better estimates may be obtained with constructions of E that better exploit structure of T.
B Equivariant Neural Processes
In this section, we outline the equivariant NP models we use throughout our experiments, and how their approximately equivariant counterparts can easily be constructed using the results from Section 3. In all cases, the asymptotic space and time complexity remains unchanged.
B.1 Convolutional Conditional Neural Process
The ConvCNP [Gordon et al., 2019] translation equivariant NP and constructed in exactly the form shown in Equation 2, with db X² + IR an RBF kernel with learnable lengthscale and Φ(y.) = 1J". The functional embedding discretised and passed pointwise through an MLP to some final functional representation ē(D) X + RR,,, where X denotes the discretised input domain. implemented as CNN through which ē(D) is passed through together with another RBF kernel tip which maps back to a continuous function space. We provide pseudo-code for forward pass through the ConyCNP in Algorithm 1.
16

---

Algorithm 1: Forward pass through the Con-
Algorithm 2: Forward pass through the Con-
vCNP (1) for off-the-grid data.
vCNP (T) for off-the--ritt thee-grid data.
Input: p p= (CNN, 1/y,, 89, and density 7
Input: o (CNN, v/y). tb, density y, and fixed
Context
input t. Context
D. (X.:: Y)) = {(xan: (on))ee
D. = (X.: Y.) = {(xemn Ye.n) No
and target X (xe,m m
and target X₁ fxe. m==1
begin
begin
lower, upper + range UX)):
lower, upper + range (Xt U
T 1 glowr, upper, ?);
[& 1 t- gloee, upper, 7);
h Ide n= [1,
ye.n] J" v(x; Xx,n);
h, - No (1,y (Xx: Xxn)::
h: + h !
[1)
h + h /h!")
h; 6 MLP(h:):
h: t- (
(f(x) !T= + CNN (f&,, H.
(f(x))) + CNN ((&, h. T=)T
Om Σx f(x))) -.);;
8. Σ f(x))0 8,);
return p( mm
return p( (m))
end
end
To achieve approximate translation equivariance, we construct a learnable function & X - R D. using an MLP which represents the fixed inputs. We sum these together with the resized functional embedding, resulting in the overall implementation p(e(D), ..,..., (n) CNN(E(D) + (). Note that the number of fixed inputs is B D2 We provide pseudo-code for a forward pass through the approximately translation equivariant ConvCNP in Algorithm 2.
B.2 Equivariant Conditional Neural Process
The EquivCNP [Kawano et al., 2021] is a generalisation the ConvCNP to more general group equivariances. As with the ConvCNP, it constructed in the form shown in Equation 2. We achieve E- cquivariance with d : X² - I an RBF kernel with learnable lengthscale and Φ(y:) = ly? 1| The functional embedding discretised and passed pointwise through an MLP to
final functional representation ē(D) X >> T D=, where X denotes the discretised input domain. / implemented as a group equivariant CNN [Cohen and Welling, 20] together with another RBF kernel vpp which maps back continuous function space. The approximately E-equivariant EquivCNP constructed in an analogous manner to the approximately translation equivariant ConvCNP.
B.3 Translation Equivariant Transformer Neural Process
The TE-TNP [Ashman et al., 2024] a translation equivariant NP. However, unlike the ConvCNP, the function embedding representation differs slightly to that in Equation 2, and is given by
e(D)() X R= cat
an (z (()))) W Wv.h Wo
(20)
where
cd(ad W Wo. d(x)),
on (Zo, ((y:), x1)
(21)
5== en(an Wa. W (ols))
Here, y implemented using an MLP. This is a partially evaluated translation equivariant muat cross-attention (TE-MHCA) operation. with attention weights computed according to Equation 21. We refer to Φ(y:) as the initial token embedding for y.. is implemented using translation equivariant multi-head self-attention (TE-MHSA) operations which update the token representations. These are combined with TE-MHCA operations which map the token representations space. We provide pseudo-code for forward pass through the TE-TNP in Algorithm 3.
continuous function
17

---

Algorithm 3: Forward pass through the TE-
Algorithm 4: Forward pass through the TE-
TNP (2))
TNP (T).
Input: P (TE--MSAA TE-MHSA (9) JE..
Input: p (TE-MHSA TE-MHSA (0)
6. Context
6, fixed input Context
D. = (Xe: c)) = (xx,.n Yea)
D. = (X., Y.) = ((xe, Ve,) LN hem
and target X, = {xe.m (Xc,mmmm
and target X, )
begin
begin
Ze,n ((yc.n);
Zcnn @(Ve.n) t(xx..):
2t.m < Z0i
Zt,m <- Z0i
L..... do for
for
= I..... L do
2t
TE-MHCA (zi.m, Ze, Xt.m, Xo);
TE (Zi,m, Los Xt.ms X.):
Zonu 1 + TE-MHSA (O) (Zen X.):
Zc,n J + TEMHSA (Zo, Xo)i
end
end
On MLP(Z
8m MLP(Z..))
return pl. (m)):
return pl (mm)m))
end
end
Unlike the ConvCNP and EquivCNP, do discretise the input domain of the functional representation. It is not clear how one would sum together a fixed input and the functional embedding without requiring infinite summations over the entire input domain. Thus, take alternative approach in which the initial token representations for context set are modified by summing together with the fixed imput value the corresponding location. This modification still conforms the form given in Equation 8, and is simple to implement. We provide pseudo-code for a forward pass through the approximately translation equivariant TE-TNP in Algorithm
B.4 Pseudo-Tokee Token Translation Equivariant Transformer Neural Process
The pseudo-token TE-TNP (PT-TE-TNP) is an alternative the TE-TNP which avoids the O(N2) computational cost of regular transformers through the use of pseudo-tokens [Feng et al., 2022, Ashman et al., 2024, Lee al., 2019]. The functional embedding representation for the PT-TE-TNP is given by
(D)))) X D= cat Σ
on (zo, Um, "x)Ju W V Wo
(22)
where
N
Um cat
an ( (um,0, o(y)), Vm vm. vm.
" Wo
(23)
and
N N
Vm Vm, +
(24)
The attention mechanism used in the PT-TE-TNP the same as the TE-TNP, and given in Equation 21. Intuitively, the initial token embeddings for the dataset D are 'summarised" by the pseudo-tokens {um 1 and pseudo-token input locations (vm Wm) p implemented using series of TE-MHSA and TE-MHCA operations which update the pseudo-tokens and form a mapping from pseudo-tokens to a continuous function space. We provide pseudo-code for a forward pass through the PT-TE-TNP in Algorithm 5. Note that this assumes that the perceiver-style approach taken [Jaegle et al., 2021], as in the latent bottlenecked attentive NP Feng et al. [2022]. The induced set transformer approach of Lee al. [2019] can also be used.
18

---

Algorithm 5: Forward pass through the PT-
Algorithm 6: Forward pass through the PT-
TE-TNP (2).
TE-TNP (T).
Input: p=
Input: P
(TE-MSA (9, TE-MHCA)) TE-MHCA
TE=1. (TE-MHSA (9, TECA TE-MHCA, (9)
1E=1
V [V U (u.,))==n Zo.
6. V = (mmo)s=11
{us,0) ) Zo-
Context
Context
De = (Xe: Y.) = {(xo.n: Yon)llm
D. = (X Y.) {(xem ( Non)
and larget X, (xee. m==
and larget X (x). 3 M
begin
begin
Zenn f o(yan);
Zen t @(ye,))
Vm <- Vm, + 1/N E Xo,si
Vm Vm. + 1/N Eit Xe,ii
y f- u1,05
u,
u8o t(V.)):
Zt.m f Z0;
Bt.m
2o:
for = L..... do
for L
L....: L....LLo do
Uj f- TEEMHCAA (8) (u,, Zo, V3, X.);
Uj f TE-MHCA (8) (u,, Ze: Vj, X.):
{u, /u f TE-MHSA (O) (U, V);
{1 f- TE-MHSA (9) (U, V):
Ze,m
Zt.m
TE-MHCA (2) (Zi.m, U. Zi,m; V);
TE-MHCA (Zi,m, U, Zi,m. V);
end
end
Om + MLP(Zi.m):
6n MLP(2(.m):
return p(.10m))
return p(-10m.)
end
end
We incorporate the fixed inputs by modifying summing together the initial pseudo-token values with the fixed inputs evaluated at the corresponding pseudo-token input location. We provide pseudo-code for forward pass through the approximately translation equivariant PT-TE-TNP in Algorithm 6.
B.5 Relaxed Convolutional Conditional Neural Process
The RelaxedCOnvCNNP is equivalent to the ConvCNP with the RelaxedCNN of Wang et al. [2022a] used in place of standard CNN. The RelaxedCNN replaces standard convolutions with relaxed convolutions, which we discuss in more detail in Appendix C. In short, the relaxed convolution is defined as
( Σ (0)
(k+ t ())
(25)
This modifies the kernel weights (u(u-uv) with the input-dependent function wi()), which is our fixed input. To enable the fixed inputs be zeroed out, we modify this slightly as
/ JG 1=1
(ks & (u))
(26)
so that when (v) we recover the standard G-equivariant convolution. We make (v(v) learnable function parameterised by an MLP, as with the other approaches.
C Unification of Existing Approximately Equivariant Architectures
In this section, we demonstrate that the input-dependent kernel approaches of Wang et al. [2022a] and van der Ouderaa et al. [2022] are special cases of our approach. Throughout this section, we consider scalar-valued functions. We define the action of g E G on f:G. R
(g (w) (l))
(27)
We begin with the approach of Wang et al. [2022a], which defines the relaxed group convolution as
/ Σ (1))
(I f))))
(28)
J
19

---

where uu(e) is the left Haar measure on G. Observe that this replaces single kernel with set of kernels {ky Mi whose contributions linearly combined with coefficients that vary with v G. This can be expressed as
(A: * f))
Σ
o)dd(())
Σ un) E(S, ,,, ...,, wr) (29)
((f.w))(a))
where correspond to the fixed inputs and E is used to denote equívariant operators w.r.t. the group G. Clearly each E, is G-equivariant with respect its inputs, as applying a transformation g G to the inputs gives
Eigg f.g mi)(u) /
(()))(())((()))))))))))))))))))) gv)dp())))
(30)
= =9 Fa(J, ())
Since the sum of G-equivariant functions is itself G-equivariant, E(; wi,..., wr.(()) G- equivariant with respect to its inputs.
The approach of van der Ouderaa al. [2022] is more general than that of Wang et al. [2022a], and relax strict equivariance through convolutions with input-dependent kernels:
(k [)(w) = k(v (()
(31)
Define fixed input t(u) u. We can express the above convolution as G-equivariant operation on t and f:
((,()(u)= .f (()))))))
(32)
To demonstrate G-equivariance, consider applying a transformation y E G to the inputs:
/ k(U ((o)
Elg f.g. -t(())
/ k(U -u,
(33)
[ k(o (g (o)
= g (().)(u).
Thus, this method is also a special case of ours.
D Achieving e-Approximate G-Equivariance
Wang et al. [2022b] introduce a useful metric for assessing the degree to which mappings are approximately equivariant. Specifically, let x: X Y denote some mapping between G-spaces X and Y. Let d: Y x Y + R be metric on Y. We say is c-approximately G-equivariant if for all g E and x X
d(T(gr), gT(x)) Se.
(34)
Let TAE-NP S + C(X', θ) be a learned neural process that breaks equivariance with fixed additional inputs (4x ( and let NEAP: S - C(X²' +) the strictly equivariant neural process obtained by setting 4, = 0 in TA If the predictions of Aa.nnP are not too different from those of TENNP. because the fixed additional inputs only affect the predictions in a limited way, then in this appendix we argue that TAE-P e-approximately equivariant.
20

---

To assume that predictions of TAE-NP are not too different from those of TE-NP. assume that there exists some metric d and 0 such that, for all D, d(TAENP(D), TTNP I c. For example. the metric can be the root-mean-squared-emns (RMSE) between the mean functions + [z)( of two stochastic processes by integrating over the input domain X. In practice, we will find that AETNE-PP(() TAE-NP has a finite receptive field-in the sense that the oulput stochastic process has only lo- cal dependencies then, provided we zero out the additional fixed inputs as described in Section 3.1, ((TAAENP(D) TE)) S e will hold over the entire input domain.
(D)) Σ holds for D over the training domain. If we assume that
Finally, to show that TAE-NP is e-approximately equivariant, we use the triangle inequality:
d( ((gTAEN((D), (AENN(())))))) S d ((gTAENP(D))
(D) + d(gTENNP (D), TE ()D)
Se ()) (ANNN))))
(35)
S 2c,
(36)
where d(gTEeN(D)) TE-NP( = 0 because TE-NP is strictly equivariant.
E Experiment Details and Additional Results
E.1 Synthetic 1-D Regression
We consider synthetic 1-D regression task using samples drawn from Gaussian processes (GPs) with the Gibbs kernel with an observation noise of 0.2. This kernel a non-stationary generalisation of the squared exponential kernel, where the lengthscale parameter becomes function of position 1(x):
21(x)((xx) k(x, 21; - (1 (TOP +(=))+((>'))
(x x')² [(x))² 1(x1)²
exp
We consider 1-D space with two regions with constant, but different lengthscale one with ( - 0.1. and one with = 4.0. The lengthscale changepoint is situated æ O. We randomly sample with a 0.5 probability the orientation (left/right) of the low/high lengthscale region. Formally, ((x) = (0.13 + 4(1 4)1-=))ālr < 0] (0.1 - 3) + 43) 0), where 3 ~ Bern(..5). For each task, we sample the number of context points N ~ [f[1, 64) and set the number of target points to N₂ 128. The context range [x,,min: Xx.max] (from which the context points are uniformly sampled) is an interval of length 4, with its centre randomly sampled according to 4-7.7 for the ID task, and according to 2/13.27 for the OOD task. The target range is Xt, X1, mma] = [Xe.min - Xe + 1]. This is also applicable during testing, with the test dataset consisting of 80,000 datasets.
We use an embedding / token size of D= = 128 for the bNP-based models and D = 64 for the ConvCNP-basee based ones, and decoder consisting of an MLP with two hidden layers of dimension Dz. The decoder parameterises the mean and pre-softplus variance of a Gaussian likelihood with heterogeneous noise. Model specific architectures
follows:
The initial context tokens are obtained by passing the concatenation [r,y, 1] through an MLP with two hidden layers of dimension D.. The initial target tokens are obtained by passing the concatenation [x. 0, through the same MLP. The final dimension of the input acts as a 'density" channel to indicate whether layers of self-attention and cross-attention blocks, each with H = 8 attention heads with dimensions Dv Dox 16. In each of the attention blocks, apply a residual connection consisting of layer- normalisation to the input tokens followed by the attention mechanism. Following this, there is another residual connection consisting of layer- normalistaiii followed by pointwise MLP with two hidden layers dimension D..
not observation is present. The TNP encoder consists of nine
TNP (T)) For the TNP (7) model we follow Ashman et al. [2024]. The architecture is similar to the TNP model, with the attention blocks replaced with their translation equivariant counterparts. For the translation equivariant attention mechanisms, implement of RH x R + g an MLP with
21

---

two hidden layers of dimension D,. The initial context token embeddings are obtained by passing the context observations through an MLP with two hidden layers of dimension D.. The initial target token embeddings are sampled from standard normal. Pseudo-code for a forward pass through the TNP (7) can be found in Appendix B.3.
TNP (T) The architecture of the TNP (T) similar to that of the TNP (T), with the exception of the extra fixed inputs that are added to the context token representation. These are obtained by first performing a Fourier expansion to the context token locations, and then passing the result through an MLP with two hidden layers of dimension D2. We use four Fourier coefficients, zero out the fixed inputs outside of [-7, 7]. and during training we drop them out with a probability of 0.5. Pseudo-code for a forward pass through the TNP (T) can be found in Appendix B.3.
CNnyCNP (7) For the ConvCNP model, we use a CNN with 9 layers. We use C² 64 channels, kernel size k = with a stride of one. The input domain is discretised with 46 points per unit. The decoder uses five different learned lengthscales to map the output of the CNN back to continuous function. Pseudo-code for forward pass through the ConyCNP (7) can be found in Appendix B.
Conv CNP (T) The ConvCNP (T) closely follows the ConvCNP (I) model, with the main difference being in the input to the model. For the approximately equivariant model, we sum up the output of the resized functional embedding with the representation of the fixed inputs. The latter is obtained by passing the input locations (that lie on the grid) through and MLP with two hidden layers of dimension C. We consider C' fixed inputs, we zero them out outside of [-7, 7] and during training we drop them out with a probability of 0.1. Pseudo-code for a forward pass through the ConvCNP (T) can be found in Appendix B.1.
RelaxedConvCNP (T) We use identical architecture the ConvCNP(T), with regular convo- lutional operations replaced by relaxed convolutions. We use L = 1 kernels such that the total parameter count remains similar (see Appendix B.5)
EquivCNP (E) We use an identical architecture to the ConvCNP (T), with symmetric convolutions in place of regular convolutions. Pseudo-code for forward pass through the EquivCNP (E) can be found in Appendix B.2.
EquivCNP (E) We use an identical architecture to the ConvCNP (7), with symmetric convolutions in place of regular convolutions. Pseudo-code for forward pass through the EquivCNP (E) can be found in Appendix B.2.
Training Details and Compute For all models, we optimise the model parameters using AdamW [Loshchilov and Hutter, 2017] with learning rate of 5 x 10-4 and batch size of 16. Gradient value magnitudes are clipped at 0.5. We train for a maximum of 500 epochs, with each epoch consisting of 16,000 datasets (10,000 iterations per epoch). We evaluate the performance of each model on 80,000 test datasets. We train and evaluate all models on a single 11 GB NVIDIA GeForce RTX 2080 Ti GPU.
Additional Results In Figure we compared the predictive distributions of the eight considered models for test dataset where the context range spanned both the low and high-lengthscale regions. In Figure 3 and Figure 4 we provide additional examples where the context range only spans one region (the low-lengthscale one in Figure 3 and high-lengthscale one in Figure 4).
Figure 3 shows that the non-equivariant model (TNP) does not produce well-calibrated uncertainties far away from the context region. The equivariant models underestimate the uncertainty near the change point location x giving rise to overly-confident predictions at the transition between the low and high-lengthscale regions. For the TNP (T)) the uncertainties remain low throughout the entire high- lengthscale region. In contrast, the approximately equivariant models manage more accurately capture the uncertainties beyond the transition point, into the high-lengthscale region. This indicates that the approximately equivariant models are better suited cope with non- -tatiiniiossis in the data.
When the context region only spans the high-lengthseale region, both the strictly and approximately equivariant models output predictions that closely follow the ground truth. However, the non- equivariant model completely fails to generalise. outputting almost symmetric predictions about the origin (x = 0). We hypothesise this is because of its inability to generalise, resulting from the lack of suitable inductive biases.
22

---

Context
Model
TE Model
Ground truth
(a) TNP.
(b) ConyCNP (7).
(c) EquivCNP (E).
(d) TNP
(e) ConvCNP (?)) (f) RelaxeCConvCNP (7))
(g) EquivCNP (E).
(h) TNP (
Figure 3: comparison between the predictive distributions on a single synthetic 1D regression dataset of the TNP- ConvCNP-, and EquivCNP-based models with different inductive biases (non- cquivariant, equivariant, or approximately equivariant). Unlike Figure the context range only spans the low-lengthscale region. For the approximately equivariant models, plot both the model prediction (blue), as well as the predictions obtained without using the fixed inputs, which results in strictly equivariant model (red). The approximately equivariant models the only ones able to correctly capture the uncertainties around the lengthscale change point (z = 0).
Context
Model
TE Model ---
Ground truth
TNP.
(b) ConyCNP (7).
(c) ENU (E).
( TNP (7).
(e) ConvCNP (T).
(f) RelaxedConvCNF (T). (g) EquivCNP (E))
(h) TNP (T).
Figure 4: A comparison between the predictive distributions on a single synthetic 1D regression dataset of the TNP-, ConyCNP-, and EquivCNP-basee models with different inductive biases (non- equivariant, equivariant, or approximately equivariant). The context range only spans the high- lengthscale region. For the approximately equivariant models, plot both the model prediction (blue), as well as the predictions obtained without using the fixed inputs, which results in a strictly cquivariant model (red). Both the strictly and approximately equivariant models output predictions that closely resemble the ground truth. but the non -ou-eauiia TNP model completely fails to generalise.
23

---

Table 3: Average log- likelihoods (1) for the synthetic I-D GP experiment when tested on context
sets. Ground truth log- leeeliilihood is 0. 0.2806 = 0.0005.
Model
Context log-likelihood (1)
TNP
0.2296 0.0007
TNP (2)
0.2396 0.0013
TNP (7)
0.2344 0.0020
ConvCNP (7)
0.2362 0.0007
ConvCNP (7)
0.2218 + 0.0009
RelaxedConvCNP (7)
0.2381 + 0.0008
EquivCNP (E)
0.2213 + 0.0007
EquivCNP (E)
0.1992 + 0.0010
Log-likelihood context data We assess the model's ability to reconstruct the context set by setting the target set equal to the context set and computing the log-likelihood. Note that, unlike the log-likelihood values from Table 1, where the target range expands beyond the context range, now the model is just tested within the context range, leading to higher overall values.
The results in Table show that the model is able to accurately model the context sets, with log- likelihood values close to the ground truth log-likelihood of the data of 0)28066 + 0.0005.
Analysis of equivariance deviation We next analyse the approximately equivariant behaviour of our models. More specifically, we show that while breaking equivariance to a certain degree, the models still share similarities with the fully-cquivariant models-thus being characterised as approximately equivariant. This is already visually depicted in Figure 1, where the predictive means of the approximately equivariant models are similar to the predictive means of the corresponding equivariant models up to some 'small perturbation'.
We also show this quantitatively, by introducing the equivariance deviation
Mequiv
Mapprox-eape equiv
Dequiv
E
Mequiv
where Mupprox-equiv represents the predictive mean of the approximately equivariant model, and Pequiv the predictive mean of the same model with zeroed-out basis functions (i.e. the equivariant component of the approximately equivariant model). this case, we consider the L₃ norm (i.e. p = 1).
Table 4: Equivariance deviation (Asquiv) the approximately equivariant models in the 1-D synthetic GP experiment.
Model
Asquv
TNP (T)
0.0896 + 0.0011
ConvCNP (T)
0.0823 + 0.0005
RelaxedConvCNP (T)
0.1460 + 0.0006
EquivCNP (15)
0.0825 0.0006
The equivariance deviations from Table are around 8 9% (with the exception of the Relaxed- ConvCNP (I) which shows a 15% deviation). This indicates that the approximately equivariant models' predictions deviate only slightly from the equivariant predictions, thus being equivariant in the approximate sense. Moreover, as pictured in Figure 1, the deviation from strict equivariance is not random, and allows the models to learn symmetry- breaking features in data-driven manner (in this case, the lengthscale changepoint).
Analysis of number of fixed inputs In this ablation, analyse the effect of the number of fixed inputs B. As mentioned in the main text, the number of fixed inputs influences the number of degrees of freedom in which the decoder deviates from G-equivariance. Thus, 0 fixed inputs leads to

---

G-equivariant model, while in the limit of infinitely many fixed inputs, the decoder can become fully non con-equivariant
In the main experiment we used 64 fixed inputs for the ConvCNP (T) and 4 Fourier coefficients for TNP (T). Table shows the results when varying the number of fixed inputs B for the two models. We observe that the biggest gain in performance is achieved by going from B = 0 to B = 1 basis functions. In this dataset, the approximate equivariance is achieved by modifying one parameter (i.e. lengthscale) of the data generation process from one side of the changepoint location to the other. Indeed, for the TNP (T) we observe that the performance plateaus for B 2 1. suggesting that one Fourier coefficient is enough to capture the symmetry-breant
feature of the dataset. For the
ConvCNP (I') the performance increases up to B 2 4, and for B 1 the test log-likelihoods are within 3 standard deviations for all variants. Thus, the empirical findings are in agreement with the data generation process, indicating that deviation in one a couple of degrees of freedom suffices to capture the departure from strict equivariance. What is more, we see that using more fixed inputs than necessary does not hurt the performance of the model (when considering the standard deviations). This motivates in practice the use a 'sufficiently high number' of fixed inputs, to make sure the model is given enough flexibility to correctly capture the departure from strict equivariance.
Table 5: Average test log- likelihoodds (f) of the TNP (T) and the ConvCNP (T) models when varying the number of fixed inputs. All standard deviations are 0.004.
No. fixed inputs
Model
1
8
16
TNP (T)
0..888
0..499
-0.410
-0.406
-0.408
0..03
CNnVCNP (T)
-0.499
-0.451
-0.439
-0.430
-0.433
-0.432
E.2 Smoke Plumes
The smoke plume dataset consists of 128 x 128 2-D smoke simulations for different initial conditions generated through PhiFlow[Holl et al., 2020]. Hot smoke is emitted from circular region at the bottom, and the simulations output the resulting air flow in a closed box. We also introduce fixed obstacle at the top of the box. To obtain a variety of initial conditions we sample the radius of the smoke source uniformly according to r ~ 2|, 30). Moreover, we randomly choose its position among three possible x-axis locations: (30, 64, 110} (but we keep the position fixed at 5). Finally, we sample the buoyancy coefficient the medium according to B ~ ((00 0.5). The closed box, the fixed spherical obstacle, and the position of the smoke inflow (sampled out of three possible locations) break the symmetry in this dynamical system.
For each initial condition, we run the simulation for 35 time-steps with a time discretisation of At 0.5, and only keep the last state as one datapoint. We show in Figure 5 examples of such states. In total, we generate samples for 25,000 initial condition, and we use 20,000 for training, 2,500 for validation, and the remaining for test.
The inputs consist of [32, 32) regions sub-sampled from the [128, 128) grid. Each dataset consist of a maximum of N 1024 datapoints, from which the number of context points are sampled according to N₂ ~ {10, 250}. with the remaining points set as the target points.
We use an embedding / token size of D2 128 for the PT-TNP-basse models and D. 16 for the ConyCNP-basee models. The decoder consists of an MLP with two hidden layers of dimension Dz- The decoder parameterises the mean and pre-softplus standard deviation of Gaussian likelihood with heterogeneous noise. Model specific architectures are as follows: PT-TNP For the PT-TNP models we use the same architecture dimensions as the TNP deseribed in Appendix E. We use the IST-style implementation of the PT-TNP [Ashman et al., 2024], with initial pseudo-token values sampled from a standard normal distribution. We use 128 pseudo-tokens.
PT-TNP (T) The PT-TNP (T) models adopt the same architecture choices as the TNP (T) described in Appendix E.I. The initial pseudo-tokens and pseudo-input-location are sampled from a standard
25

---

Figure 5: Examples of smoke simulations from the smoke plume dataset for six different combina- tions of smoke radius r and buoyancy B. For each such combination, we show the resulting state for all of the three possible x-axis locations. The inputs to our models are randomly sampled 32 x 32 patches (indicated in red) from the 128 128 states.
normal. We use 128 pseudo-tokens. Pseudo-code for forward pass through the PT-TNP (7) can be found in Appendix B.4.
PT-TNP (T) The architecture of the PT-TNP (T) is similar to that of the PT-TNP (T), with the exception the extra fixed inputs that are added to the pseudo-token representation. These are obtained by passing the pseudo-toker locations through an MLP with two hidden layers of dimension D2. We use D2 fixed inputs and apply dropout to them with a probability of 0.5. Pseudo-code for forward pass through the PT-TNP (T) can be found in Appendix B.4.
U-Net [Ronneberger et al., 2015] architecture for the CNN with layers. We use Cin = 16 input channels and Caut 16 output channels, with the number of channels doubling / halving apply pooling with size two, and between each upwards layer we linearly up-sample to recover the size. We
ConvCNP (T) For the ConvCNP model,
the way down / up. Between each downwards layer
kernel size of 9 with stride of one. We use the natural discretisation of the 128 x 128 grid.
ConyCNP (?) We use the same architecture as the ConvCNP (7)) The fixed inputs are obtained by passing the discretised grid locations through MLP with two hidden layers of dimension C... We use Cin fixed inputs and apply dropout with probability 0.5.
RelaxedConv (T) We use an identical architecture to the ConvCNP (T), with regular con- volutional operations replaced by relaxed convolutions. We use = 1 kernels such that the total parameter count remains similar. We obtain the additional fixed inputs by passing the effective discretised grid each layer (after applying the same pooling up-sampling operations the U-Net) through an MLP with two hidden layers of dimension Cin-
stecrable E-cquivariant CNN architecture consisting of nine layers, each with = 16 input output channels. We use a kernel size of = and stride of one. We discretise the continuous rotational symmetries to integer multiples of 2n/8. We use the natural discretisation of the 128 x 128 grid.
Equiv CNP (E) For the EquivCNP (E) model,
EquivCNP (E) We use the same architecture as the EquivCNP (E) with the same fixed input architecture as the ConvCNP (T).
Training Details and Compute For all models, we optimise the model parameters using AdamW [Loshchilov and Hutter, 2017] with learning rate of 5 x 10-1 For the ConvCNP 7 and T, EquivCNP T' and T. and non-equivariant PT-TNP models we use batch size of 16, while for the PT-TNP T and 7 models we use batch size of 8. Gradient value magnitudes are clipped at 0.5. We train for maximum of 500 epochs, with each epoch consisting of 16,000 datasets for a batch size of 16, and
26

---

8,000 datasets for batch size of 8 (10,000 iterations per epoch). We evaluate the performance of each model on 80,000 test datasets. We train and evaluate all models on a single 11 GB NVIDIA GeForce RTX 2080 Ti GPU.
We show in Figure 6 comparison between the predictive means, well as the absolute difference between them and the ground-truth (GT), for all the models in Table 1. For PT-TNP and ConvCNP, the predictions of the equivariant models are more blurry, whereas the approximately equivariant models better capture the detail surrounding the obstacle or the flow boundary, For the EquivCNP we did not observe significant difference between the equivariant and approximately equivariant model.
Additional Results
GT.
PT-TNP (Ø, T, and 7).
ConvCNP (T and T).
EquivCNP (E and E).
Context.
Figure 6: A comparison between the predictive distributions of the equivariant and approximately cquivariant versions of the three classes of models: PT-TNP, ConvCNP, and EquivCNP. From left to right we show: the ground-truth (GT), the non-equivariant PT-TNP, PT-TNP (T). PT-TNP (T). ConvCNP (I)) ConvCNP (I"), GelaxedConnCCN (?), EquivCNP (E), and EquivCNP (É). The top row shows the mean the predictions, while the bottom row shows the absolute difference between the predicted mean of each model and the ground-truth.
E.3 Environmental Data
The environmental dataset consists of surface air temperatures derived from the fifth generation the European Centre for Medium-Range Weather Forecasts (ECMWF) atmospheric reanalyses (ERAS) [Copernicus Climate Change Service, 2020]. The data has a latitudinal and longitudinal resolution of 0.5%, and temporal resolution of an hour. We consider data collected in 2018 and 2019 from regions in Europe (latitude longitude range of [350, 60%] / [10", 45%)) and in the US (latitude longitude range of [-200-; -80) / [30%, 50"). In both the 2-D and 4-D experiment, we train on Europe's 's 2018 data and test on both Europe's and the US's 2019 data.
For the 2-D experiment, the inputs consist of latitude and longitude values. Individual datasets are obtained by sub-sampling the larger regions, with each dataset consisting of a [32, 32] grid spanning 16" across each axis. For the 4-D experiment, the inputs consist of latitude and longitude values, as well as time and surface elevation. Each dataset consists of [4, 16, 16) grid spanning 8 across each axis and 4 days. In both experiments, each dataset consists of a maximum of N 1024 datapoints, from which the number of context points are sampled according N. U1 100 }}. with the remaining set as target points.
For all models, D2. The decoder parameterises the mean and pre-softplus variance of a Gaussian likelihood with heterogeneous noise. Model specific architectures follows:
decoder consisting of MLP with two hidden layers of dimension
PT-TNP Same as Appendix E.2.
PT-TNP (T') Same as Appendix E.2.
PT-TNP (7) Same as Appendix E.2 with 128 pseudo-tokens and fixed inputs zeroed outside 2018 and the latitude longitude range of [35", 60%) [10", 45%].
ConyCNP (I) Same as Appendix E.2.
ConvCNP (7) Same as Appendix E.2 with fixed inputs zeroed outside 2018 and the latitude / longitude range of [35", 60") [10", 45%].
27

---

RelaxedConyCNF (T') Same as Appendix E.2 with fixed inputs zeroed outside 2018 and the latitude longitude range of [35", 60% [10P, 45%].
EquivCNP (F.) Same as Appendix E.2.
EquivCNP (E) Same as Appendix E.2 with fixed inputs zeroed outside 2018 and the latitude/ / longitude range of [35", 60% [10", 45%].
Training Details and Compute For models, we optimise the model parameters using Adam [Loshchilov and Hutter, 2017] with learning rate of 5 x 10-2 and batch size of 16. Gradient value magnitudes are clipped at 0.5. We train for a maximum of 500 epochs, with each epoch consisting of 16,000 datasets (10,000 iterations per epoch). We evaluate the performance of each model on 16,000 test datasets. We train and evaluate all models on a single 11 GB NVIDIA GeForce RTX 2080 Ti GPU.
Additional Results To demonstrate the importance of dropping out the fixed inputs during training with finite probability, we compare the performance of two ConvCNP (T) models on the 2-D experiment in Table 6: one with a dropout probability of 0.0, and the other with 0.5. Due to limited time, we were only able to train each model for 300 epochs (rather than for 500 epochs, hence the difference in results for the ConvCNP (7) model with dropout probability of = 0.5 to those shown in Table 2). Nonetheless, we observe that having a finite dropout probability is important for the model to be able to generalise OOD (i.e. the US).
Table 6: Average test log-likelihoode (f) for the 2-D environmental regression experiment, p denotes the probability of dropping out the fixed inputs during training.
Model
Europe (1)
US (1)
ConvCNP (ī, p = 0.0)
1.19 0.01
0.51 = 0.02
ConvCNP (T.p= 0.5)
1.16 + 0.01
0.16 = 0.02
Log-likelihood of context data We perform the same assessment as in Appendix E. to check whether the model is able accurately reconstruct the context data.
The results in Table 7 show that the models have good performance when tested on the context sets, achieving better log-likelisoods than in Table 2, where the models tested on target sets.
Analysis of equivariance deviation We repeat the equivariance deviation analysis from Appendix E. on the 2-D environmental regression dataset. In this that we only compute this on Europe, since the predictions on US are obtained by zeroing-out the basis function, leading
the Le instead of the L₁ norm. Note
equivariance deivation of O. The results are shown in Table 8.
All equivariance deivations are between 2 4%, indicating that the models only slightly deviate from the equivariant prediction. However, as shown in Figure 2, this deviation allows them to capture the
Table 7: Average test log-likelihoods (f) for the 2-D environmental regression experiment when tested on context sets.
2-D Regression
Model
Europe (1)
US US)
PT-TNP
1.74 0.01
PT-TNP (T)
1.66 + 0.01
1.28 + 0.01
PT-TNP (7)
1.76 0.01
1.47 0.01
ConvCNP (T)
1.20 0.02
0.34 0.02
ConvCNP (7)
1.50
0.97 0.01
RelaxedConvCNP (?)
1.29 - 0.01
0.86 0.01
EquivCNP (E)
2.03 1 0.01
1.76 0.01
EquivCNP (F.
2.05 0.01
1.69
28

---

Table 8: Equivariance deviation (Aequi) of the approximately equivariant models in the 2-D environ- mental regression experiment.
Model
Augiv
PT-TNP (7)
0.0406 1 0.0005
ConCCNP (
0.00237 + 0.0004
RelaxedConvCNP
0.0239 0.0004
EquivCNP (E)
0.0242 0.0005
symmetry-breaking components of the weather dataset, with one important such component being the topography.
Analysis number of fixed inputs Finally, we investigate the influence of the number of fixed inputs B on the performance of the model. We vary the number of fixed inputs used in the ConvCNP (7) from the 2-D regression experiment. The fixed inputs are linearly transformed 16 inputs, which are then summed together with the input into the CNN.
The results are shown in Table 9, where we see that, similarly to the 1-D GP experiment, the largest gain in performance obtained by going from to a single fixed input. Using more than fixed input provides diminishing gains, but, importantly, never hurts performance. The saturation in performance for B 2 1 is expected in this dataset, given that the information that is primarily missing is the topography. However, in practice, given that the performance of the model does not decrease with increasing B (i.e. the model learns to ignore unnecessary fixed inputs), one can choose a 'sufficiently high' value for B.
Table 9: Average test log-likel lieliiieihoodds (f) of the ConvCNP (7) with different number of additional fixed inputs for the 2-D environmental regression experiment. We show the results when the models tested on Europe.
No. fixed inputs
Region
0
2
4
8
16
Europe
1.11 + 0.01 1.19 + 0.01
1.19 + 0.01
1.20 + 0.01
1.20 + 0.01
1.20 + 0.01
29

---

NeurIPS Paper Checklist
The checklist is designed encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the
relevant information is Not Available.
Please provide short (1-2 sentence) justificatin right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]" ", is perfectly acceptable to answer "[No] "provided proper justification is given (c.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
general, answering
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS paper checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only the provided macros for your answers.
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: All novel theorems include either the full proof, or a reference the full proof in the Appendix. We show supporting empirical evidence to justify our claims and contributions in Section
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
30

---

Justification:
Limitations are discussed in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approxinations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low images are taken in low lighting. Or speech-to-text system might be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: All novel theorems include either the full proof, or reference to the full proof in the Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
We provide full experimental details in the Appendix.
Guidelines:
31

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide open access to the codebase, containing instructions to reproduce the data and replicate the main experimental results. The environmental data is public.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
muuides/CodeSubmi
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https: hwmbbllc/guiiesss httpc///ggiiee///eeebmmmmn (Co.eeumiissiiiiecoooc ssionPolicy) for more details.
/
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
32

---

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We provide full experimental details in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification:
We report standard errors for all results across the entire test set.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide details of compute for each experiment in the Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
33

---

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /tterr//neuussspee httppp///aarrss...pppiii///////e cecpiiicccullllieee ines?
Answer: [Yes]
Justification:
The paper conforms to the Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations in their jurisdiction).
Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper conducts foundational research and is not tied to particular applica- tions.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA No, they should explain why their work has societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out. For example, it is legitimate point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
34

---

Justification:
This paper does not pose such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification:
We provide complete citations for datasets used in this paper.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, saperswithcone com/ datasets has curated licenses for some datascts. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [Yes]
Justification:
We released documented code.
Guidelines:
The answer NA means that the paper does not release assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
35

---

Answer:
[NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
36


Reviews:
[{"review_id": "gtgwwidzm3", "reviewer": "Reviewer_7yUH", "summary": "The work provides an alternative approach to obtain a loose equivariance constraint. The authors established an interesting relationship between equivariant and non-equivariant models, showing that equivariant models with enough fixed input can approximate any non-equivariant model. Subsequently, the authors suggest using additional fixed input in neural processes, which is a common practice in transformers and NeRFs. The proposed method achieved better performance on multiple benchmark datasets.", "questions": "Line 103: how is the action of the group on the dataset $D$ defined?\n\nEquation 2, line 111: How is the G-invariance of the kernel defined?\n\nLine 143: $E_n$ is defined to take $n$ additional input; however, line 146 is defined to map on $H^{1+2n}$, which I think means $2n$ additional input. An explanation would be appreciated.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The work provided a nice theoretical result that offers an interesting reinterpretation of existing practices.\n\n2. The proposed technique is empirically evaluated on a wide range of datasets.", "weaknesses": "The work provides a nicer reinterpretation of the existing practice of adding fixed/learnable positional embedding. Despite this reinterpretation, the work does not quantify the degree to which these additional fixed inputs hurt the equivariance property of the equivariant models.\n\nNor does it provide a scheme to automatically learn the number of fixed inputs for controlling the equivariance restriction. \n\nThus, the work indeed proposes an approximately equivariant model, however, with unknown approximation.\n\nThis severely limits the contribution of the work.\n\nLine 91: different notation should be used to denote the group action on $X$ and $Y$ \n\nLine 102: the symbol $f$ is already used in line 87 to denote the group element. The use of different symbols will fascinate reading.\n\nLine 115: the notation of the encoder matches with a notation of the identity element of the group.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for acknowledging our theoretical contributions, as well as the empirical evaluation on a variety of datasets. We hope that our work justifies the advantages of relaxing model symmetries and proves that the proposed framework is an effective approach for achieving this. We address your concerns below.\n\n**Weaknesses**\n\n**Quantifying degree of non-equivariance** We believe that our experiments carefully analyse the effects of including additional inputs in equivariant architectures through both quantitative results and visualisations (i.e. we run both the versions with and without the additional fixed inputs, so the results quantify the effect of including additional fixed inputs). We provide results for both the synthetic GP and environmental data experiments, showing how our models break equivariance and how they compare to their strictly equivariant counterparts.\n\nNevertheless, we acknowledge that we do not characterise precisely how and to which degree additional inputs break equivariance. (Throughout the exposition, however, we do repeat the quite rudimentary intuition that more additional inputs should break equivariance more, e.g. lines 163-171). In the conclusion, we mention this limitation, which paves the way for interesting future work.\n\n**Automatically learn the number of fixed inputs** - We respectfully disagree that not having a way to automatically learn the number of fixed inputs is a severe limitation of the work. In practice, the number of additional inputs can very reasonably be tuned as a hyperparameter alongside all other architectural choices, like the number of layers and widths of MLPs.\n\nIn our experiments, setting the number of fixed inputs to a “reasonably high number” tends to work well. We suspect that the dropout procedure (“Recovering equivariance out-of-distribution” on line 224) regularises any unnecessary additional inputs. The additional ablation provided in the common response supports this hypothesis.\n\n**L91** - We acknowledge that using different symbols for the group actions on $\\mathcal{X}$ and $\\mathcal{Y}$ would further disambiguate the notation. However, the notation in our submission is already heavy, as pointed out by other reviewers, and we remark that overloading $g$ in this way is both accurate (by defining $\\mathcal{X}$ and $\\mathcal{Y}$ both as $G$-spaces) and standard. For example, see the section “Formalization” of the entry “Equivariant map” on Wikipedia.\n\n**L102** - Many thanks for pointing this out—we will modify the notation such that the symbol $g’$ is used to denote an alternative group element.\n\n**L115** - Thanks for pointing this out, we will modify the notation.\n\n**Questions**\n\n**L103** - We apologise for the confusion, as there is a typo on line 102, where we explain how the action of the group on the dataset $\\mathcal{D}$ is defined. $g\\mathcal{D} \\in S$ consists of input-output pairs ($g\\mathbf{x}_n, \\mathbf{y}_n$), rather than  pairs ($g\\mathbf{x}_n, \\mathbf{x}_y$). \n\n**L111** - G-invariance of the kernel $\\psi$ is defined by the property that $\\psi(\\mathbf{x}_n, \\mathbf{x}_m) = \\psi(g\\mathbf{x}_n, g\\mathbf{x}_m)$ for all $g \\in G$. We will add this definition to the writing.\n\n**L143/L146** - We do apologise for the confusion here: there are indeed $2n$ additional inputs, owing to $n$ $t$’s and $n$ $e$’s. We thank the reviewer for spotting the mistake, and will correct it in the revised version.\n\nIn light of the above, we would be grateful if you would reconsider your score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LkNeUaL96c", "reviewer": "Reviewer_Ks9f", "summary": "This work considers approximately equivariant models --- which may better model or learn real-world tasks than exactly equivariant models --- especially in the context of neural processes. A new approximately equivariant method is developed, which uses an exactly equivariant model along with fixed inputs that break equivariance. This method can be used to parameterize approximately equivariant neural processes, by modifying existing approaches in simple ways. Experiments on 1-D regression, smoke plumes, and environmental data show benefits of the approach.", "questions": "1. In equation (3), I think think the number of inputs should be larger than $n$, right?\n2. Perhaps I missed this, but do you initialize the models to be equivariant? Via setting $t$ to be the zero function?\n3. It would be worth noting connections of your approximate equivariance approach to some related prior work. You mention the connection to positional embeddings in Transformers on Page 5, but this also applies to other domains such as Vision Transformers [Dosovitskiy et al. 2021]; this is more specifically investigated in the context of approximate equivariance by Lim et al. 2023. Also, the approach of taking additional inputs is related to the symmetry breaking approach of Xie et al. 2024 and the probabilistically / approximately equivariant approach of Kim et al. 2023.\n\nReferences:\n* Wang et al. 2022a. Approximately Equivariant Networks for Imperfectly Symmetric Dynamics\n* Dosovitskiy et al. 2021. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale\n* Lim et al. 2023. Positional Encodings as Group Representations: A Unified Framework\n* Xie et al. 2024. Equivariant Symmetry Breaking Sets\n* Kim et al. 2023. Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance", "limitations": "Some discussion in conclusion", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Nice flexible framework for approximate equivariance, which unifies existing methods. Similar ideas could be useful more broadly in learning with (approximate) symmetries.\n2. The approach has nice ways of controlling the degree of equivariance, by number of inputs (\"degrees of freedom\" in Section 3.1), or empirically doing things like regularizing the effect of the additional inputs towards zero. The method of setting $t_b(x) = 0$ during test time is also nice.\n3. Good empirical gains on several datasets, by making simple tweaks to several NP methods.", "weaknesses": "1. Notation and definitions are heavy, which is somewhat understandable given the subject matter, but I do think it can be improved. For instance, compact operators are not defined.\n2. From what I can tell, the smoke plumes experiment in Section 5.2 follows a similar setup to that of Wang et al. [2022a], but this is not mentioned. Also it would be good to note how you chose e.g. the parameters of the simulations and the choice to include an obstacle, or whether this was mostly arbitrary (not a big issue).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for acknowledging our theoretical and empirical contributions in the field of approximate equivariance, with a focus on applications in neural processes. In particular, we are pleased that you appreciate the importance of setting the additional inputs to zero at test time outside the training domain, which we believe is one of the key components for the model’s ability to generalise. We address your concerns below.\n\n**Weaknesses**\n\n**Heavy notation** - We acknowledge that the paper might be notation and definition heavy at times given the nature of the subject we tackle, but we put in effort towards including as many intuitive explanations as possible—for example, the discussion after Theorem 2 (lines 154 - 171).\n\nIf you have further suggestions to improve notational clarity, we would be open to them.\n\n**Definition of compactness** -  We apologise for not including a definition of compactness and we will update the text accordingly.\n\n**Smoke plume setup** - You are right, the setup does hold similarities to that of Wang et al. [2022a], but there are also a few key differences, especially in the way we break the strict equivariance. These are all outlined in Appendix D.2, but here we provide a comparison with the setup from Wang et al. [2022a].\n\n**Similarities**\n\n- We also use the PhiFlow library in order to generate the smoke simulations.\n\n- We use the same solver techniques (i.e., the code for the solver is inspired from the repository associated with Wang et al. [2022a]), with the same parameters for the time discretisation and choice of optimiser.\n\n**Differences**\n\n- Wang et al. [2022a] sample the smoke inflow location out of $35$ locations, we randomly sample out of $3$ locations $((30, 5), (64, 5), (110, 5)$ on a ($128, 128$) domain).\n\n- We randomly sample the buoyancy force at each simulation as $B \\sim \\mathcal{U}[0.1, 0.5]$, and keep it constant throughout the entire sample domain. In contrast, Wang et al. [2022a] have two subdomains, with a different buoyancy force for each subdomain, and they keep those values constant for all their simulations.\n\n- Besides the boundary of the simulation (which are the same in our and Wang et al [2022a]’s setup), we also consider an obstacle with its centre at $(64, 100)$ and a radius of $20$.\n\n- We train and evaluate the model on randomly sampled $32 \\times 32$ patches of the $128 \\times 128$ domain corresponding to the final simulation state, as opposed to the entire state ($128 \\times 128$). \n\nThus, although we took inspiration from Wang et al. [2022a] setup, there are some key differences—we break the symmetry of the system through the closed boundary (as in Wang et. al [2022a]), through the fixed spherical obstacle through which smoke cannot pass, and by sampling the inflow position. The latter is also performed by Wang et. al [2022a], but they consider more possible inflow positions. In the limit of randomly sampling the inflow position, this no longer breaks the symmetry of the system, so we decided to only use 3 positions in order to depart from strict symmetry more.\n\n**Questions**\n\n**Number of inputs** - Many thanks for pointing out this mistake—the number of inputs should be $2n + 1$, which we will correct in the revised version of the paper.\n\n**Model initialisation** - We do not initialise the models to be equivariant, but there are two techniques that we use in order to recover equivariance out-of-distribution, where the training data does not contain enough information in order for the model to correctly learn how to depart from strict equivariance. More specifically, we:\n- During training, randomly set the additional inputs to $0$ (corresponding to strict equivariance) with a fixed probability (usually $0.5$), such that the model learns a good strict equivariant representation of the underlying system when the additional inputs are zeroed out.\n\n- During testing, we always set the additional inputs outside the training domain to $0$, such that the model reverts to the predictions of the equivariant component in the regions where it has not seen any data during training (where strict equivariance is the most reasonable inductive bias we can assume).\n\nSee the paragraph “Recovering equivariance out-of-distribution” at line 224.\n\n**Connections to related work** - You are correct that this work has connections to, and offers insights for, many other important methods used within ML such as positional encodings in ViT. Many thanks for pointing out Lim et al. 2023, Xie et al. 2024 and Kim et al. 2023. These three works are indeed related to our work, and we will include them as citations in the related work section.\n\nThanks again for your feedback—we hope that we have adequately addressed your concerns. If so, we would be grateful if you could consider increasing your score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CZKuTZOtRf", "reviewer": "Reviewer_WhxG", "summary": "The paper describes a new framework for soft/approximate equivariance, based on the functional analysis of compact on hilbert spaces, assuming unitary group actions.\n\nThis is then applied to equivariant neural processes.", "questions": "0) see the un-numbered remark at the beginning of \"weaknesses\" section, about relevance to \"neural processes\" to the presented methods.\n\n1) line 49-50 can the authors be more specific about what it means that \"any equivariant [...] fixed inputs\"? A reference to the actual result would suffice.\n\n2) line 69-70 \"represented with stochastic processes\" -- is it not \"represented by\"? (this is a genuine question, not a rethorical one)\n\n3) The statement of theorem 1 is not well written. If $\\Phi$ has image in $C(\\mathcal X,\\Theta)$ and if eq. (2) holds, then $\\rho$ has image in the same space, not in $\\mathbb R^{2D}$. Similar incompatibilities also hold for the other domain / codomain spaces for operators $\\rho, e, \\psi$ in that theorem. Also, the discussion in lines 114-116 is incompatible in terms of domain/codomain spaces. This makes the theorem hard or impossible to understand.\n\n4) Still for eq. (2), it is not clear why $\\psi$ has two arguments, is it a kernel operator? the notation from (2) should be expanded and actually explained with a degree of consistency/precision.. pretty please..\n\n5) lines 130-133, one can just require that the action of $G$ is by linear unitary operators, which I think is equivalent to what the authors claim? Also, why do you say \"acting linearly\" at line 130, and then write the additivity explicitly at line 132, as if it were a new requirement and not part of the preceding one?\n\n6) I think that before Proposition 1, or actually even in the introduction itself, one should emphasize the role of compact operators on Hilbert spaces. This is a strong requirement of this theory, and a strong limitation, so it should be highlighted. (I'll add an obvious remark: Don't worry, nobody will consider you worthless if you are fair about the limitations, on the contrary, it speaks highly of you if you do so!)\n\n7) line 156: \"$T\\simeq CNN(\\cdot, t_1,\\dots,t_n)$\" means what exactly? What is the kind of multilinear CNN's that this refers to? Any reference for them? I think that currently this kind of CNN is not well developed, and it is not that trivial as an extension of usual linear CNNs. I think this should be mentioned. If one looks at appendix C, lines 592 and following, that's just the simplest case of such CNNs, so I maybe the authors could spend some time to point this out, and to describe a little bit the difficulties and issues with such multilinear-CNN theory. \n\n8) in the paragraph lines 205-211 I think that it should be stressed (or at the very least mentioned) that $T$  is assumed to be compact, and can't be more general than that.\n9) same mention of \"compact\" should be inserted also in paragraph lines 249-261\n10) same mention of \"compact\" operators should be inserted in seciton 6 (conclusions section)", "limitations": "these are addressed in the \"weaknesses\" section.\n\nI think that the main limitations are in efficiency and scalability of multi-linear CNNs and in the fact that operators of interest are far from compact and not easy to approximate with finite-range ones.", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "The proposed method is very general and conceptually well grounded, and to the extent I could verify they are also new.\n\nThe experiments are well presented and seem promising in terms of accuracy (not sure about efficiency).\n\nEdit after rebuttal: the added experiments are convincing enough in terms of accuracy as well.", "weaknesses": "Edit after rebuttal: I think that the method is a good addition to the literature, and my concerns 1-3 below have been adequately answered in the rebuttal and \"global rebuttal\" parts.\n\n-------------------\n\nThe main enigma about this paper, is: does the technique used really warrant putting \"neural processes\" in the paper's title? It seems like the authors develop a large setup for general approximate equivariance, and the application to neural processes is just one of many applications. \n\nWhat is restricting the applicability of the method, from general equivariant neural networks, to the class of neural processes? This (meaning, the emphasis on neural processes, and the paper's title) is a big distraction in reading the paper, as one tries to find a reason why the setting is restricted like that and one doesn't actually find a convincing answer.\n\n1) The assumption of having tasks with underlying compact operators seems to be swept under the carpet, and not properly discussed. \nIn functional analysis \"compact\" is roughly equivalent to \"approximable by finite-rank, as pointed out at several points in the paper, however the authors don't discuss what limitation this could imply. If this is beyond the scope of this work, still I think that it needs highlighting and *at the very least* it requires pointing out very clearly that this is a future direction to be investigated.\n\n2) when authors work with multilinear CNNs (see question 7 below), this seems not well described, and remains probably too mysterious. Also, I think that their implementation becomes clumsy and inefficient: is that so, or do the authors have a justification and complexity control of why not? This is not required to do in much detail, but maybe just to mention that in order not to hide possible underlying difficulties of the method.\n\n3) the parameter $k_n$ from Theorem 3 is not well behaved, so it is not clear if this theorem is useful in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for acknowledging the novelty of our approach, as well as its theoretical groundness, coupled with encouraging empirical evaluations. We address your concerns below.\n\n**Focus on NPs** - Thank you for highlighting an important point about the general purpose of our approach. The reason why we focussed on the applications in NPs is because we identified a limitation of neural processes that hurts performance in practice. Our goal was to alleviate this limitation and hence developed theory and a method to enable them to model approximate equivariance. We noticed that the theory can be applied more generally and tried to reflect that in the presentation of our paper, but our intended scope is to analyse approximate equivariance in NPs. \n\nOne point we would like to make is that NPs are a rather large class of models (and hence, not that restrictive), as many meta-learning methods can be seen as close relatives/variations of NPs.\n\n**Compactness** - Compactness is a notion conventionally only used for linear operators. For nonlinear operators, which is the main focus of the paper, we focus on the (for linear operators equivalent) notion that the inputs and outputs of $T$ can be approximated with a single finite-dimensional basis (Prop. 2 in App A). To see whether this notion is reasonable to assume, suppose e.g. the case where $\\mathbb{H} = L^2([a, b])$, which would be appropriate for data living on an interval. Considering the Fourier basis, “compactness” here would mean that, uniformly over the input and output functions of $T$, the magnitude of high-frequency basis elements go to zero as the frequency increases. To us, this is a reasonable assumption.\n\nOur intention was definitely not to sweep this issue under the rug. We will add the above argument to the paragraph in lines 188–194 and emphasise its importance appropriately.\n\n**Multilinear CNNs** - The example with the CNN is to provide intuition for how Theorem 2 could work in the nonlinear case (note that in Theorem 2 $E_n$ is nonlinear, so a nonlinear CNN is suitable in approximating linear non-equivariant operators). We do not intend to propose multilinear CNNs as a practical method—that is, both here and in the experiments, we consider nonlinear CNNs, for which there exist a variety of ways in which the additional fixed inputs can be included into the CNN (e.g. as additional input channels). However, note that multilinear CNN are just one-layer CNNs (without nonlinearities) that take in more than one channel. For example, consider a one-layer CNNs that operates on images with a red, blue, and green channel: CNN(red, blue, green). Extending it with even more channels would give the example in the paper: CNN(red, blue, green, $t_1$, $t_2$, …). We will clarify in the writing that this is only an example to provide intuition. In terms of exposition, Theorem 2 is intended as a step-up to Theorem 3; see the first paragraph of Section 3.\n\n**$k_n$ not well behaved** - The utility of Theorem 3 is to motivate the construction of augmenting with additional, fixed inputs, and we show in the experiments that this construction can yield practical benefits. We agree that Theorem 3 cannot be used to characterise exactly how many additional inputs are needed, because the estimate on $k_n$ is poor. As we mention in App A, better proof techniques and/or additional assumptions may enable better estimates on $k_n$. For example, assuming linearity in Theorem 3 would yield that $k_n$ grows linearly in $n$, which is the result of Theorem 2. We acknowledge this as a limitation of our approach in the Conclusion section, and we are excited to pursue this line of work in order to more accurately identify how the behaviour of $k_n$ depends on $n$ for more general nonlinear operators. We suspect, unfortunately, that such estimates require strong additional assumptions on $T$ and might depend on unknowable constants in practice.\n\nNevertheless, please refer to the additional experiment that analyses the performance of the models as we increase the number of additional fixed inputs. As shown in the environmental data case (which is a real-life dataset), in this case $k_n$ is not ill-behaved.\n\n**L49-50** - This refers to Theorem 2 and 3, which we will make clear in the revised version of the paper.\n\n**L69-70** - We are happy to rephrase this to “represented by”, to avoid possible confusion as to whether or not predictions are “represented alongside” stochastic processes.\n\n**Theorem 1** - Many thanks for pointing out this mistake. The first $\\rho$ should be $\\phi: \\mathcal{Y} \\rightarrow \\mathbb{R}^{2D}$, which we will correct in the revised version. Regarding the incompatibilities in the following paragraph, Theorem 1 can be obtained by choosing $\\mathcal{Z} = \\mathbb{R}^{2D}$. However, in practice we do not enforce this restriction, typically choosing $\\mathcal{Z} = \\mathbb{R}^E$ where $E >> 2D$.\n\n**$\\psi$ in Eq. 2** - Yes, as described on line 111 $\\psi$ is a G-invariant positive-definite kernel. We appreciate that the notation used in this theorem is very dense, and we will rephrase to clarify.\n\n**L130-133** - Our intention was to be extra clear. We will slightly reword to avoid the suggestion that it is a new requirement.\n\n**Compact operators** - We did not highlight compactness because it is conventionally a notion for linear operators, whereas our main result is for nonlinear operators (Theorem 3). Nevertheless, instead of compactness, you’re completely right that Theorem 3 requires other technical conditions. We will add a clear mention to these technical conditions in the introduction and discussion.\n\n**L156** - Please see the response above about multilinear CNNs above.\n\n**L205-211, L249-261** - We will clarify in both places. Thank you.\n\nWe are very thankful for your detailed feedback, and for the questions / concerns you have raised—we have no doubt that these have made for a stronger paper. We hope that we have adequately addressed your concerns.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NjfDlDDhgb", "reviewer": "Reviewer_dBf7", "summary": "This paper implemented the approximately equivariant neural process (NP) by relaxing the equivariance of the NP decoder and the relaxation is conducted by adding several leanable parameters as fixed inputs fed with the data embeddings into the decoder.", "questions": "1. Encoder should be equivariant? If so, doesn’t it lose non-equivariant information from the dataset D?\n1. How to add the fixed inputs into the dataset embedding in ConvCNP? If we just sum all, the number of the fixed inputs is obviously meaningless.\n1. When the decoder is ViT, how many model parameters are additionally necessary compared to the strictly equivariant case?\n1. It seems that the suggested relaxation is applicable to not only NP. Why did the paper choose NP to prove the relaxation idea?", "limitations": "There is no limitation of the contribution itself but experimental limitations are mentioned in the weakness.", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "The proposed method is applicable regardless of architecture and easy to be implemented because it is simply adding some parameters as inputs.", "weaknesses": "1. No supplemental material for a source code.\n1. A concept figure will be helpful to understand the intuition of the paper.\n1. Absent analysis for equivariance error of the trained models. Did the trained model really reflect the approximate equivariance inherited in the dataset?\n1. The paper reported only the log-likelihood of (maybe) the target data, which may be overfitted by the target data and underfitted by the context data.  It is necessary to report the log-likelihoods with respect to both context and target separately to show the strength of the suggested method accurately.\n1. Analysis for the number of the fixed inputs is necessary in order to prove that the number of the fixed input really controls the magnitude of the approximate equivariance.\n1. All included experiments are only regression tasks. Since NP is essentially a meta-learner, it would be better to add experiments beyond regression to demonstrate its generalization performance across various types of tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and for recognising how easily our method can be applied to achieve approximate equivariance.\n\n**Supplemental code** - As mentioned in the Paper Checklist, we intend to provide open access to the data and code prior to publication. \n\n**Concept figure** - Please see the PDF attached in the common response for a sketch concept figure. If you find this figure helpful, we will include a refined version in the revised paper.\n\n**Equivariance error** - We believe we do provide qualitative analysis of the model’s behaviour with regards to equivariance. For example, in the GP experiment, we provide both the model predictions, as well as the predictions of the model if we set the additional fixed inputs to zero (corresponding to the equivariant component of the model). Fig. 1 e), f), g), and h) show how the additional fixed inputs help the models better capture the transition from the low-lengthscale to the high-lengthscale region (also see lines 290-293).\n\nAnother example is provided in the environmental data experiment. One significant variable influencing the surface air temperature is the elevation. If the elevation map is not provided to the model, we expect the model to depart from strict translation equivariance in a way that correlates with the elevation. This is exactly what we test in the 2-D Spatial Regression task, where we see that the model can infer the effect of local topographical features on air temperature through the non-equivariant component. This can be observed by comparing the difference between the approximately equivariant predictions and the equivariant component of the model (Fig. 2f, i) with the true elevation map (Fig. 2c). Remarkably, the non-equivariant component (i.e. difference between predictions and equivariant component) learns a representation that closely resembles the elevation map, which is probably one of the main factors that leads to departure from strict equivariance in this setup.\n\n**Log-likelihood of target data** - In the usual meta-learning framework for neural processes, the goal is to predict the target set given the context set, so models only parametrise the probability of the target set given the context. For example, see [1] and many follow-up publications [2, 3, 4, 5, 6].\n\n[1] Garnelo, M. et al. (2018). Conditional Neural Processes. \n\n[2] Garnelo, M. et al. (2018). Neural Processes.\n\n[3] Kim, H. et al. (2019). Attentive Neural Processes.\n\n[4] Gordon, J. et al. (2019). Convolutional Conditional Neural Processes. \n\n[5] Bruinsma, W.P. et al. (2023). Autoregressive Conditional Neural Processes.\n\n[6] Nguyen, T., & Grover, A. (2022). Transformer Neural Processes: Uncertainty-Aware Meta Learning Via Sequence Modeling.\n\n**Analysis of number of fixed inputs** - Please refer to our common response.\n\n**Limited to regression** - NPs can indeed be used for other types of tasks beyond regression, like classification. The focus of our submission is on approximate equivariance, and we are primarily interested in applications in regression for studying this property (e.g. fluid dynamics simulation and environmental data modelling). This is also the approach that other related work studying approximate equivariance take, focussing primarily on regression problems [1]. Nevertheless, we acknowledge that applications in classification could also be investigated, similarly to [2, 3], but this is beyond the intended scope of our work.\n\n[1] Wang, R. et al. (2022). Approximately Equivariant Networks for Imperfectly Symmetric Dynamics.\n\n[2] Elsayed, G. F. et al. Revisiting spatial invariance with low-rank local connectivity. ICML 2020.\n\n[3] van der Ouderaa T. et al. Relaxing equivariance constraints with non-stationary continuous filters. NeurIPS 2022.\n\n**Equivariant encoder** - Encoders of equivariant NP architectures are indeed equivariant. Importantly, these encoders do not lose information, because they are usually so-called homeomorphisms, which means that they map every different dataset to a different encoding ([1] , App.A of [2]). Equivariance then means, for example, that the encoding is shifted whenever the dataset is shifted. Equivariance does not break bijectivity (i.e. does not lose information), but says something about how the encoding behaves as the dataset is transformed by a symmetry.\n\n[1] Zaheer M. et. al (2017). Deep Sets. NIPS'17.\n\n[2] Gordon, J. et. al. (2019). Convolutional Conditional Neural Processes. ICLR 2020.\n\n**Adding fixed inputs into the dataset embedding** - We provide implementation details in Appendix D.  When adding the fixed inputs to the dataset embedding for the ConvCNP, we achieve this by adding a different fixed input to each input channel into the CNN. This isn’t the only route through which additional fixed inputs can be included—if we were to treat them as additional channels into the CNN, then the number of additional fixed inputs is not restricted.\n\n**ViT decoder** - In general, the only additional parameters required to modify existing equivariant architectures are those that are required to define the additional fixed inputs. This is typically much smaller than the number of parameters used in the original architecture. For example, for the 2-D environmental data experiment, the regular ConvCNP ($T$) has $15.886 \\times 10^6$ parameters, whereas the ConvCNP ($\\widetilde{T}$) has $15.887\\times 10^6$ parameters (an increase of < 0.01%). The regular TNP ($T$) has $4.006 \\times 10^4$ parameters, whereas the TNP ($\\widetilde{T}$) has $4.823 \\times 10^4$ parameters (an increase of around 20%).\n\n**Focus on NPs** - Please see response to reviewer WhxG regarding **Focus on NPs**.\n\nWe appreciate that our analysis combines results from the neural process literature with more formal mathematical analysis, and consequently might be more dense than usual. Nevertheless, we hope that you find our rebuttal informative, and we would appreciate it if you would consider increasing your score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We would like to thank all reviewers for the time taken to review the paper, for their feedback and useful suggestions, and we hope to address all their concerns through our responses. We are pleased that, in general, our method was seen as a general way of achieving approximate equivariance, while also being “conceptually well grounded”, “new”, and a “flexible framework” that “unifies existing methods”. Moreover, we are content that the reviewers also appreciated the effectiveness of our method, noting the “good empirical gains on several datasets”. 

However, there were some concerns that we aim to address through this general comment, as well as individually with each reviewer. A weakness mentioned by all reviewers is that we do not characterise precisely how many additional input channels are needed and how their number influences the degree of approximate equivariance. We acknowledge this weakness, which is a hard theoretical research question that we aim to investigate in future work. Our suspicion is that no practical answer exists, just like there is no practical way to know how wide an MLP should be for a particular data problem. We, however, do not think that this diminishes the utility of the proposed approach, because the number of additional input channels can be tuned as a hyperparameter. In practice, choosing a “sufficiently high number” tends to just work fine, which is what we did in our empirical evaluations. 

Moreover, we have run an ablation study for the 2D environmental regression experiment where we vary the number of additional fixed inputs into the ConvCNP ($\widetilde{T}$) model. For this particular experiment, the information that is primarily missing is topography. Since topography is only a single variable, we expect that a low number of additional inputs will saturate performance. The predicted log-likelihoods for Europe are as follows:

No. fixed inputs: 0 / 1 / 2 / 4 / 8 / 16

Log-likelihood: 1.11 / 1.19 / 1.19 / 1.20 / 1.20 / 1.20

We observe that the largest gain is observed when going for zero to a single fixed input, and that there are diminishing improvements thereafter. Importantly, observe that increasing the number of additional inputs never hurts performance. This demonstrates that the model is able to ignore additional inputs it doesn’t need. It also suggests that one can usually just set the number of additional inputs to a “sufficiently high number”. We suspect that this is due to dropout being included in the training procedure, so that the closest equivariant component is always learnt. We are currently running the same ablation for the synthetic GP experiment, and will provide these results in due course.

These results can also be visualised in Table 1 in the attached PDF.

Meta-review: The paper presents a general formalism for the construction of approximately equivariant operators agnostic to the choice of symmetry group and model architecture, which is then used to construct approximately equivariant neural processes. The method essentially uses an exactly equivariant model together with fixed inputs that break equivariance. The reviewers generally had a positive appraisal of the work, which I mostly share--I think the paper is a good contribution to the literature. However, the paper should be further revised for the camera ready. Some points are highlighted below. 

First off, the authors should consider incorporating the various clarifications that they have made during the discussion phase, making use of the extra page. Incorporating these points could amount to a moderately large revision. Beyond the broad discussion points, generally, the discussion to related work in the paper can also be improved. While the paper cites some approximately equivariant networks-based works in the first page, the paper can sometimes read like these works were an afterthought. The idea of approximate equivariance is now embedded in the literature for a few years and this should be prominently reflected in the narrative--this does not come through in reading the abstract, for example. The paper will also benefit in including additional references suggested by reviewer Ks9f. In the same vein, another work about generalization and approximation error characterization that is more relevant to the paper than some papers cited is: "Approximation-generalization trade-offs under (approximate) group equivariance", M Petrache, S Trivedi, NeurIPS 2023, as it is not specific to NNs. Reviewers 7yUH and dBf7 have also raised a valid concern about equivariance error quantification. While it is understandable that a detailed analysis is out of scope for the paper, this limitation should be highlighted more clearly in the paper. Finally, the concerns highlighted by dBf7 about the log-likelihood of context reconstruction should be taken note of. The authors have provided a response in the rebuttal. It would be beneficial to include these results in the paper with some accompanying discussion.
Decision: Accept (poster)