Title: Approximately Equivariant Neural Processes
Venue: NeurIPS 2024
Accepted: yes
Keywords: equivariance

neural processes

meta learning

deep learning

probabilistic methods
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

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 realworld regression experiments, showing that approximately equivariant NP models can outperform both their non-equivariant and strictly equivariant counterparts.

1 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 equivariance serves as a useful inductive bias-however, it is clear that it is not strictly translation equivariant due to geographical and seasonal variations. More generally, whilst there are aspects of modelling problems which are universal and exhibit symmetries such as equivariance (e.g. atmospheric physics) there are 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

Richard E. Turner

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 are neural processes [NPs; Garnelo et al., 2018a,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 et al., 2022a, van der Ouderaa et al., 2022, Finzi et al., 2021, Romero and Lohit, 2022]-which restricts their applicability to certain symmetry groups-and often require modifications to the loss function used to train the models [Finzi et al., 2021, Kim et al., 2023a]. As there exist NPs that utilise a 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-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 [TE-TNP; Ashman et al., 2024], and the relational CNP [RCNP; Huang et al., 2023].

We outline our core contributions as follows:

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 an equivariant mapping with fixed additional inputs (Theorems 2 and 3). This provides insight into how to construct approximately equivariant models, and generalises several existing approaches to relaxing equivariant constraints.

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 a 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 , y ) ∈ X × Y denotes an input-output pair. Let S = ⋃ ∞ N =0 ( X × Y ) N be a collection of all finite data sets, which includes the empty set ∅ . Let S M = ( X × Y ) M be the collection of M input-output pairs and S ≤ M = ⋃ M m =1 S m be the collection of at most M pairs. We denote a context and target set with D c , D t ∈ S , where |D c | = N c , |D t | = N t . Let X c ∈ ( X ) N c , Y c ∈ ( Y ) N c be the inputs and corresponding outputs of D c , and let X t ∈ ( X ) N t , Y t ∈ ( Y ) N t be defined analogously. We denote a single task as ξ = ( D c , D t ) = (( X c , Y c ) , ( X t , Y t )) . Let P ( X ) denote the collection of Y -valued stochastic processes on X . Let Θ 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-based parametrisations of prediction maps π : S → P ( X ) from data sets S to predictions P ( X ) , 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 ( Y | X , D ) = ∏ n p ( y n | x n , D ) . We denote all parameters of a CNP by ω . CNPs are trained in a meta-learning fashion, in which the expected predictive log-probability is maximised:

Here, the expectation is taken with respect to the distribution over tasks p ( ξ ) . In practice, we only have access to a finite number of tasks for training, so the expectation is approximated with an 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 or the group of rotations. Mathematically, a group G is a set endowed with a binary operation G × G → G (denoted as multiplication) such that (i) ( fg ) h = f ( gh ) for all f, g, h ∈ G ; (ii) there exists an identity element e ∈ G such that eg = ge = g for all g ∈ G ; and (iii) every element g ∈ G has an inverse g -1 ∈ G . A G -space is a space X for which there exists a function G × X → X called a group action (again denoted by multiplication) such that (i) ex = x for all x ∈ X ; and (ii) f ( gx ) = ( fg ) x for all f, g ∈ G and x ∈ X . The notion of group equivariance is used to describe mappings for which, when the input to the mapping is transformed by some g ∈ 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 ρ ( gx ) = gρ ( x ) for all g ∈ G and x ∈ X .

2.3 Group-Equivariant Conditional Neural Processes

The property of G -equivariance is particularly useful in NPs when modelling ground-truth stochastic processes that exhibit G -stationarity:

Definition 2 ( G -stationary stochastic process) . Let X be a G -space. We say that a stochastic process P ∈ P ( X ) is G -stationary if, for f ∼ P and all g ∈ G , f ( · ) is equal in distribution to f ( g · ) .

Let P ∈ P ( X ) be a ground-truth stochastic process. For a dataset D ∈ S , define π P ( D ) by integrating P against some density π ′ P ( D ) , such that d π P ( D ) = π ′ P ( D ) d P . π ′ P ( D ) is the RadonNikodym derivative of the posterior stochastic process with respect to the prior; intuitively, π ′ P ( D )( f ) is proportional to the likelihood, π ′ P ( D )( f ) = p ( D | f ) /p ( D ) , so it determines how the data is observed (e.g. under which noise?). Assume that π ′ P ( ∅ ) ∝ 1 so that π P ( ∅ ) = P . We say that π ′ P is G -invariant if π ′ P ( g D ) ◦ g = π ′ P ( D ) .

Proposition 1 ( G -stationarity and G -equivariance) . The ground-truth stochastic process P is G -stationary and π ′ 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 ∈ 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 a 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 is formalised by Elesedy and Zaidi [2021], Bulusu et al. [2021], Petrache and Trivedi [2023].

For a CNP, assume that every marginal p ( · | x , D ) is in some fixed parametric family with parameters θ ( x ; D ) ∈ Θ . For example, θ could consist of the mean and variance for a Gaussian distribution. Let C ( X , Θ) denote the set of continuous functions X → Θ . For a CNP, define the associated parameter map Φ: S → C ( X , Θ) by Φ( D )( x ) = θ ( x ; D ) . Intuitively, the parameter map Φ maps a dataset D to a function Φ( D ) giving the parameters Φ( D )( x ) for every input x . Assume that X is a G -space. We turn S and C ( X , Θ) into G -spaces by applying g to the inputs: g D ∈ S consists of the input-output pairs ( g x n , y n ) , and gθ ( x ; D ) ∈ C ( X , Θ) is defined by gθ ( x ; D ) = θ ( g -1 x ; D ) .

Definition 3 ( G -equivariant CNP) . A CNP is G -equivariant if the associated parameter map Φ is G -equivariant: Φ( g D ) = g Φ( D ) for all datasets D .

To parametrise a G -equivariant CNP, we must parametrise the associated parameter map Φ . Here we present a general construction by Kawano et al. [2021].

Theorem 1 (Representation of G -equivariant CNPs, Theorem 2 by Kawano et al. [2021]) . Let Y ⊆ R D be compact. Consider an appropriate collection S ′ ≤ M ⊆ S ≤ M . Then a function Φ: S ′ ≤ M →

C ( X , Θ) is continuous, permutation-invariant, and G -equivariant if and only if it is of the form

for some continuous ϕ : Y → R 2 D , an appropriate G -invariant positive-definite kernel ψ : X 2 → R (i.e. ψ ( x n , x m ) = ψ ( g x n , g x m ) for all g ∈ G ), and some continuous and G -equivariant ρ : H → C ( X , Θ) with H an appropriate G -invariant space of functions (i.e. closed under g ∈ G ).

