Approximately Equivariant Neural Processes

Authors: matthew  ashman
cristiana  diaconu
adrian  weller
wessel  bruinsma
richard e turner
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 15:19:31

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.

Main paper: 
 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 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:1. We demonstrate that under certain technical conditions, such as Hölder regularity and compactness, any non-equivariant mapping between function spaces can be approximated by 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. 2. We apply this result and the insights it provides to construct approximately equivariant versions of several popular equivariant NPs. The modifications required are very simple, yet effective: we demonstrate improved performance relative to both strictly equivariant and non-equivariant models on a number of spatio-temporal regression tasks. 
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 ) Nc , Y c ∈ (Y) Nc be the inputs and corresponding outputs of D c , and let X t ∈ (X ) Nt , Y t ∈ (Y) Nt 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.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: ω ML = arg max ω L ML (ω) where L ML (ω) = E p(ξ) Nt n=1 log p ω (y t,n |x t,n , D c ) .(1)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] .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) (f g)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) = (f g)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.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) dP . π ′ P (D) is the Radon-Nikodym 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 (gD) • g = π ′ P (D). Proposition 1 (G-stationarity and G-equivariance). The ground-truth stochastic process P is Gstationary 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: gD ∈ S consists of the input-output pairs (gx 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: Φ(gD) = 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Φ(D) = ρ(f enc (D)) where f enc ((x 1 , y 1 ), . . . , (x m , y m )) = m i=1 ϕ(y i )ψ( • , x i ) (2)for some continuous ϕ : Y → R 2D , an appropriate G-invariant positive-definite kernel ψ : X2 → R (i.e. ψ(x n , x m ) = ψ(gx n , gx 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. 
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, forf 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 2n )where E n is G-equivariant and t 1 , . . . , t 2n 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+2n → H, n ≥ 1, and a sequence of functionsEQUATION)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 = 2n 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 T P 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 T P 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 T P n ∥ → 0, and that T is (c, α)-Hölder for c, α > 0, in the following sense:∥T (u) -T (v)∥ ≤ c∥u -v∥ α for all u, v ∈ H.(5) 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+kn → H, n ≥ 1, and a sequence of functionsEQUATIONand sup u∈H:∥u∥≤MEQUATION)If assumption ( * ) does not hold, then the conclusion holds with E n ( • , t 1 , . . . , t kn ) replaced by E n (P n • , t 1 , . . . , t kn ). 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 i2πkx . Then e k (x -τ ) = e i2π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 T P 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 ∥T gf -gT f ∥ as a measure of how equivariant a mapping is [Wang et al., 2022a] .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 ρ:EQUATIONwhere 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 Gequivariance 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. 
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. 
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] .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.Table 1 : Average test log-likelihoods (↑) for the synthetic 1-D GP and 2-D smoke experiments. For the 1-D dataset we used the regular TNP, while for the 2-D experiment we used the PT-TNP. Results are grouped together by model class. Best in-class result is bolded. 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.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 .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. 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. • 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. 
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. Proof. Assume that ∥T -P n T P 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 boundEQUATIONWe 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,EQUATIONwhich proves that ∥T -P n T P n ∥ → 0.Theorem 3 (Approximation of non-equivariant operators.). Let T : H → H be a continuous, possibly nonlinear operator. Assume that ∥T -P n T P n ∥ → 0, and that T is (c, α)-Hölder for c, α > 0, in the following sense:EQUATIONMoreover, 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,EQUATION)If assumption ( * ) does not hold, then the conclusion holds with E n ( • , t 1 , . . . , t kn ) replaced by E n (P n • , t 1 , . . . , t kn ). We provide a proof in Appendix A.Proof. Let ε > 0. Choose n ∈ N such that ∥T -P n T P n ∥ < 1 2 ε, and choose h > 0 such thatch α = 1 2 ε. Set h n = h/ √ 2n.Consider the following collection of vectors:EQUATION)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:E : H × H 2|A| → H, E(u, a 1 , t 1 , . . . , a |A| , t |A| ) = |A| i=1 k(u, a i )t i |A| i=1 k(u, a i ) (12)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, byEQUATION)Now set t i = (P n T P 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 ,EQUATIONEQUATION(ii)≤ |A| i=1 k(P n u, a i )c∥P n u -a i ∥ α |A| i=1 k(P n u, a i ) (iii) ≤ ch α = 1 2 ε (17)where (i) uses the triangle inequality; (ii) (P n T P n )(u) = (P n T P n )(P n u) and (c, α)-Hölder continuity of P n T P 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 ,∥T (u) -E(P n u, . . . )∥ ≤ ∥T (u) -(P n T P n )(u)∥ + ∥(P n T P n )(u) -E(P n u, . . . )∥ ≤ ε. (18)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 , • ⟩, soEQUATIONTherefore, in the case that condition ( * ) holds, P n can be absorbed into the definition of E n .From the proof, we find thatk n = 2|A| = 2(1 + 2⌈ √ 2nM/h⌉) n ≤ 2(2 + 2 √ 2nM/h) n and ch α = 1 2 ε, which implies that h = (ε/2c) 1/α , so k n ≤ 2(2 + 2 √ 2nM (ε/2c) -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 .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. 
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. 
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). 
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). 
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). 
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] 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. • 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. 
 
 Appendix: 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.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 ẽ(D) : X → R Dz , where X denotes the discretised input domain. ρ implemented as a CNN through which ẽ(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 Con-vCNP (T ) for off-the-grid data. ; hi ← Nc n=1 [1, y T c,n ] T ψ(xi -xc,n); h (1) i ← h (1) i /h (0) i ; hi ← MLP(hi); {f (xi)} T i=1 ← CNN {xi, hi} T i=1 ; θm ← T i=1 f (xi)ψp(xt,m -xi); return p(•|θm); endAlgorithm 2: Forward pass through the Con-vCNP ( T ) for off-the-grid data. ; hi ← Nc n=1 [1, y T c,n ] T ψ(xi -xc,n); h (1) i ← h (1) i /h (0) i ; hi ← MLP(hi) + t( xi); {f (xi)} T i=1 ← CNN {xi, hi} T i=1 ; θm ← T i=1 f (xi)ψp(xt,m -xi); return p(•|θm); endTo achieve approximate translation equivariance, we construct a learnable function t : X → R Dz 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(ẽ(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.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 Eequivariance 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 Dz , 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.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 byEQUATION)whereEQUATIONHere, µ 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 3: Forward pass through the TE-TNP (T ). 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.Input: ρ = {TE-MHSA (ℓ) , TE-MHSA (ℓ) } L ℓ=1 , ϕ. Context Dc = (Xc, Yc) = {(xc,n, yc,n)} NcThe 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 byEQUATIONwhereEQUATIONandEQUATION)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 PT-TE-TNP (T ).Input: ρ = {TE-MHSA (ℓ) , TE-MHCA (ℓ) 1 , TE-MHCA (ℓ) 2 } L ℓ=1 , ϕ, V = {vm,0} M j=1 , U = {uj,0} M m=1 , z0. Context Dc = (Xc, Yc) = {(xc,n, yc,n)} Nc n=1 , and target Xt = {xt,m} N t m=1 begin zc,n ← ϕ(yc,n); vm ← vm,0 + 1/N Nc i=1 xc,i; uj ← uj,0; zt,m ← z0; for ℓ = 1, . . . , L do uj ← TE-MHCA (ℓ) 1 (uj, Zc, vj, Xc); {uj} M j=1 ← TE-MHSA (ℓ) (U, V); zt,m ← TE-MHCA (ℓ)2 (zt,m, U, zt,m, V); end θm ← MLP(zt,m); return p(•|θm); end Algorithm 6: Forward pass through the PT-TE-TNP ( T ).Input: ρ = {TE-MHSA (ℓ) , TE-MHCA (ℓ) 1 , TE-MHCA (ℓ) 2 } L ℓ=1 , ϕ, V = {vm,0} M j=1 , U = {uj,0} M m=1 , z0. Context Dc = (Xc, Yc) = {(xc,n, yc,n)} Nc n=1 , and target Xt = {xt,m} N t m=1 begin zc,n ← ϕ(yc,n); vm ← vm,0 + 1/N Nc i=1 xc,i; uj ← uj,0 + t(vm); zt,m ← z0; for ℓ = 1, . . . , L do uj ← TE-MHCA (ℓ) 1 (uj, Zc, vj, Xc); {uj} M j=1 ← TE-MHSA (ℓ) (U, V); zt,m ← TE-MHCA (ℓ) 2 (zt,m, U, zt,m, V); end θm ← MLP(zt,m); return p(•|θm); endWe 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.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 asEQUATION)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 asEQUATIONso 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.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 asEQUATION)We begin with the approach of Wang et al. [2022a] , which defines the relaxed group convolution asEQUATIONwhere µ(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(k * f )(u) = L l=1 G f (v)w l (v)k l (u -1 v)dµ(v) E l (f,w l )(u) = L l=1 E l (f, w l ) = E(f, w 1 , . . . , w L )(u) (29)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 givesE l (g • f, g • w l )(u) = G f (g -1 v)w l (g -1 v)k l (u -1 v)dµ(v) = G f (v)w l (v)k l (u -1 gv)dµ(v) = G f (v)w l (v)k l ((g -1 u) -1 v)dµ(v) = g • E l (f, w l )(u). (30) 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:EQUATIONDefine a fixed input t(u) = u. We can express the above convolution as a G-equivariant operation on t and f :EQUATION)To demonstrate G-equivariance, consider applying a transformation g ∈ G to the inputs:EQUATIONThus, this method is also a special case of ours.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 . LetEQUATION)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:EQUATION)where d gπ E-NP (D), π E-NP (gD) = 0 because π E-NP is strictly equivariant.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):k(x, x ′ ; l) = 2l(x)l(x ′ ) l(x) 2 + l(x ′ ) 2 exp - (x -x ′ ) 2 l(x) 2 + l(x ′ ) 2We 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 Dx → 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.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.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.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.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)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.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.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. For the approximately equivariant models, we plot both the model prediction (blue), as well as the predictions obtained without using the fixed inputs, which results in a strictly equivariant model (red). Both the strictly and approximately equivariant models output predictions that closely resemble the ground truth, but the non-equivariant TNP model completely fails to generalise.Table 3 : Average log-likelihoods (↑) for the synthetic 1-D GP experiment when tested on context sets. Ground truth log-likelihood is 0.2806 ± 0.0005. 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.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∆ equiv = E µ equiv -µ approx-equiv µ equiv p ,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 Relaxed-ConvCNP ( 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).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. 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.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 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.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.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.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 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. (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.We note that the form of the embedded dataset fenc(D) differs slightly in the RCP and TE-TNP. However, in both cases the form in Equation 2 can be recovered as special cases. We discuss this more in Appendix B.For example, the rectangle covering the maximum longitude and latitude of geographical locations seen during training of an environmental model.An implementation of our models can be found at cambridge-mlg/aenp.A longitude / latitude range of [35 • ,• ] / [10 • , 45 • ] and [-120 • , -80 • ] / [30 • , 50 • ], respectively. 


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

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

Reviews: review: paper_summary: paper_summary: The work provides an alternative approach to obtain a loose equivariance constraint. The authors established an interesting relationship between equivariant and non-equivariant models, showing that equivariant models with enough fixed input can approximate any non-equivariant model. Subsequently, the authors suggest using additional fixed input in neural processes, which is a common practice in transformers and NeRFs. The proposed method achieved better performance on multiple benchmark datasets.
strength_weakness: strength_weakness: The work provides a nicer reinterpretation of the existing practice of adding fixed/learnable positional embedding. Despite this reinterpretation, the work does not quantify the degree to which these additional fixed inputs hurt the equivariance property of the equivariant models.

Nor does it provide a scheme to automatically learn the number of fixed inputs for controlling the equivariance restriction. 

Thus, the work indeed proposes an approximately equivariant model, however, with unknown approximation.

This severely limits the contribution of the work.

Line 91: different notation should be used to denote the group action on $X$ and $Y$ 

Line 102: the symbol $f$ is already used in line 87 to denote the group element. The use of different symbols will fascinate reading.

Line 115: the notation of the encoder matches with a notation of the identity element of the group.
questions: questions: Line 103: how is the action of the group on the dataset $D$ defined?

Equation 2, line 111: How is the G-invariance of the kernel defined?

Line 143: $E_n$ is defined to take $n$ additional input; however, line 146 is defined to map on $H^{1+2n}$, which I think means $2n$ additional input. An explanation would be appreciated.
limitations: limitations: yes
paper_summary: paper_summary: This work considers approximately equivariant models --- which may better model or learn real-world tasks than exactly equivariant models --- especially in the context of neural processes. A new approximately equivariant method is developed, which uses an exactly equivariant model along with fixed inputs that break equivariance. This method can be used to parameterize approximately equivariant neural processes, by modifying existing approaches in simple ways. Experiments on 1-D regression, smoke plumes, and environmental data show benefits of the approach.
strength_weakness: strength_weakness: 1. Notation and definitions are heavy, which is somewhat understandable given the subject matter, but I do think it can be improved. For instance, compact operators are not defined.
2. 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).
questions: questions: 1. In equation (3), I think think the number of inputs should be larger than $n$, right?
2. Perhaps I missed this, but do you initialize the models to be equivariant? Via setting $t$ to be the zero function?
3. 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.