Theorem 1 naturally gives rise to architectures which can be deconstructed into two components: an encoder and a decoder . The encoder, f enc : S → Z , maps datasets to some embedding space Z . The decoder, ρ : Z → C ( X , Θ) , takes this representation and maps to a function that gives the parameters of the CNP's predictive distributions: p ( y | x , D ) = p ( y | ρ ( f enc ( D ))( x )) where ρ ( f enc ( D ))( x ) ∈ Θ . For simplicity, we often view the decoder as a function Z × X → Θ and more simply write ρ ( f enc ( D ) , x ) . In Theorem 1, both the encoder f enc and decoder ρ are G -equivariant. Many neural process architectures are of this form, including the ConvCNP [Gordon et al., 2019], the EquivCNP [Kawano et al., 2021], 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, or approximated by, an equivariant mapping with additional, fixed functions as input. These results motivate a simple construction for approximately equivariant models, which we use to construct approximately equivariant 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 ( H , ⟨ · , · ⟩ ) be a Hilbert space of functions X → R . Let G be a group acting linearly on H from the left. For every g ∈ G and f ∈ H , applying g to f gives another function: gf ∈ H . By linearity, for f 1 , f 2 ∈ H , g ( f 1 + f 2 ) = gf 1 + gf 2 . Assume that H is separable and that H is G -invariant, meaning that, for all f 1 , f 2 ∈ H and g ∈ G , ⟨ gf 1 , gf 2 ⟩ = ⟨ f 1 , f 2 ⟩ . If H = L 2 ( X ) and X is a separable metric space, then H 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 ∈ B is compact if the image of every bounded set is relatively compact. Let ( e i ) i ≥ 1 be an orthonormal basis for H . Define P n = ∑ n i =1 e i ⟨ e i , · ⟩ . The following proposition is well known:

Proposition 2 (Finite-rank approximation of compact operators; e.g., Corollary 6.2 by Brezis [2011].) . Let T ∈ B . Then T is compact if and only if there exists a sequence of operators of finite rank ( T n ) n ≥ 1 ⊆ B such that ∥ T -T n ∥ → 0 . In particular, one may take T n = P n T , so T is a compact operator if and only if ∥ T -P n T ∥ → 0 .

Intuitively, compactness of an operator implies that its inputs and outputs can be well approximated with a finite-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 n ( · , t 1 , . . . , t 2 n ) where E n is G -equivariant and t 1 , . . . , t 2 n are the additional fixed inputs.

Theorem 2 (Approximation of non-equivariant linear operators.) . Let T ∈ B . Assume that T is compact. Then there exists a sequence of continuous nonlinear operators E n : H 1+2 n → H , n ≥ 1 , and a sequence of functions ( t n ) n ≥ 1 ⊆ H such that every E n is G -equivariant,

Proof. Observe that P n Tf = ∑ n i =1 e i ⟨ T ∗ e i , f ⟩ = ∑ n i =1 e i ⟨ τ i , f ⟩ , with τ i = T ∗ e i ∈ H . Define the continuous nonlinear operators E n : H 1+2 n → H as E n ( f, τ 1 , e 1 , . . . , τ n , e n ) = ∑ n i =1 e i ⟨ τ i , f ⟩ . By G -invariance of H , these operators are G -equivariant:

Since T is compact, ∥ T -P n T ∥ → 0 (Proposition 2), so ∥ T -E n ( · , τ 1 , . . . , τ n , e 1 , . . . , e n ) ∥ → 0 , which proves the result.

As an example, consider G to be the group of translations such that E n is translation equivariant. Translation-equivariant mappings between function spaces can be approximated with a CNN [Kumagai and Sannai, 2020]. Therefore, Theorem 2 gives that T ≈ CNN( · , t 1 , . . . , t k ) where t 1 , . . . , t k 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 k 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-equivariant operator.

Theorem 2 is not exactly what we set out to achieve: we wish to construct approximately equivariant mappings, not use equivariant mappings to approximate non-equivariant models. However, instead of applying Theorem 2 directly to T , consider the decomposition T = T equiv + T non-equiv where T equiv is in some sense the best 'equivariant approximation' of T and T non-equiv is the residual. Then, if we approximate T non-equiv with E non-equiv using Theorem 2, and approximate T equiv with some G -equivariant mapping E equiv directly, we have T ≈ E equiv + E non-equiv ( · , t 1 , . . . , t k ) . If k = 0 , this approximation roughly recovers E equiv, the best equivariant approximation of T ; and, if k is increased, the equivariance starts to break and the approximation starts to wholly approximate T . Specifically, in Theorem 2, k = 2 n determines the dimensionality of the range of the approximation E n . Therefore, a small k means that that one would deviate from T equiv 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 ∈ B , it is true that T is compact if and only if ∥ T -P n TP n ∥ → 0 (Proposition 3 in Appendix A). In generalising Theorem 2 to nonlinear operators, we shall use this equivalent condition. Roughly speaking, ∥ T -P n TP n ∥ → 0 says that both the domain and range of T admit a finite-dimensional approximation, and the proof then proceeds by discretising these finite-dimensional approximations.

Theorem 3 (Approximation of non-equivariant operators.) . Let T : H → H be a continuous, possibly nonlinear operator. Assume that ∥ T -P n TP n ∥ → 0 , and that T is ( c, α ) -Hölder for c, α > 0 , in the following sense:

Moreover, assume that the orthonormal basis ( e i ) i ≥ 1 is chosen such that, for every n ∈ N and g ∈ G , span { e 1 , . . . , e n } = span { ge 1 , . . . , ge n } , meaning that subspaces spanned by finitely many basis elements are invariant under the group action ( ∗ ). Let M > 0 . Then there exists a sequence ( k n ) n ≥ 1 ⊆ N , a sequence of continuous nonlinear operators E n : H 1+ k n → H , n ≥ 1 , and a sequence of functions ( t n ) n ≥ 1 ⊆ H such that every E n is G -equivariant,

and

If assumption ( ∗ ) does not hold, then the conclusion holds with E n ( · , t 1 , . . . , t k n ) replaced by E n ( P n · , t 1 , . . . , t k n ) . We provide a proof in Appendix A.

Condition ( ∗ ) says that the orthonormal basis ( e i ) i ≥ 1 must be chosen in a 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 H = L 2 ( S 1 ) where S 1 = R / Z is the one-dimensional torus ( [0 , 1] with endpoints identified); G the one-dimensional translation group; and ( e i ) i ≥ 1 the Fourier basis: e k ( x ) = e i 2 πkx . Then e k ( x -τ ) = e i 2 πkτ e k ( x ) ∝ e k ( x ) , so span { e k ( · -τ ) } = span { e k } . This example shows that the result holds for translation equivariance, which is the symmetry that we will primarily consider in the experiments. Importantly, Theorem 3 requires ∥ T -P n TP n ∥ → 0 to hold. For this example, this condition roughly means that the dependence of T on high-frequency basis elements goes to zero as the frequency increases.