References:
* Wang et al. 2022a. Approximately Equivariant Networks for Imperfectly Symmetric Dynamics
* Dosovitskiy et al. 2021. An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
* Lim et al. 2023. Positional Encodings as Group Representations: A Unified Framework
* Xie et al. 2024. Equivariant Symmetry Breaking Sets
* Kim et al. 2023. Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance
limitations: limitations: Some discussion in conclusion
paper_summary: paper_summary: The paper describes a new framework for soft/approximate equivariance, based on the functional analysis of compact on hilbert spaces, assuming unitary group actions.

This is then applied to equivariant neural processes.
strength_weakness: strength_weakness: Edit after rebuttal: I think that the method is a good addition to the literature, and my concerns 1-3 below have been adequately answered in the rebuttal and "global rebuttal" parts.

-------------------

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

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

1) The assumption of having tasks with underlying compact operators seems to be swept under the carpet, and not properly discussed. 
In 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.

2) 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.

3) the parameter $k_n$ from Theorem 3 is not well behaved, so it is not clear if this theorem is useful in practice.
questions: questions: 0) see the un-numbered remark at the beginning of "weaknesses" section, about relevance to "neural processes" to the presented methods.

1) 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.

2) line 69-70 "represented with stochastic processes" -- is it not "represented by"? (this is a genuine question, not a rethorical one)

3) 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.

4) 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..

5) 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?

6) 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!)

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

8) 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.
9) same mention of "compact" should be inserted also in paragraph lines 249-261
10) same mention of "compact" operators should be inserted in seciton 6 (conclusions section)
limitations: limitations: these are addressed in the "weaknesses" section.

I 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.
paper_summary: paper_summary: This paper implemented the approximately equivariant neural process (NP) by relaxing the equivariance of the NP decoder and the relaxation is conducted by adding several leanable parameters as fixed inputs fed with the data embeddings into the decoder.
strength_weakness: strength_weakness: 1. No supplemental material for a source code.
1. A concept figure will be helpful to understand the intuition of the paper.
1. Absent analysis for equivariance error of the trained models. Did the trained model really reflect the approximate equivariance inherited in the dataset?
1. 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.
1. 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.
1. 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.
questions: questions: 1. Encoder should be equivariant? If so, doesn’t it lose non-equivariant information from the dataset D?
1. 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.
1. When the decoder is ViT, how many model parameters are additionally necessary compared to the strictly equivariant case?
1. It seems that the suggested relaxation is applicable to not only NP. Why did the paper choose NP to prove the relaxation idea?
limitations: limitations: There is no limitation of the contribution itself but experimental limitations are mentioned in the weakness.