In Theorem 3, the sequence ( k n ) n ≥ 1 determines how many additional fixed inputs are required to obtain a good approximation. For linear T (Theorem 2), k n grows linearly in n . In the nonlinear case, k n 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 k n in n depends on T . A promising idea is to consider ∥ Tgf -gTf ∥ as a measure of how equivariant a mapping is [Wang et al., 2022a].

3.1 Approximately Equivariant Neural Processes

Theorem 3 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 f enc and decoder ρ . The construction that we consider in this paper is to insert additional fixed inputs into the decoder ρ :

where the decoder ρ : H 1+ B → C ( X , Θ) now takes in the dataset embedding f enc ( D ) as well as B additional fixed inputs ( t b ) B b =1 ⊆ H . These become additional model parameters and break G -equivariance of the decoder. The more are included, the more non-equivariant the decoder becomes. Crucially, Theorem 3 shows that including sufficiently many additional inputs eventually recovers any non-equivariant decoder. Conversely, by only including a few of them, the decoder deviates from G -equivariance in only a few degrees of freedom. Hence, the number B of additional fixed inputs determines to which extent the decoder can become non-equivariant.

There are a myriad of ways in which ( t b ) B b =1 can be incorporated into existing architectures for ρ . If ρ is a CNN or 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 our implementation. For more general ρ , 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 transformer-based LLMs [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 or pixels, as their position is important to the overall structure. Following Theorem 3, we can add word- or pixelposition-specific 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 out-of-distribution. We stress that equivariance is crucial for models to generalise beyond the training distribution-this is the 'shared' component that is inherent to the system we are modelling. Whilst the non-equivariant component is able to learn local symmetrybreaking 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-equivariant component outside of the training domain and instead reverts to equivariant predictions. We achieve this in the following way: (i) During training, set t b = 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: t b ( 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 to recovering equivariance out-of-distribution also connects to Wang et al.'s (2022b) notion of ϵ -approximate equivariance. Let π AE-NP be the learned neural process, and let π E-NP be the neural process obtained by setting t b = 0 in π AE-NP. If the predictions of π AE-NP are not too different from those of π E-NP, because the fixed additional inputs only affect the predictions in a limited way, then π AE-NP is ϵ -approximately equivariant in the sense of Wang et al. [2022b]. See Appendix D.

4 Related Work

Group equivariance and equivariant neural processes. Group equivariant deep learning architectures are ubiquitous in machine learning, including CNNs [LeCun et al., 1989], group equivariant CNNs [Cohen and Welling, 2016a], 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 a significant body of research demonstrating improved sample complexity and generalisation capabilities [Mei et al., 2021, Elesedy and Zaidi, 2021, Bulusu et al., 2021, Zhu et al., 2021]. There exist a 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, 2016a,b]. 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 et al., 2022b] has demonstrated that both equivariant and approximately equivariant models achieve better generalisation bounds than data augmentation. Kim et al. [2023b] propose a similar approach to data augmentation, replacing a 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 a 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. 4 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. T ( n ) becomes T ). 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-based models [Kawano et al., 2021]: E and ˜ E . For experiments involving a large number of datapoints, we replace the TNP-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 a 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 kernel, except the lengthscale ℓ ( x ) is a 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 equivariant counterparts. We construct each dataset by first sampling a change point, either side of which the GP lengthscale is either small ( ℓ ( x ) = 0 . 1 ) or large ( ℓ ( 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 c ∼ U{ 1 , 64 } and the number of target points is set as N t = 128 . See Appendix E for a complete description.

(a)

TNP.

(b)

ConvCNP (

T

).

(c)

EquivCNP (

E

).

(d)

TNP (

T

).

(e)

ConvCNP (

˜

T

).

(f)

RelaxedConvCNP (

˜

T

).

(g)

EquivCNP (

˜

E

).

(h)

TNP (

˜

T

).

Context

Model

TE Model===

Ground truth

Context

Model

Ground truth

-2

0

-2

0

2

2

0

2

2

0

Context

Model

1

Ground truth

D

0

-2

0

2

0

0

-2

0

We evaluate the log-likelihood on both the in-distribution (ID) training domain and on an out-ofdistribution (OOD) setting in which the test domain is far away from the change point. Table 1 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 ID regime; and 2) generalise as well as the equivariant models when tested OOD. We provide an illustrative comparison of the predictive distributions for each model in Figure 1. More examples can be found in Appendix E. 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-lengthscale 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.

We show in Appendix E.1 that the approximately equivariant models deviate only slightly from strict equivariance, thus being equivariant in 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.1. The biggest improvement is obtained when going from 0 to 1 basis function, highlighting the importance of relaxing strict equivariance. We generally find that a few basis functions suffice to capture the symmetry-breaking features, but using more fixed inputs than necessary does not hurt performance, thus justifying in practice the use of a 'sufficiently high number' of fixed inputs.

5.2 Smoke Plumes

There are inherent connections between symmetries and dynamical systems, yet it is also true that real world dynamics rarely exhibit perfect symmetry. Motivated by this, we investigate the utility of

approximate equivariance in the context of modelling symmetrically-imperfect simulations generated from partial differential equations. We consider a setup similar to Wang et al. [2022a]-we construct a dataset of 128 × 128 2-D smoke simulations, computing the air flow in a closed box with a 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 ∼ U{ 5 , 30 } and the buoyancy coefficient B ∼ U [0 . 1 , 0 . 5] . 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 × 32 patch from each state to construct a dataset. We sample the number of context points according to N c ∼ U{ 10 , 250 } and set the remaining datapoints as target points. Table 1 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-breaking 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 are collected at a latitudinal and longitudinal resolution of 0 . 5 ◦ , and temporal resolution of an hour. We also have access to the surface elevation at each coordinate, resulting in a 4-D input ( x n ∈ R 4 ). We consider measurements collected from Europe and from central US. 5 . We train each model on Europe's 2018 data, and test on both Europe's and central US' 2019 data. Because the CNN-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.

2-D Spatial Regression. We sample datasets spanning 16 ◦ across each axis. Each dataset consists of a maximum of N = 1024 datapoints, from which the number of context points are sampled according to N c ∼ U{⌊ N/ 100 ⌋ , ⌊ N/ 3 ⌋} . The remaining datapoints are set as target points. Table 2 presents the average test log-likelihood 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. Surprisingly, 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-based models for a single 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-truth. 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 2f 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-equivariant component. As with the 1-D GP experiment, we analyse the degree of approximate equivariance and the performance with an increasing number of fixed inputs in Appendix E.3, drawing similar conclusions.

4-D Spatio-Temporal Regression. In this experiment, we sample datasets across 4 days with measurements every day and spanning 8 ◦ across each axis. Each dataset consists of a maximum of N = 1024 datapoints, from which the number of context points are sampled according to N c ∼ U{⌊ N/ 100 ⌋ , ⌊ N/ 3 ⌋} . The remaining datapoints are set as target points. We provide results 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.

1400

(a)

Context.

(b)

Ground-truth.

(c)

Elevation.

(d)

PT-TNP:

T

.

(e)

PT-TNP:

˜

T

.

(f)

PT-TNP:

˜

T

-

T

.

(g)

EquivCNP:

E

.

(h)

EquivCNP:

˜

E

.

(i)

EquivCNP:

E

-

E

.

270

275

280

285

270

275

280

285

200

400

009

800

1000

1200

46.5N

46.5°N

46.5N

45°N

45°N

45°N

43.5°N

43.5°N

43.5°N

42°N

42°N

42°N

40.5°N

40.5°N

40.5°N

270

275

280

285

270

275

280

285

46.5°N

46.5°N

46.5°N

45°N

45°N

45N

43.50N

43.5°N

43.5°N

42N

42N

42°N

40.5°N

40.5°N

40.5°N

270

275

280

285

270

275

280

285

b

-2

46.5°N

46.5°N

46.5°N

45°N

45°N

45°N

43.5°N

43.5°N

43.5°N

42°N

42°N

42°N

40.5°N

40.5°N

40.5°N

17.5°E20°E22.5°E25°E27.5°E30°E32.5°E

17.5°E20°E22.5°E25°E27.5°E30°E32.5°E

17.5°E20°E22.5°E25°E27.5°E30°E32.5°E

˜

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 nonequivariant and strictly equivariant counterparts on a 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.

References

Matthew Ashman, Cristiana Diaconu, Junhyuck Kim, Lakee Sivaraya, Stratis Markou, James Requeima, Wessel P Bruinsma, and Richard E Turner. Translation-equivariant transformer neural processes. In International conference on machine learning . PMLR, 2024.

Haim Brezis. Functional analysis, Sobolev spaces and partial differential equations , volume 2. Springer, 2011.

Wessel P. Bruinsma. Convolutional Conditional Neural Processes . PhD thesis, Department of Engineering, University of Cambridge, 2022. URL https://www.repository.cam.ac.uk/ handle/1810/354383 .

Srinath Bulusu, Matteo Favoni, Andreas Ipp, David I Müller, and Daniel Schuh. Generalization capabilities of translationally equivariant neural networks. Physical Review D , 104(7):074504, 2021.

Andrew Carr and David Wingate. Graph neural processes: Towards Bayesian graph neural networks. arXiv preprint arXiv:1902.10042 , 2019.

Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning , pages 2990-2999. PMLR, 2016a.

Taco S Cohen and Max Welling. Steerable cnns. arXiv preprint arXiv:1612.08498 , 2016b.

Copernicus Climate Change Service. Near surface meteorological variables from 1979 to 2019 derived from bias-corrected reanalysis, 2020. URL https://cds.climate.copernicus.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 at scale. In International Conference on Learning Representations , 2021. URL https://openreview. net/forum?id=YicbFdNTTy .

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 bottlenecked attentive neural processes. arXiv preprint arXiv:2211.08458 , 2022.

Marc 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.

Marc Finzi, Gregory Benton, and Andrew G 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-8295, 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 on machine learning , pages 1704-1713. PMLR, 2018a.

Marta Garnelo, Jonathan Schwarz, Dan Rosenbaum, Fabio Viola, Danilo J Rezende, SM Eslami, and Yee Whye Teh. Neural processes. arXiv preprint arXiv:1807.01622 , 2018b.

Mark N 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 E Turner. Convolutional conditional neural processes. arXiv preprint arXiv:1910.13556 , 2019.

Peter Holderrieth, Michael J 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://openreview. net/forum?id=HyeSin4FPB .

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 preprint arXiv:2306.10915 , 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 arXiv:2102.08759 , 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:1901.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 probabilistic symmetrization for architecture agnostic equivariance. Advances in Neural Information Processing Systems , 36:18582-18612, 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 a back-propagation network. Advances in neural information processing systems , 2, 1989.

Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set transformer: A framework for attention-based permutation-invariant 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 arXiv:1711.05101 , 2017.

Song Mei, Theodor Misiakiewicz, and Andrea Montanari. Learning with invariances in random features and kernel models. In Conference on Learning Theory , pages 3351-3418. PMLR, 2021.

Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. arXiv preprint arXiv:2207.04179 , 2022.

Mircea Petrache and Shubhendu Trivedi. Approximation-generalization trade-offs under (approximate) 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-Assisted Intervention-MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18 , pages 234-241. Springer, 2015.

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: 33818-33830, 2022.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz 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.

A Proof of Theorem 3

Proposition 3. Let T ∈ B . Then T is compact if and only if ∥ T -P n TP n ∥ → 0 .

Proof. Assume that ∥ T -P n TP n ∥ → 0 . Then T is the limit of a sequence of finite-rank operators, so T is compact (Proposition 2).

On the other hand, assume that T is compact. Then ∥ T -P n T ∥ → 0 (Proposition 2). Use the triangle inequality to bound

We already know that ∥ T -P n T ∥ → 0 , and it is true that ∥ P n ∥ = 1 . Finally, since T is compact, the adjoint T ∗ is also compact. Hence, again by Proposition 2,

which proves that ∥ T -P n TP n ∥ → 0 .

Theorem 3 (Approximation of non-equivariant operators.) . Let T : H → H be a continuous, possibly nonlinear operator. Assume that ∥ T -P n TP n ∥ → 0 , and that T is ( c, α ) -Hölder for c, α > 0 , in the following sense:

Moreover, assume that the orthonormal basis ( e i ) i ≥ 1 is chosen such that, for every n ∈ N and g ∈ G , span { e 1 , . . . , e n } = span { ge 1 , . . . , ge n } , meaning that subspaces spanned by finitely many basis elements are invariant under the group action ( ∗ ). Let M > 0 . Then there exists a sequence ( k n ) n ≥ 1 ⊆ N , a sequence of continuous nonlinear operators E n : H 1+ k n → H , n ≥ 1 , and a sequence of functions ( t n ) n ≥ 1 ⊆ H such that every E n is G -equivariant,

and

If assumption ( ∗ ) does not hold, then the conclusion holds with E n ( · , t 1 , . . . , t k n ) replaced by E n ( P n · , t 1 , . . . , t k n ) . We provide a proof in Appendix A.

Proof. Let ε > 0 . Choose n ∈ N such that ∥ T -P n TP n ∥ < 1 2 ε , and choose h > 0 such that ch α = 1 2 ε . Set h n = h/ √ 2 n . Consider the following collection of vectors:

By construction of A , for every u ∈ H such that ∥ u ∥ ≤ M , there exists an a ∈ A such that ∥ P n u -a ∥ 2 ≤ nh 2 n = 1 2 h 2 < h 2 .

Let ˜ k : R → [0 , ∞ ) be a continuous function with support equal to ( -h, h ) . Set k : H × H → [0 , ∞ ) , k ( u, v ) = ˜ k ( ∥ u -v ∥ ) . Consider the following map:

where 0 / 0 is defined as 0 . This map E is continuous: the numerator is a continuous H -valued function, and the denominator is continuous R -valued function which is non-zero wherever the numerator is non-zero. Moreover, by G -invariance of H , k is G -invariant, so E is G -equivariant:

Now set t i = ( P n TP n )( a i ) . Consider some u ∈ H such that ∥ u ∥ ≤ M . By construction of A , ∑ | A | i =1 k ( P n u, a i ) > 0 . Therefore, by ( c, α ) -Hölder continuity of T ,

where (i) uses the triangle inequality; (ii) ( P n TP n )( u ) = ( P n TP n )( P n u ) and ( c, α ) -Hölder continuity of P n TP n ; and (iii) that ∥ P n u -a i ∥ ≥ h implies that k ( P n u, a i ) = 0 .

Therefore, for all u ∈ H such that ∥ u ∥ ≤ M ,

Finally, if condition ( ∗ ) holds, then P n is G -equivariant, which we show now. Let u ∈ H and consider g ∈ G . Since span { e 1 , . . . , e n } = span { ge 1 , . . . , ge n } and ge 1 , . . . , ge n forms an orthonormal basis for span , { e 1 , . . . , e n } , we can also write P n = ∑ i =1 ge i ⟨ ge i , · ⟩ , so

Therefore, in the case that condition ( ∗ ) holds, P n can be absorbed into the definition of E n .

From the proof, we find that k n = 2 | A | = 2(1 + 2 ⌈ √ 2 nM/h ⌉ ) n ≤ 2(2 + 2 √ 2 nM/h ) n and ch α = 1 2 ε , which implies that h = ( ε/ 2 c ) 1 /α , so k n ≤ 2(2 + 2 √ 2 nM ( ε/ 2 c ) -1 /α ) n . 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] is a translation equivariant NP and is constructed in exactly the form shown in Equation 2, with ψ : X 2 → R an RBF kernel with learnable lengthscale and ϕ ( y i ) = [ y T i ; 1 ] T . The functional embedding is discretised and passed pointwise through an MLP to some final functional representation ˜ e ( D ) : ˜ X → R D z , where ˜ X denotes the discretised input domain. ρ implemented as a CNN through which ˜ e ( D ) is passed through together with another RBF kernel ψ p which maps back to a continuous function space. We provide pseudo-code for a forward pass through the ConvCNP in Algorithm 1.