Comments: title: Awaiting your response
No comment notification
Response to the Authors
Answer Part 2
Answer Part 1
Rebuttal
about the compactness + need more details for question 7
comment: Dear dBf7,

This is a gentle reminder that we are still awaiting your response. Please note that the discussion period is ending in a few hours.

In case it helps, we would like to clarify a technical point that we have left implicit. By showing that the approx. equiv. predictions are approximately equal to equivariant predictions, we show that `d(approx-np(D), equiv-np(D))` is small for all `D` (where `d` is some metric). Then, by the triangle inequality, we find that `d(T approx-np(D), approx-np(T D))` is small for all `T` and `D`:

```
    d(T approx-np(D), approx-np(T D))
        <= d(T approx-np(D), T equiv-np(D)) + d(T equiv-np(D), equiv-np(T D))
            + d(equiv-np(T D)), approx-np(T D))
```

where the first and third terms are small because `d(approx-np(D), equiv-np(D))` is small for all `D`, and the second term is zero because `equiv-np` is exactly equivariant. Therefore, the metric which we computed exactly implies equivariance in an approximate sense.
Here are the additional results for the environmental data.

**Context log-likelihoods environmental data**

Table 1: Average context loglikelihoods for the 2-D environmental regression experiment. Results are grouped together by model class.
_________________________________
Model | Europe | US
__________________
PT-TNP | 1.74 $\pm$ 0.01 | - |

PT-TNP ($T$) | 1.66 $\pm$ 0.01 | 1.28 $\pm$ 0.01

PT-TNP ($\widetilde{T}$) | 1.76 $\pm$ 0.01 | 1.47 $\pm$ 0.01
______________________________________
ConvCNP ($T$) | 1.20 $\pm$ 0.02 | 0.34 $\pm$ 0.02

ConvCNP ($\widetilde{T}$) | 1.50 $\pm$ 0.01 | 0.97 $\pm$ 0.01

RelaxedConvCNP ($\widetilde{T}$) | 1.29 $\pm$ 0.01 | 0.86 $\pm$ 0.01
__________________
EquivCNP ($E$) | 2.03 $\pm$ 0.01 | 1.76 $\pm$ 0.01

EquivCNP ($\widetilde{E}$) | 2.05 $\pm$ 0.01 | 1.69 $\pm$ 0.01

--------------------


Similar to the additional results for the synthetic 1-D regression experiments, these results show that all models achieve higher log-likelihood on the context data, demonstrating that our models are not “underfitting” to the context data.

**Equivariance error environmental data**

Table 2: Equivariance error on the 2-D environmental regression experiment. Note that all equivariance errors for the US are zero, as the additional fixed inputs are set to 0.
_________________
Model | Europe
_____________________
PT-TNP ($\widetilde{T}$) | 0.0406 $\pm$ 0.0005

ConvCNP ($\widetilde{T}$) | 0.0237 $\pm$ 0.0004

RelaxedConvCNP ($T$) | 0.0239 $\pm$ 0.0004

EquivCNP ($\widetilde{E}$) | 0.0242 $\pm$ 0.0005
_________________

These vary between 2-4%, again indicating that the model’s predictions only deviate slightly from the equivariant predictions.

We hope that our previous comment, together with these additional results, have adequately addressed all of your concerns. If so, we would encourage you to reconsider your score.
We have not received any notification after posting the previous two comments. With this comment, we just wanted to make sure that the reviewer gets notified about the new comments.
**Equivariance Error**

While you acknowledge that we (1) show that the approximately equivariant models outperform the equivariant models (with added inputs set to zero) and (2) approximately equivariant model outperform the non-equivariant models too, we understand that your main concerns are that we (a) do not show that our model is equivariant in an approximate sense and (b) do not show that the model’s approximate equivariance reflects the “actual approximate equivariance”. We address these points in order.

(a): Figures 1 and 2 show that the predictive means of the approximately equivariant model are equal to the predictive mean of the corresponding equivariant model plus a “small perturbation”. In other words, visually, the approximate equivariant predictive mean only slightly deviates from the corresponding equivariant predictive mean. Therefore, given that these figures are illustrative of the general behaviour of the models, we believe that we have clearly qualitatively shown that the predictions of the approximately equivariant models are indeed equivariant in an approximate sense.

To also argue this quantitatively, we have computed equiv_deviation = $\mathbb{E}$[norm(equiv_mean - approx_equiv_mean) / norm(equiv_mean)], which quantifies exactly how much on average the predictive means deviate from the corresponding equivariant predictive mean. The results are as follows:

Table 2. Equivariance error for the GP experiment.
| TNP ($\widetilde{T})$     | ConvCNP ($\widetilde{T}$)        | RelaxedConvCNP ($\widetilde{T}$) | EquivCNP ($\widetilde{T}$) |
| --------------------- | ------------------------- | -------------------------------- | -------------------------- |
| 0.0896 $\pm$ 0.0011                        | 0.0823 $\pm$ 0.0005                                | 0.1460 $\pm$ 0.0006                                              | 0.0825 $\pm$ 0.0006                                  |
|                                            |                                                    |                                                                  |                                                      |



These percentages are around 8-9% (only the RelaxedConvCNP shows 15% difference, which is based on the approach from Wang et al. [2022a], rather than our approach), which means that, on average, the models’ predictions deviate only slightly from the equivariant predictions.

(b): Whether the models learn the “actual approximate equivariance” is hard to determine. For example, in the GP experiments, what would the actual approximate equivariance be? In addition, what would the actual approximate equivariance be in the climate experiments?

While this is a hard question, we agree that it is an important question, which is why we attempted to answer this question in the following way: the models learn the “actual approximate equivariance” if the perturbation w.r.t. the corresponding equivariant prediction is “consistent with the structure of the problem”. In the GP experiments, we show that the perturbation exactly corresponds with the length-scale change. In the climate experiments, we show that the perturbation exactly corresponds with the key missing information that primarily breaks stationarity of the weather: orography. Therefore, we believe that we have reasonably demonstrated that the models approximate the “actual approximate equivariance”.

By having provided the log-likelihoods for the reconstruction of the context data and having provided numerical evidence that the models' predictions are indeed equivariant in an approximate sense, we hope to have addressed your concerns and consequently hope that you will reconsider your score.
Thank you for your reply and we hope to address your remaining concerns by providing the metrics you requested.

**Context Log-Likelihood**

We computed the log-likelihoods for reconstructing the context data in the GP experiment, and we will provide the results on the environmental data by the end of the discussion period. Please see the results below, which we will include in the supplementary material.

Table 1. Context log-likelihood for the GP experiment. Ground truth log-likelihood is $0.2806 \pm 0.0005$.
| TNP                            | TNP ($T$)                            | TNP ($\widetilde{T}$) |
| ------------------- | ------------------------------ | ------------------------------------ |
| 0.2296 $\pm$ 0.0007                    | 0.2396 $\pm$ 0.0013                                          | 0.2344 $\pm$ 0.0020                                                      |
| **ConvCNP ($T$)**                      | **ConvCNP ($\widetilde{T}$)**                                | **RelaxedConvCNP ($\widetilde{T}$)**                                     |
| 0.2362 $\pm$ 0.0007                    | 0.2218 $\pm$ 0.0009                                          | 0.2381 $\pm$ 0.0008                                                      |
| **EquivCNP ($E$)**                     | **EquivCNP ($\widetilde{E}$)**                               |                                                                          |
| 0.2213 $\pm$ 0.0007                    | 0.1992 $\pm$ 0.0010                                          |                                                                          |
|                                        |                                                              |                                                                          |

Note that these log-likelihoods are much higher than target log-likelihoods reported in the main body, because the models can accurately reconstruct the context data. Moreover, they are close to the ground truth log-likelihood ($0.2806 \pm 0.0005$), indicating that underfitting is unlikely.

**“However, high uncertainty in the reconstruction of context indicates merely inferring the whole sequence from a pattern in the context, which is not a stochastic process. (...) using imputation models over neural processes.”**

A neural process is defined as a neural network architecture that maps a context set to a stochastic process. This stochastic process can then be evaluated at any target inputs to produce predictions for any target set. All models in our submission are of this form. Therefore, by construction, all neural process models derive their predictions from an underlying stochastic process, and this underlying stochastic process can always be queried to e.g. analyse statistics of the underlying time series.

The premise of the neural process framework is that learning to predict the target sets given the context sets makes this underlying learned stochastic process converge to the “true stochastic process”. The sampling distributions of the context and target data should be set up in a way that enables this convergence in the limit of infinite data.

**“Target performance is not easily manipulated by variance, but it can be artificially enhanced by sacrificing context performance, as seen in TNP.”**

If all possible target sets include all possible context sets too, which is usually the case, then we would argue that this is false: in this case, worse context performance directly implies worse target performance, because context sets are target sets. A model could make a trade-off where it specifically performs worse in context set reconstruction in favour of better performance in interpolation further away, but this is unlikely to happen as the objective usually weights all target sets equally. In our experience, we have not seen this happen, unless the neural process severely underfits (which does happen for some architectures, like the original CNP based on just deep sets). In our submission, we do not believe that any of the models severely underfit in any of the experiments, which can be verified by the visualised predictions.
Thank you for providing more information, but most of it seems to carry the same nuance as the first rebuttal.

**Context Log-Likelihood**

As you mentioned, the context log-likelihood is not a direct performance metric. However, high uncertainty in the reconstruction of context indicates merely inferring the whole sequence from a pattern in the context, which is not a stochastic process. Instead, it is more like regression using a neural network that takes the context as input and produces the target as output. For example, in a task like prediction from irregularly sampled time series, the reason we use a stochastic process rather than other frameworks is to analyze the statistics of the time series, including both the context and target data, especially since the context itself may contain observational noise. If the goal were simply to build a model that reconstructs the target well, I would prefer using imputation models over neural processes. Target performance is not easily manipulated by variance, but it can be artificially enhanced by sacrificing context performance, as seen in TNP.

**Equivariance Error**

I understand that you showed the approximately equivariant model outperformed the equivariant model (with added inputs set to zero). However, my point is more about comparing the fully non-equivariant model with the approximately equivariant model. The approximately equivariant model should outperform the fully non-equivariant model because it leverages the inductive bias. While you did make this comparison, it does not guarantee that your model is indeed an approximately **equivariant** model, as you did not demonstrate that it truly reflects the actual approximate equivariance. I think your paper focused more on how to create an approximately **equivariant** model, rather than just breaking equivariance to overcome the strict constraints of equivariant architectures and add flexibility. If your paper is solely about breaking equivariance to introduce flexibility, then the experiments I suggested might not be necessary.

To be honest, I do not understand why you do not report the log-likelihoods during rebuttal period, especially since you agree that that is also a diagnostic metric. In my experience, **measuring the context log-likelihood takes less than 10 minutes**.

I think the paper should be reviewed after including the context log-likelihoods and the equivariance errors.
Thanks for the response. I want to maintain my score.
**Equivariance error analysis.**

**“(...) it does not showed that how much it breaks and whether it can indeed utilize the inductive bias of the weakly broken symmetry or not.”**

In the experiments, we compare the approximately equivariant model to the strictly equivariant model obtained by setting the additional inputs to zero. This reveals exactly what the effect of the additional inputs is and in what way they break equivariance. In particular, the bottom row of Figure 1 shows predictions for both the approximately equivariant model and the equivariant model obtained by setting the additional inputs to zero. Note that the equivariant model unnecessarily inflates the uncertainty, because translation equivariance prevents the model from recognising a fixed $x$-location, so it cannot learn that a change-point always happens at $x=0$.


The same comparison is done in Figure 2 for the environmental data experiment, (d) versus (e) and (g) versus (h). In the same figure, (f) and (i) show how much the predictions change as a consequence of the additional inputs. This difference exactly correlates with the missing information, the orography in (c).


We believe that these comparisons show that the additional inputs break equivariance to an appropriate extent. For example, in the bottom row of Figure 1, note that the approximately equivariant and equivariant predictions are very similar before the length-scale change. Including the additional inputs only changes the predictions after the length-scale change. 

**“If you show such analysis, you can say that the reason why your method works well is not because the method fully break the inductive bias and just empower the expressivity of the neural network.”**

The additional basis functions do strictly increase the expressivity of the model, exactly because the basis functions break equivariance of the model. That is, if the basis functions are zero, then the model is fully equivariant; and if the basis functions are non-zero, then the model is not equivariant. Because everything is continuous, the transition between these two “modes” is smooth and depends on the magnitude of the basis functions, which the model automatically learns.

The main practical question that remains is to whether this indeed corresponds to useful, practical, and appropriate improvements. We believe that is clearly shown in Figures 1 and 2.

We agree that it would be very interesting to more rigorously quantify the degree to which equivariance is broken with a metric like in RPP. However, for the current submission, we believe to have demonstrated that additional inputs (a) have a principled theoretical motivation, (b) break equivariance in interpretable and appropriate ways (as argued above), and (c) give considerable performance improvements on real-world data (experimental results). Although such an analysis would be interesting, we believe to have shown that additional inputs are a practical method that works in the intended way and have therefore left it for future work.

Nevertheless, if you insist that a more rigorous quantification is essential, then we would be willing to compute a metric very similar to EquivError from RPP for Figures 1 and Figure 2. For example, for Figure 2, this would divide the norm for panel (f) by the average of the norms of panels (d) and (e). By the colour bars, we expect this value to be around 2%. We would include these numbers in the captions.
Thank you for taking the time to go through our rebuttal. We address your remaining concerns below.


**Log-likelihood of context data.**
The practical importance of neural processes is determined by how they are used in practice, after training. In the usual framework, after training, there will be an unseen context set $D^*_c$ and an unseen target set $D^*_t$, and practical utility of the neural process is determined by how well the model predicts $D^*_t$ given $D^*_c$, usually measured in terms of log-likelihood. The appropriate way of measuring this performance is to (a) hold out some context–target pairs from training; and (b), after training, to compute the log-likelihood of the held-out target sets given the context sets.


In this framework, the probabilities $p(D_c | D_c)$ can be reasonably used as a diagnostic, but these probabilities do not necessarily correlate with good performance on the task of interest (predicting $D^*_t$ given $D^*_c$). For example, consider a NP perfectly trained on data sampled from a GP with a Matern kernel, and suppose $(D^*_c, D^*_t)$ is now sampled from a GP with a periodic kernel. Then the NP will achieve near-optimal $p(D^*_c | D^*_c)$, because the NP can reconstruct data it is conditioned on; but the NP will achieve disastrous $p(D^*_t | D^*_c)$, because the NP has never interpolated periodic data before.


It is often the case that the sampling distributions of the context and target data are the same. In that case, a model that maximises $\log p(D_t | D_c)$ will be guaranteed to reconstruct the context data well, because a context set is like a target set.