Algorithm 1: Forward pass through the ConvCNP ( T ) for off-the-grid data.

Input:

ρ

,

)

= (CNN

Context

D

c

c

X

= (

and target

begin

lower upper

,

{

i

x

˜

}

h

T

i

=1

i

h

h

(1)

i

i

{

N

∑

p

, ψ

,

c

Y

X

, and density

ψ

{

x

) =

(

t

c,n c,n

y

,

)

γ

}

t,m

x

=

{

range (

t

N

m

=1

}

X

t

grid

c

n

=1

lower

(

∪

,

[1

,

h

(1)

i

/

f

h

MLP(

(˜

i

x

i

h

}

)

θ

m

T

i

=1

∑

return

T

p

(

i

=1

f

m

y

T

X

c

upper

T

c,n

]

(0)

i

;

)

;

CNN

(

(˜

i

x

)

ψ

p

)

;

ψ

, γ

(˜

i

x

-

{

˜

x

(

i

,

x

h

t,m

;

)

)

x

i

;

c,n

T

i

}

-

=1

˜

)

)

x

i

)

;

.

N

c

n

;

;

=1

·|

θ

end end

To achieve approximate translation equivariance, we construct a learnable function t : ˜ X → R D z using an MLP which represents the fixed inputs. We sum these together with the resized functional embedding, resulting in the overall implementation ρ ( e ( D ) , t 1 , . . . , t B ) = CNN(˜ e ( D ) + t ) . Note that the number of fixed inputs is B = D z . 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 of the ConvCNP to more general group equivariances. As with the ConvCNP, it is constructed in the form shown in Equation 2. We achieve E -equivariance with ψ : X 2 → R an RBF kernel with learnable lengthscale and ϕ ( y i ) = [ y T i ; 1] T . The functional embedding is discretised and passed pointwise through an MLP to some final functional representation ˜ e ( D ) : ˜ X → R D z , where ˜ X denotes the discretised input domain. ρ is implemented as a group equivariant CNN [Cohen and Welling, 2016a] together with another RBF kernel ψ p which maps back to a continuous function space. The approximately E -equivariant EquivCNP is 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] is a translation equivariant NP. However, unlike the ConvCNP, the function embedding representation differs slightly to that in Equation 2, and is given by