**“A good likelihood only in the target data may still seem favorable because its variance easily covers the true target data”**

Because probability densities integrate to one, the likelihood will penalise both too high and too low variance. In fact, the optimal likelihood is achieved if and only if the variance is equal to the true variance and the mean equal to the true mean. The likelihood therefore cannot be “cheated” by artificially increasing the variance. To make this mathematically rigorous, see e.g. Prop 1 by Foong et al., 2020.


In summary, the right way to measure practical performance is by computing log-likelihoods of held-out target data given context data, and this metric cannot be “cheated” by artificially increasing the variance. Nevertheless, we agree that the log-likelihood for reconstructing the context set can be used as a diagnostic metric.


If you insist that this diagnostic metric is important, then we are willing to include it as an additional result in the supplementary material.

Finally, as a clarification regarding implementation - **“(...) we often set the variance of the output as 0.1+0.9⋅Softplus($\sigma$), where $\sigma$ is the decoder output for variance. However, the Transformer Neural Process (TNP) sets it as $\exp⁡(\sigma)$”**

Our implementation of the TNP parameterises the standard deviation of the normal distribution as $\sigma_{min} + \text{Softplus}(z)^{1/2}$, where $z$ is the output of the decoder for the variance and $\sigma_{min}$ a hyperparameter usually set at 0.1. This is mentioned in the Appendix lines 700-701: “The decoder parameterises the mean and pre-softplus variance [correction: standard deviation] of a Gaussian likelihood with heterogeneous noise.”

**References**

Foong, Andrew Y. K. et al. Meta-Learning Stationary Stochastic Process Prediction with Convolutional Neural Processes. NeurIPS 2020.
Thank you so much for the detailed clarification. I understood much of it, but I still have some concerns regarding the log-likelihood and the equivariance error.

**The log-likelihood of reconstruction for the context data is important to evaluate the practical performance of neural processes.**

For example, we often set the variance of the output as $0.1+0.9\cdot\text{Softplus}(\sigma)$ where $\sigma$ is the decoder output for variance. However, the Transformer Neural Process (TNP) sets it as $\exp(\sigma)$, which differed from its baselines. This setting results in high variance that can cover lots of the target data but remains high even for the context data, leading to significant underfitting in the context data. Indeed, TNP shows a log-likelihood in context reconstruction that is 10 times worse than its baselines. A good likelihood only in the target data may still seem favorable because its variance easily covers the true target data, but it simply maintain high uncertainty for all data, which is not practical as a meta-learner. For these reasons, many previous works also report the log-likelihood or error of context data reconstruction. Please refer to Figure 3 in ANP [1], Table 1 in BNP [2], and Table 1 in MPNP [3].

If you agree with my opinion, I hope you consider including the log-likelihood or error of context data reconstruction, which is not difficult to measure.

[1] Kim, H. et al. (2019), Attentive Neural Processes.

[2] Lee, J. et al. (2020), Bootstrapping Neural Processes.

[3] Lee, H. et al. (2023), Martingale Posterior Neural Processes.

**Equivariance error analysis, either in theory or experiment, is necessary.**

What I mean by equivariance error analysis is, for example, adopting a simple function like $y=\text{CosineSimilarity}(\boldsymbol{x}\_1,\boldsymbol{x}\_2)$ (rotation-invariant), then weakly breaking the symmetry, such as by using $y = \text{CosineSimilarity}(\boldsymbol{x}\_1, \boldsymbol{x}\_2+0.01\cdot\boldsymbol{1})$, and generating a synthetic dataset. Note that since we already know the true function and the actual equivariance error, we can compare the equivariance error learned from the synthetic dataset with the true error.

The paper theoretically showed that adding fixed learnable inputs breaks the equivariance but it does not showed that how much it breaks and whether it can indeed utilize the inductive bias of the weakly broken symmetry or not. If you show such analysis, you can say that the reason why your method works well is not because the method fully break the inductive bias and just empower the expressivity of the neural network. Please refer to Figure 3 in RPP [1] and Figure 3 in PER [2].

[1] Finzi, M. et al. (2021), Residual Pathway Priors for Soft Equivariance Constraints.

[2] Kim, H. et al. (2023), Regularizing Towards Soft Equivariance Under Mixed Symmetries.
We greatly appreciate your engagement with our work. Thank you again for taking the time to review the paper in detail—without doubt, your feedback has led to an improvement in the exposition of our research.
Thanks for the clarification. Given your willingness to clarify the "compactness" part, the above clarifications, and the part of the "global rebuttal" on $k_n$ behavior, that we didn't discuss, I'm more confident to raise my score to 8.
Many thanks for your reply. You raise an important point that Wang et al. should be cited when detailing the set up for this experiment---we shall modify the paper accordingly. 

Thank you again for your feedback, and please don't hesitate to ask any remaining questions should they arise.
No problem at all! 

**CNNs**

So, as mentioned, we use a regular CNN which consists of multiple convolutional layers, each of which takes as input a $C_{in}$ dimensional feature map $f: \mathcal{X} \rightarrow \mathbb{R}^{C_{in}}$, where $\mathcal{X}$ denotes a discretised input domain (i.e. grid) and convolves it with a kernel $\psi: \mathcal{X} \rightarrow \mathbb{R}^{C_{out}\times C_{in}}$ : $$(f \ast \psi)(x)= \Sigma_{x'\in \mathcal{X}} f(x')\psi(x - x').$$ 

here, $C_{in}$ and $C_{out}$ denote the number of input and output channels, respectively, and $\ast$ denotes the convolution operation.  In practice, the kernel has a finite receptive field, meaning that $\psi(x - x') = 0$ when $x - x'$ exceeds some value. After each convolutional layer there is a nonlinearity $\phi: \mathbb{R} \to \mathbb{R}$ which acts point-wise on the feature map value. We have omitted details such as strides, padding and bias, which are also used when implementing. Note that we provide precise details on our hyperparameters choices (e.g. number of layers, number of channels at each layer, receptive field) in the Appendix. 

When used in ConvCNPs, the input into the CNN is the output of the 'SetConv' encoder $e: \mathcal{S} \to \mathbb{H}$ which maps datasets $\mathcal{D} \in \mathcal{S}$ to feature maps $f: \mathcal{X} \to \mathbb{R}^{E} \in \mathbb{H}$ .

**Additional fixed inputs**

As we have described, there are many possible ways to include additional fixed inputs $t_i: \mathcal{X} \to \mathbb{R}^{E}$ . The simplest approach is to include them as additional channels into the first convolutional layer, which in practice is achieved by concatenating $f$ with the fixed inputs. The first convolutional layer is then $$(\operatorname{cat}(f, t_1, \ldots, t_B) \ast \psi)(x) = \sum_{x' \in \mathcal{X}} \operatorname{cat}(f, t_1, \ldots, t_B)(x') \psi(x - x').$$ Let $T_{\tau}f = f(\cdot - \tau)$. Note that  $\operatorname{cat}(T_{\tau}f, t_1, \ldots, t_B) \neq T_{\tau}\operatorname{cat}(f, t_1, \ldots, t_B)$ , hence translation equivariance is broken. An alternative approach, which we take in the paper, is to add the additional channels to $f$: $$((f + t_1 + \ldots + t_B) \ast \psi)(x) = \sum_{x' \in \mathcal{X}} (f + t_1 + \ldots + t_B)(x') \psi(x - x').$$ Again, as $T_{\tau}f + t_1 + \ldots + t_B \neq T_{\tau}(f + t_1 + \ldots + t_B)$, translation equivariance is broken. 

Remarkably, two of the most prominent approaches to achieving approximate equivariance [1, 2] can be understood as alternative ways of including additional fixed inputs into a CNN to break equivariance. We discuss this in more detail in Appendix C. We believe that this marks an important contribution of our work: to first understand approximate equivariance, to unify existing methods, and to use this to develop a simple and effective approach to building approximate equivariance in *any* operator.

**References**

[1]: Wang, Rui, Robin Walters, and Rose Yu. "Approximately equivariant networks for imperfectly symmetric dynamics." International Conference on Machine Learning. PMLR, 2022.

[2]: van der Ouderaa, Tycho, David W. Romero, and Mark van der Wilk. "Relaxing equivariance constraints with non-stationary continuous filters." Advances in Neural Information Processing Systems 35 (2022): 33818-33830.
Ok sorry, I had forgotten that I was the one to use the term "multilinear" in the first place. I was assuming that the CNN was linear in $t_1,t_2,\dots$, but actually you never said it was.

But anyway, can you give some details on the architecture of the CNN abbreviated at line 156 of the paper?
We greatly appreciate you taking the time to go through our rebuttal. Please don't hesitate to leave additional comments/questions. We are more than happy to answer them.

**Compactness**: we will update the paper with the above discussion on compactness.

**Multilinear CNNs**: With regards to multilinear CNNs: it seems as though there's potentially some communication error here. To be clear, we do not use, nor propose the use of, multilinear CNNs---we use regular (nonlinear) CNNs throughout. Admittedly, we hadn't heard of a "multilinear CNN" before it was mentioned in your review, and were hoping that you would be able to provide a reference. 

Our rebuttal assumed you implied the use of a single-layer CNN without any nonlinearities (which would indeed make it a multiple input CNN with linear operations on the input). The only reference we can find in the literature ([1]) takes this approach.

[1] Pinson, Hannah, Joeri Lenaerts, and Vincent Ginis. "Linear CNNs discover the statistical structure of the dataset using only the most dominant frequencies." International Conference on Machine Learning. PMLR, 2023.
Thank you for the reply, I'm slowly going through the things you wrote back, so this may not be my only comment (sorry). 

About compactness, if you write something like the reply you put, I'll be satisfied. I was worried only about the "sweeping under the rug" nothing else.

One part that i still do not understand is what are multilinear CNNs. 

a) is there a reference for multilinear CNNs in the literature?
b) can you give a formula/pseudocode/details for a special case that is a multilinear CNN but not a standard CNN please? Because I'm not fully sure that I understand.
We thank the authors for their rebuttal. The clarifications are useful for me.

On the Wang et al. [2022a] point, thanks for clarifying the differences. My point is that, given that you take inspiration from and use the same tools as the Wang et al. setup, you should note and cite this in the paper.

Everything else looks good to me, I maintain my score.