where

Here, µ is implemented using an MLP. This is a partially evaluated translation equivariant multi-head cross-attention (TE-MHCA) operation, with attention weights computed according to Equation 21. We refer to ϕ ( y i ) as the initial token embedding for y i . ρ 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 to a continuous function space. We provide pseudo-code for a forward pass through the TE-TNP in Algorithm 3.

,

Algorithm 2: Forward pass through the ConvCNP ( ˜ T ) for off-the-grid data.

Input:

ρ

= (CNN

)

,

input

t

, ψ

p

. Context

D

c

= (

X

c

,

Y

c

) =

and target

begin

lower upper

,

{

i

x

˜

}

h

T

i

=1

i

h

h

(1)

i

i

{

N

∑

X

ψ

, density

{

(

x

t

=

{

x

c,n

,

y

γ

, and fixed c,n

t,m range (

}

N

t

m

=1

X

t

grid

c

n

lower

(

∪

,

[1

,

=1

(1)

i

h

/

f

h

MLP(

(˜

i

x

i

h

}

)

θ

m

T

i

=1

T

∑

return

p

(

i

=1

f

θ

m

T

c

X

upper

]

T

c,n

y

(0)

i

;

) +

t

( ˜

ψ

x

CNN

, γ

(˜

i

x

-

i

)

(

(˜

x

i

)

ψ

{

p

)

;

(

;

˜

x

x

i

h

,

t,m

)

}

,

)

;

)

;

x

i

}

-

c,n

T

i

=1

˜

)

N

c

n

=1

;

)

x

i

)

;

;

·|

Algorithm 3: Forward pass through the TETNP ( T ).

begin

Algorithm 4: Forward pass through the TETNP ( ˜ T ).

begin

Unlike the ConvCNP and EquivCNP, we do not 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, we take an alternative approach in which the initial token representations for context set are modified by summing together with the fixed input value at the corresponding location. This modification still conforms to 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 4.

B.4 Pseudo-Token Translation Equivariant Transformer Neural Process

The pseudo-token TE-TNP (PT-TE-TNP) is an alternative to the TE-TNP which avoids the O ( N 2 ) computational cost of regular transformers through the use of pseudo-tokens [Feng et al., 2022, Ashman et al., 2024, Lee et al., 2019]. The functional embedding representation for the PT-TE-TNP is given by

where and

The attention mechanism used in the PT-TE-TNP is the same as in the TE-TNP, and is given in Equation 21. Intuitively, the initial token embeddings for the dataset D are 'summarised' by the pseudo-tokens { u m } M m =1 and pseudo-token input locations { v m } M m =1 . ρ is implemented using a 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 is taken [Jaegle et al., 2021], as in the latent bottlenecked attentive NP of Feng et al. [2022]. The induced set transformer approach of Lee et al. [2019] can also be used.

Algorithm 5: Forward pass through the PTTE-TNP ( T ).

Input:

begin

z

c,n

v

ρ

{

=

TE-MHSA

(

ϕ

,

V

v

=

{

m,

Context

D

c

X

= (

c

and target

ℓ

0

,

Y

X

)

,

TE-MHCA

}

j M

(

ℓ

)

1

=1

c

t

c,n

ϕ

(

y

m

u

j

v

z

m,

u

t,m j,

0

0

;

for

z

ℓ

= 1

u

0

;

, . . . , L

j

{

do

(

ℓ

,

U

=

{

u

) =

=

{

(

x

{

c,n t,m

x

∑

,

}

c

N

i

=1

t

N

m

=1

x

)

1

j

TE-MHCA

j

=1

(

u

c,i

;

,

(

Z

ℓ

c

,

)

,

(

z

0

j,

y

,

TE-MHCA

(

ℓ

M

m

}

c,n

2

,

0

.

z

=1

)

}

N

c

n

v

U

=1

j

,

,

V

t,m

,

,

X

)

;

V

)

;

c

TE-MHSA

}

u

j M

z

t,m

(

ℓ

)

2

TE-MHCA

(

z

,

U

MLP(

end

θ

m

z

t,m return

p

·|

(

θ

m

)

)

;

;

t,m

;

)

+1

/N

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 a forward pass through the approximately translation equivariant PT-TE-TNP in Algorithm 6.

B.5 Relaxed Convolutional Conditional Neural Process

The RelaxedConvCNP is equivalent to the ConvCNP with the RelaxedCNN of Wang et al. [2022a] used in place of a 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

This modifies the kernel weights k l ( u -1 v ) with the input-dependent function w l ( v ) , which is our fixed input. To enable the fixed inputs to be zeroed out, we modify this slightly as

so that when t l ( v ) = 0 we recover the standard G -equivariant convolution. We make t l ( v ) a 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 ∈ G on f : G → R as

We begin with the approach of Wang et al. [2022a], which defines the relaxed group convolution as

)

;

)

Algorithm 6: Forward pass through the PTTE-TNP ( ˜ T ).

Input:

}

L

ℓ

=1

,

begin

z

c,n

,

TE-MHCA

(

ρ

{

=

TE-MHSA

ϕ

m,

,

V

v

=

{

Context

D

c

c

X

= (

and target

ℓ

0

,

Y

X

)

}

j M

(

ℓ

1

=1

c

t

c,n

ϕ

(

y

v

m

u

j

v

z

m,

u

t,m j,

0

for

z

ℓ

= 1

u

;

)

0

+1

+

/N

t

(

v

0

;

, . . . , L

j

{

do

(

ℓ

,

U

u

=

{

) =

=

x

(

{

{

c,n t,m

x

∑

,

N

y

t

m

}

c

N

i

;

=1

=1

x

)

1

j

TE-MHCA

j

=1

(

u

j M

c,i

;

,

(

Z

ℓ

c

,

)

,

j,

0

c,n

)

,

TE-MHCA

(

ℓ

M

m

}

2

,

0

.

z

=1

)

}

N

c

n

v

(

=1

j

,

,

U

z

t,m

V

,

,

X

c

)

;

V

)

;

TE-MHSA

}

u

z

t,m

(

ℓ

)

2

TE-MHCA

(

z

,

U

MLP(

end

θ

m

z

t,m return

p

(

·|

θ

m

)

)

;

;

t,m

m

)

)

;

)

}

L

ℓ

=1

,

where µ ( v ) is the left Haar measure on G . Observe that this replaces a single kernel with a set of kernels { k l } L l =1 , whose contributions are linearly combined with coefficients that vary with v ∈ G . This can be expressed as

where w l correspond to the fixed inputs and E is used to denote equivariant operators w.r.t. the group G . Clearly each E l is G -equivariant with respect to its inputs, as applying a transformation g ∈ G to the inputs gives

Since the sum of G -equivariant functions is itself G -equivariant, E ( f, w 1 , . . . , w L )( u ) is G -equivariant with respect to its inputs.

The approach of van der Ouderaa et al. [2022] is more general than that of Wang et al. [2022a], and relax strict equivariance through convolutions with input-dependent kernels:

Define a fixed input t ( u ) = u . We can express the above convolution as a G -equivariant operation on t and f :

To demonstrate G -equivariance, consider applying a transformation g ∈ G to the inputs:

Thus, this method is also a special case of ours.

D Achieving ϵ -Approximate G -Equivariance

Wang et al. [2022b] introduce a useful metric for assessing the degree to which mappings are approximately equivariant. Specifically, let π : X → Y denote some mapping between G -spaces X and Y . Let d : Y × Y → R be a metric on Y . We say π is ϵ -approximately G -equivariant if for all g ∈ G and x ∈ X

Let π AE-NP : S → C ( X , Θ) be a learned neural process that breaks equivariance with fixed additional inputs ( t b ) B b =1 , and let π E-NP : S → C ( X , Θ) be the strictly equivariant neural process obtained by setting t b = 0 in π AE-NP. If the predictions of π AE-NP are not too different from those of π E-NP, because the fixed additional inputs only affect the predictions in a limited way, then in this appendix we argue that π AE-NP is ϵ -approximately equivariant.

To assume that predictions of π AE-NP are not too different from those of π E-NP, assume that there exists some metric d and ϵ > 0 such that, for all D , d ( π AE-NP ( D ) , π E-NP ( D )) ≤ ϵ . For example, the metric can be the root-mean-squared-error (RMSE) between the mean functions x ↦→ E [ f ( x )] of two stochastic processes by integrating over the input domain X . In practice, we will find that d ( π AE-NP ( D ) , π E-NP ( D )) ≤ ϵ holds for D over the training domain. If we assume that π AE-NP has a finite receptive field-in the sense that the output stochastic process has only local dependencies-then, provided we zero out the additional fixed inputs as described in Section 3.1, d ( π AE-NP ( D ) , π E-NP ( D )) ≤ ϵ will hold over the entire input domain.

Finally, to show that π AE-NP is ϵ -approximately equivariant, we use the triangle inequality:

where d ( gπ E-NP ( D ) , π E-NP ( gD ) ) = 0 because π E-NP is strictly equivariant.

E Experiment Details and Additional Results

E.1 Synthetic 1-D Regression

We consider a 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 is a non-stationary generalisation of the squared exponential kernel, where the lengthscale parameter becomes a function of position l ( x ) :

We consider a 1-D space with two regions with constant, but different lengthscale - one with l = 0 . 1 , and one with l = 4 . 0 . The lengthscale changepoint is situated at x = 0 . We randomly sample with a 0.5 probability the orientation (left/right) of the low/high lengthscale region. Formally, l ( x ) = (0 . 1 β +4(1 -β )) δ [ x < 0] + (0 . 1(1 -β ) + 4 β ) δ [ x ≥ 0] , where β ∼ Bern (0 . 5) . For each task, we sample the number of context points N c ∼ U{ 1 , 64 } and set the number of target points to N t = 128 . The context range [ x c, min , x c, max ] (from which the context points are uniformly sampled) is an interval of length 4, with its centre randomly sampled according to U [ -7 , 7] for the ID task, and according to U [13 , 27] for the OOD task. The target range is [ x t, min , x t, max ] = [ x c, min -1 , x c, max +1] . This is also applicable during testing, with the test dataset consisting of 80,000 datasets.

We use an embedding / token size of D z = 128 for the TNP-based models and D z = 64 for the ConvCNP-based ones, and a decoder consisting of an MLP with two hidden layers of dimension D z . The decoder parameterises the mean and pre-softplus variance of a Gaussian likelihood with heterogeneous noise. Model specific architectures are as follows:

TNP The initial context tokens are obtained by passing the concatenation [ x, y, 1] through an MLP with two hidden layers of dimension D z . The initial target tokens are obtained by passing the concatenation [ x, 0 , 0] through the same MLP. The final dimension of the input acts as a 'density' channel to indicate whether or not an observation is present. The TNP encoder consists of nine layers of self-attention and cross-attention blocks, each with H = 8 attention heads with dimensions D V = D QK = 16 . In each of the attention blocks, we 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 a layer-normalisation followed by a pointwise MLP with two hidden layers of dimension D z .

TNP ( T ) For the TNP ( T ) 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, we implement ρ ℓ : R H × R D x → R H as an MLP with

two hidden layers of dimension D z . The initial context token embeddings are obtained by passing the context observations through an MLP with two hidden layers of dimension D z . The initial target token embeddings are sampled from a standard normal. Pseudo-code for a forward pass through the TNP ( T ) can be found in Appendix B.3.

TNP ( ˜ T ) The architecture of the TNP ( ˜ T ) is 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 D z . 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.

ConvCNP ( T ) For the ConvCNP model, we use a CNN with 9 layers. We use C = 64 channels, a kernel size k = 21 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 a continuous function. Pseudo-code for a forward pass through the ConvCNP ( T ) can be found in Appendix B.1.

ConvCNP ( ˜ T ) The ConvCNP ( ˜ T ) closely follows the ConvCNP ( T ) 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 an identical architecture to the ConvCNP( T ), with regular convolutional 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 a forward pass through the EquivCNP ( E ) can be found in Appendix B.2.

EquivCNP ( ˜ E ) We use an identical architecture to the ConvCNP ( ˜ T ), with symmetric convolutions in place of regular convolutions. Pseudo-code for a 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 a learning rate of 5 × 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 1 we compared the predictive distributions of the eight considered models for a 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 = 0 , 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 to more accurately capture the uncertainties beyond the transition point, into the high-lengthscale region. This indicates that the approximately equivariant models are better suited to cope with non-stationarities in the data.

When the context region only spans the high-lengthscale region, both the strictly and approximately equivariant models output predictions that closely follow the ground truth. However, the nonequivariant 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.

(a)

TNP.

(b)

ConvCNP (

T

).

(c)

EquivCNP (

E

).

(d)

TNP (

T

).

(e)

ConvCNP (

˜

T

).

(f)

RelaxedConvCNP (

˜

T

).

(g)

EquivCNP (

˜

E

).

(h)

TNP (

˜

T

).

Context

Model

TEModel-=-

Ground truth

D

2

-2

7--9-8-01-

7-t-9-8-0I-0I

4

10

7--9-8-01-

8

7--9-8-01-01

024

10

2

7-t-9-8-01-

10

-10-8

-6

2

10

-10

-8

-2

0

4

8

10

7-t-9-8-01-

0

2

4

6

8

10

(a)

TNP.

(b)

ConvCNP (

T

).

(c)

EquivCNP (

E

).

(d)

TNP (

T

).

(e)

ConvCNP (

˜

T

).

(f)

RelaxedConvCNP (

˜

T

).

(g)

EquivCNP (

˜

E

).

(h)

TNP (

˜

T

).

Context

Model

Ground truth

3

2

1

0

-8

-6

-42

8

10

-6

4

2

-10

-8

-4

-2

D

-10

2

4

10

3

2

1

0

4-2

0

4

10

-8

-6

-4

-2

0

2

4

10

0

10

2

4

Log-likelihood of 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 3 show that the model is able to accurately model the context sets, with loglikelihood values close to the ground truth log-likelihood of the data of 0 . 2806 ± 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-equivariant 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

where µ approx-equiv represents the predictive mean of the approximately equivariant model, and µ equiv the predictive mean of the same model with zeroed-out basis functions (i.e. the equivariant component of the approximately equivariant model). In this case, we consider the L 1 norm (i.e. p = 1 ).

The equivariance deviations from Table 4 are around 8 -9% (with the exception of the RelaxedConvCNP ( ˜ T ) 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 a data-driven manner (in this case, the lengthscale changepoint).

Analysis of number of fixed inputs In this ablation, we 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 a

G -equivariant model, while in the limit of infinitely many fixed inputs, the decoder can become fully non-equivariant.

In the main experiment we used 64 fixed inputs for the ConvCNP ( ˜ T ) and 4 Fourier coefficients for TNP ( ˜ T ). Table 5 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 ≥ 1 , suggesting that one Fourier coefficient is enough to capture the symmetry-breaking feature of the dataset. For the ConvCNP ( ˜ T ) the performance increases up to B ≥ 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 or 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.

E.2 Smoke Plumes

The smoke plume dataset consists of 128 × 128 2-D smoke simulations for different initial conditions generated through PhiFlow[Holl et al., 2020]. Hot smoke is emitted from a circular region at the bottom, and the simulations output the resulting air flow in a closed box. We also introduce a 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 ∼ U [5 , 30] . Moreover, we randomly choose its position among three possible x-axis locations: { 30 , 64 , 110 } (but we keep the y position fixed at 5). Finally, we sample the buoyancy coefficient of the medium according to B ∼ U [0 . 1 , 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 ∆ t = 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 c ∼ U { 10 , 250 } , with the remaining points set as the target points.

We use an embedding / token size of D z = 128 for the PT-TNP-based models and D z = 16 for the ConvCNP-based models. The decoder consists of an MLP with two hidden layers of dimension D z . The decoder parameterises the mean and pre-softplus standard deviation of a 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 described in Appendix E.1. 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.1. The initial pseudo-tokens and pseudo-input-locations are sampled from a standard

?

normal. We use 128 pseudo-tokens. Pseudo-code for a forward pass through the PT-TNP ( T ) 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 of the extra fixed inputs that are added to the pseudo-token representation. These are obtained by passing the pseudo-token locations through an MLP with two hidden layers of dimension D z . We use D z fixed inputs and apply dropout to them with a probability of 0 . 5 . Pseudo-code for a forward pass through the PT-TNP ( ˜ T ) can be found in Appendix B.4.

ConvCNP ( T ) For the ConvCNP model, we use a U-Net [Ronneberger et al., 2015] architecture for the CNN with 9 layers. We use C in = 16 input channels and C out = 16 output channels, with the number of channels doubling / halving on the way down / up. Between each downwards layer we apply pooling with size two, and between each upwards layer we linearly up-sample to recover the size. We use a kernel size of k = 9 with a stride of one. We use the natural discretisation of the 128 × 128 grid.

ConvCNP ( ˜ T ) We use the same architecture as the ConvCNP ( T ). The fixed inputs are obtained by passing the discretised grid locations through an MLP with two hidden layers of dimension C in . We use C in fixed inputs and apply dropout with probability 0 . 5 .

RelaxedConvCNP ( ˜ T ) We use an identical architecture to the ConvCNP ( T ), with regular convolutional operations replaced by relaxed convolutions. We use L = 1 kernels such that the total parameter count remains similar. We obtain the additional fixed inputs by passing the effective discretised grid at each layer (after applying the same pooling / up-sampling operations as the U-Net) through an MLP with two hidden layers of dimension C in .

EquivCNP ( E ) For the EquivCNP ( E ) model, we use a steerable E -equivariant CNN architecture consisting of nine layers, each with C = 16 input / output channels. We use a kernel size of k = 9 and stride of one. We discretise the continuous rotational symmetries to integer multiples of 2 π/ 8 . We use the natural discretisation of the 128 × 128 grid.

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 a learning rate of 5 × 10 -4 . For the ConvCNP T and ˜ T , EquivCNP T and ˜ T , and non-equivariant PT-TNP models we use a batch size of 16, while for the PT-TNP T and ˜ T models we use a batch size of 8. Gradient value magnitudes are clipped at 0.5. We train for a maximum of 500 epochs, with each epoch consisting of 16,000 datasets for a batch size of 16, and

8,000 datasets for a 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.

Additional Results We show in Figure 6 a comparison between the predictive means, as 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 a significant difference between the equivariant and approximately equivariant model.

GT.

PT-TNP (

∅

,

T

, and

˜

T

).

ConvCNP (

T

and

˜

T

).

EquivCNP (

E

and

˜

E

).

Context.

E.3 Environmental Data

The environmental dataset consists of surface air temperatures derived from the fifth generation of the European Centre for Medium-Range Weather Forecasts (ECMWF) atmospheric reanalyses (ERA5) [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 [35 ◦ , 60 ◦ ] / [10 ◦ , 45 ◦ ] ) and in the US (latitude / longitude range of [ -120 ◦ , -80 ◦ ] / [30 ◦ , 50 ◦ ] ). In both the 2-D and 4-D experiment, we train on Europe'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 a [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 to N c ∼ U{⌊ N 100 ⌋ , ⌊ N 3 ⌋} , with the remaining set as target points.

For all models, we use a decoder consisting of an MLP with two hidden layers of dimension D z . The decoder parameterises the mean and pre-softplus variance of a Gaussian likelihood with heterogeneous noise. Model specific architectures are as follows:

PT-TNP ( ˜ T ) 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 ◦ ] .

ConvCNP ( ˜ T ) Same as Appendix E.2 with fixed inputs zeroed outside 2018 and the latitude / longitude range of [35 ◦ , 60 ◦ ] / [10 ◦ , 45 ◦ ] .

RelaxedConvCNP ( ˜ T ) Same as Appendix E.2 with fixed inputs zeroed outside 2018 and the latitude / longitude range of [35 ◦ , 60 ◦ ] / [10 ◦ , 45 ◦ ] .

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 all models, we optimise the model parameters using AdamW [Loshchilov and Hutter, 2017] with a learning rate of 5 × 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 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 ( ˜ T ) model with dropout probability of p = 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).

Log-likelihood of context data We perform the same assessment as in Appendix E.1 to check whether the model is able to 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-likelihoods than in Table 2, where the models are tested on target sets.

Analysis of equivariance deviation We repeat the equivariance deviation analysis from Appendix E.1 on the 2-D environmental regression dataset. In this case we use the L 2 instead of the L 1 norm. Note that we only compute this on Europe, since the predictions on US are obtained by zeroing-out the basis function, leading to an equivariance deivation of 0 . 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

symmetry-breaking components of the weather dataset, with one important such component being the topography.

Analysis of 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 ( ˜ T ) from the 2-D regression experiment. The fixed inputs are linearly transformed to 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 is obtained by going from 0 to a single fixed input. Using more than 1 fixed input provides diminishing gains, but, importantly, never hurts performance. The saturation in performance for B ≥ 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 .

NeurIPS Paper Checklist

The checklist is designed to 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 a short (1-2 sentence) justification 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 it (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] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.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"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a 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.

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 use 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 to the full proof in the Appendix. We show supporting empirical evidence to justify our claims and contributions in Section 5.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in 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 or 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 can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it 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]

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 a 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 approximations only holding locally). The authors should reflect on 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 a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of 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 acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important 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 a complete (and correct) proof?

Answer: [Yes]

Justification: All novel theorems include either the full proof, or a 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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it 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:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to 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 a 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 one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a 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 a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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 to the data and code, with sufficient instructions 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.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

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) is 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, hyperparameters, 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 a 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, or 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, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run 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).

It 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 one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not 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 in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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.

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 of compute required for each of the individual experimental runs as well as 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 it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: The paper conforms to the Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. 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 applications.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no 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 to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an 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 to 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 as 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 to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: 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 (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: This paper does not pose such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a 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 is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or 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?

Answer: [Yes]

Justification: We provide complete citations for all datasets used in this paper.

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are 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

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

Justification: We released documented code.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model 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 as details about compensation (if any)?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, 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

Question: 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?

Answer: [NA]

Justification: This paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not 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 NeurIPS 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe proposed method is applicable regardless of architecture and easy to be implemented because it is simply adding some parameters as inputs.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThere is no limitation of the contribution itself but experimental limitations are mentioned in the weakness.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nEdit 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.\n\nQUESTIONS:\n0) 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)\n\nLIMITATIONS:\nthese 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.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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).\n\nQUESTIONS:\n1. 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\n\nLIMITATIONS:\nSome discussion in conclusion", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nLine 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-14"}