SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Lemma 4.4. Let  $\pi$  be a prediction map,  $\mathcal{T}$  a transformation group, and  $g$  a comparison function. If  $\pi$  is relational with respect to  $g$  and  $g$  is  $\mathcal{T}$ -invariant, then  $\pi$  is  $\mathcal{T}$ -equivariant.
From Lemma 4.4 and previous definitions, we derive the main result about equivariance of RCNPs.
Proposition 4.5. Let  $g$  be the comparison function used in a RCNP, and  $\mathcal{T}$  a group of transformations. If  $g$  is  $\mathcal{T}$ -invariant, the RCNP is  $\mathcal{T}$ -equivariant.
As useful examples, the difference comparison function  $g_{\mathrm{diff}}(\mathbf{x}, \mathbf{x}^{\star}) = \mathbf{x}^{\star} - \mathbf{x}$  is invariant to translations of the inputs, and the distance comparison function  $g_{\mathrm{dist}}(\mathbf{x}, \mathbf{x}^{\star}) = ||\mathbf{x}^{\star} - \mathbf{x}||_{2}$  is invariant to rigid transformations; thus yielding appropriately equivariant RCNPs.
4.2 RCNPs are context-preserving
The previous section demonstrates that any RCNP is  $\mathcal{T}$ -equivariant, for an appropriate choice of  $g$ . However, a trivial comparison function  $g(\mathbf{x},\mathbf{x}^{\prime}) = \mathbf{0}$  would also satisfy the requirements, yielding a trivial representation. We need to guarantee that, at least in principle, the encoding procedure removes only information required to induce  $\mathcal{T}$ -equivariance, and no more. A minimal request is that the context set is preserved in the prediction map representation  $\mathbf{r}$ , modulo equivariances.
Definition 4.6. A comparison function  $g$  is context-preserving with respect to a transformation group  $\mathcal{T}$  if for any context set  $(\mathbf{X},\mathbf{Y})$  and target set  $\mathbf{X}^{\star}$ , there is a submatrix  $\mathbf{Q}' \subseteq \mathbf{Q}$  of the matrix  $\mathbf{Q}_{mnn'} = (g(\mathbf{x}_n,\mathbf{x}_m^\star),g(\mathbf{x}_n,\mathbf{x}_{n'}),\mathbf{y}_n,\mathbf{y}_{n'})$ , a reconstruction function  $\gamma$ , and a transformation  $\tau \in \mathcal{T}$  such that  $\gamma (\mathbf{Q}') = (\tau \mathbf{X},\mathbf{Y})$ .
For example,  $g_{\mathrm{dist}}$  is context-preserving with respect to the group of rigid transformations. For any  $m$ ,  $\mathbf{Q}' = \mathbf{Q}_{m:n}$  is the set of pairwise distances between points, indexed by their output values. Reconstructing the positions of a set of points given their pairwise distances is known as the Euclidean distance geometry problem [38], which can be solved uniquely up to rigid transformations with traditional multidimensional scaling techniques [39]. Similarly,  $g_{\mathrm{diff}}$  is context-preserving with respect to translations. For any  $m$  and  $n'$ ,  $\mathbf{Q}' = \mathbf{Q}_{m:n'}$  can be projected to the vector  $((\mathbf{x}_1 - \mathbf{x}_m^\star, \mathbf{y}_1), \ldots, (\mathbf{x}_N - \mathbf{x}_m^\star, \mathbf{y}_N))$ , which is equal to  $(\tau_m \mathbf{X}, \mathbf{Y})$  with the translation  $\tau_m(\cdot) = \cdot - \mathbf{x}_m^\star$ .
Definition 4.7. For any  $(\mathbf{X},\mathbf{Y}),\mathbf{X}^{\star}$ , a family of functions  $h_\theta ((\mathbf{X},\mathbf{Y}),\mathbf{X}^\star)\to \mathbf{r}\in \mathbb{R}^{d_{\mathrm{rep}}}$  is context-preserving under a transformation group  $\mathcal{T}$  if there exists  $\pmb {\theta}\in \Theta$ $d_{\mathrm{rep}}\in \mathbb{N}$ , a reconstruction function  $\gamma$ , and a transformation  $\tau \in \mathcal{T}$  such that  $\gamma (h_{\theta}((\mathbf{X},\mathbf{Y}),\mathbf{X}^{\star}))\equiv \gamma (\mathbf{r}) = (\tau \mathbf{X},\mathbf{Y})$
Thus, an encoding is context-preserving if it is possible at least in principle to fully recover the context set from  $\mathbf{r}$ , implying that no relevant context is lost. This is indeed the case for RCNPs.
Proposition 4.8. Let  $\mathcal{T}$  be a transformation group and  $g$  the comparison function used in a FullRCNP. If  $g$  is context-preserving with respect to  $\mathcal{T}$ , then the representation function  $r$  of the FullRCNP is context-preserving with respect to  $\mathcal{T}$ .
Proposition 4.9. Let  $\mathcal{T}$  be the translation group and  $g_{\mathrm{diff}}$  the difference comparison function. The representation of the simple RCNP model with  $g_{\mathrm{diff}}$  is context-preserving with respect to  $\mathcal{T}$ .
Given the convenience of the RCNP compared to FullRCNPs, Proposition 4.9 shows that we can use simple RCNPs to incorporate translation-equivariance with no loss of information. However, the simple RCNP model is not context-preserving for other equivariances, for which we ought to use the FullRCNP. Our theoretical results are confirmed by the empirical validation in the next section.
In this section, we evaluate the proposed relational models on several tasks and compare their performance with other conditional neural process models. For this, we used publicly available reference implementations of the neuralprocesses software package [1, 3]. We detail our experimental approach in Appendix B, and we empirically analyze computational costs in Appendix C.
5.1 Synthetic Gaussian and non-Gaussian functions
We first provide a thorough comparison of our methods with other CNP models using a diverse array of Gaussian and non-Gaussian synthetic regression tasks. We consider tasks characterized by functions derived from (i) a range of GPs, where each GP is sampled using one of three different kernels (Exponentiated Quadratic (EQ), Matérn- $\frac{5}{2}$ , and Weakly-Periodic); (ii) a non-Gaussian sawtooth
process; (iii) a non-Gaussian mixture task. In the mixture task, the function is randomly selected from either one of the three aforementioned distinct GPs or the sawtooth process, each chosen with probability  $\frac{1}{4}$ . Apart from evaluating simple cases with  $d_x = \{1,2\}$ , we also expand our experiments to higher dimensions,  $d_x = \{3,5,10\}$ . In these higher-dimensional scenarios, applying ConvCNP and ConvGNP models is not considered feasible. We assess the performance of the models in two distinct ways. The first one, interpolation (INT), uses the data generated from a range identical to that employed during the training phase. The second one, out-of-input-distribution (OOID), uses data generated from a range that extends beyond the scope of the training data.
Results. We first compare our translation-equivariant ('stationary') versions of RCNP and RGNP with other baseline models from the CNP family (Table 1). Comprehensive results, including all five regression problems and five dimensions, are available in Appendix D. Firstly, relational encoding of the translational equivariance intrinsic to the task improves performance, as both RCNP and RGNP models surpass their CNP and GNP counterparts in terms of INT results. Furthermore, the OOID results demonstrate significant improvement of our models, as they can leverage translational-equivariance to generalize outside the training range. RCNPs and RGNPs are competitive with convolutional models (ConvCNP, ConvGNP) when applied to 1D data and continue performing well in higher dimension, whereas models in the ConvCNP family are inapplicable for  $d_x > 2$ .
Table 1: Comparison of interpolation (INT) and out-of-input-distribution (OOID) performance of our RCNP models and CNP baselines on synthetic regression tasks with varying input dimensions. We show mean and (standard deviation) across 10 runs with different seeds. "F" denotes failed attempts that yielded very bad results. Missing entries could not be run. Statistically significantly best results are bolded. Our methods (RCNP and RGNP) perform better than their CNP and GNP counterparts in terms of both INT and OOID, and scale to higher dimension compared to ConvCNP and ConvGNP.
INT	RCNP (sta)	Weakly-periodic KL divergence(↓)	Sawtooth log-likelihood(↑)	Mixture log-likelihood(↑)
dx=1	dx=3	dx=5	dx=1	dx=3	dx=5	dx=1	dx=3	dx=5
RGNP (sta)	0.24 (0.00)	0.28 (0.00)	0.31 (0.00)	3.03 (0.06)	0.85 (0.01)	0.44 (0.00)	0.20 (0.01)	-0.10 (0.00)	-0.31 (0.03)
ConvCNP	0.03 (0.00)	0.05 (0.00)	0.08 (0.00)	3.90 (0.09)	1.09 (0.01)	1.13 (0.05)	0.34 (0.03)	0.37 (0.01)	0.04 (0.02)
ConvGNP	0.21 (0.00)	-	-	3.64 (0.04)	-	-	0.38 (0.02)	-	-
CNP	0.01 (0.00)	-	-	3.94 (0.11)	-	-	0.49 (0.15)	-	-
GNP	0.31 (0.00)	0.39 (0.00)	0.42 (0.00)	2.25 (0.02)	0.36 (0.28)	-0.03 (0.10)	0.01 (0.01)	-0.57 (0.11)	-0.72 (0.08)
0.06 (0.00)	0.08 (0.01)	0.11 (0.01)	0.83 (0.04)	0.23 (0.13)	0.02 (0.05)	0.17 (0.01)	-0.17 (0.00)	-0.32 (0.00)
OOID	RCNP (sta)	0.24 (0.00)	0.28 (0.01)	0.31 (0.00)	3.04 (0.06)	0.85 (0.01)	0.44 (0.00)	0.20 (0.01)	-0.10 (0.00)	-0.31 (0.03)
RGNP (sta)	0.03 (0.00)	0.05 (0.01)	0.08 (0.00)	3.90 (0.10)	1.09 (0.01)	1.13 (0.05)	0.34 (0.03)	0.37 (0.01)	0.04 (0.02)
ConvCNP	0.21 (0.00)	-	-	3.64 (0.04)	-	-	0.38 (0.02)	-	-
ConvGNP	0.01 (0.00)	-	-	3.97 (0.08)	-	-	0.49 (0.15)	-	-
CNP	2.88 (0.91)	1.58 (0.50)	2.20 (0.81)	F	-0.37 (0.12)	-0.22 (0.03)	F	-2.55 (1.15)	-1.71 (0.55)
GNP	F	1.47 (0.27)	0.62 (0.04)	F	F	F	F	-0.67 (0.05)	-0.72 (0.03)
We further consider two GP tasks with isotropic EQ and Matérn- $\frac{5}{2}$  kernels (invariant to rigid transformations). Within this set of experiments, we include the FullRCNP and FullRGNP models, each equipped with the 'isotropic' distance comparison function. The results (Table 2) indicate that RCNPs and FullRCNPs consistently outperform CNPs across both tasks. Additionally, we notice that FullRCNPs exhibit better performance compared to RCNPs as the dimension increases. When  $d_{x} = 2$ , the performance of our RGNPs is on par with that of ConvGNPs, and achieves the best results in terms of both INT and OOID when  $d_{x} > 2$ , which again highlights the effectiveness of our models in handling high-dimensional tasks by leveraging existing equivariances.
5.2 Bayesian optimization
We explore the extent our proposed models can be used for a higher-dimensional meta-learning task, using Bayesian optimization (BayesOpt) as our application [12]. The neural processes, ours as well as the baselines, serve as surrogates to find the global minimum  $f_{\mathrm{min}} = f(\mathbf{x}_{\mathrm{min}})$  of a black-box function. For this task, we train the models by generating random functions from a GP kernel sampled from a set of base kernels—EQ, Matérn-  $\{\frac{1}{2}, \frac{3}{2}, \frac{5}{2}\}$  as well as their sums and products—with randomly sampled hyperparameters. By training on a large distribution over kernels, we aim to exploit the metalearning capabilities of neural processes. The trained CNP models are then used as surrogates to minimize the Hartmann function [40, p.185] in three and six dimensions, a common BayesOpt test function. We use the expected improvement acquisition function, which we can evaluate analytically. Specifics on the experimental setup and further evaluations can be found in Appendix E.
Table 2: Comparison of the interpolation (INT) and out-of-input-distribution (OOID) performance of our RCNP models with different CNP baselines on two GP synthetic regression tasks with isotropic kernels of varying input dimensions.
INT	EQ KL divergence(↓)	Matérn-5/2 KL divergence(↓)
dx=2	dx=3	dx=5	dx=2	dx=3	dx=5
RCNP (sta)	0.26 (0.00)	0.40 (0.01)	0.45 (0.00)	0.30 (0.00)	0.39 (0.00)	0.35 (0.00)
RGNP (sta)	0.03 (0.00)	0.05 (0.00)	0.11 (0.00)	0.03 (0.00)	0.05 (0.00)	0.11 (0.00)
FullRCNP (iso)	0.26 (0.00)	0.31 (0.00)	0.35 (0.00)	0.30 (0.00)	0.32 (0.00)	0.29 (0.00)
FullRGNP (iso)	0.08 (0.00)	0.14 (0.00)	0.25 (0.00)	0.09 (0.00)	0.16 (0.00)	0.21 (0.00)
ConvCNP	0.22 (0.00)	-	-	0.26 (0.00)	-	-
ConvGNP	0.01 (0.00)	-	-	0.01 (0.00)	-	-
CNP	0.33 (0.00)	0.44 (0.00)	0.57 (0.00)	0.39 (0.00)	0.46 (0.00)	0.47 (0.00)
GNP	0.05 (0.00)	0.09 (0.01)	0.19 (0.00)	0.07 (0.00)	0.11 (0.00)	0.19 (0.00)
RCNP (sta)	0.26 (0.00)	0.40 (0.01)	0.45 (0.00)	0.30 (0.00)	0.39 (0.00)	0.35 (0.00)
RGNP (sta)	0.03 (0.00)	0.05(0.00)	0.11 (0.00)	0.03 (0.00)	0.05 (0.00)	0.11 (0.00)
FullRCNP (iso)	0.26 (0.00)	0.31 (0.00)	0.35 (0.00)	0.30 (0. 00)	0.32 (0.00)	0.29 (0.00)
FullRGNP (iso)	0.08 (0.00)	0.14 (0.00)	0.25 (0.00)	0.09 (0.00)	0.16 (0.00)	0.21 (0.0 00)
ConvCNP	0.22 (0.00)	-	-	0.26 (0.00)	-	-
ConvGNP	0.01 (0.00)	-	-	0.01 (0.00)	-	-
CNP	4.54 (1.76)	3.30 (1.55)	1.22 (0.09)	6.75 (2.72)	1.75 (0.42)	0.93 (0.02)
GNP	2.25 (0.61)	2.54 (1.44)	0.74 (0.02)	1.86 (0.26)	1.23 (0.17)	0.62 (0.02)
Figure 2: Bayesian Optimization. Error during optimization of a 3D/6D Hartmann function (lower is better). RCNP/RGNP improve upon the baselines, approaching the GP performance.
Results. As shown in Figure 2, RCNPs and RGNPs are able to learn from the random functions and come close to the performance of a Gaussian Process (GP), the most common surrogate model for BayesOpt. Note that the GP is refit after every new observation, while the CNP models can condition on the new observation added to the context set, without any retraining. CNPs and GNPs struggle with the diversity provided by the random kernel function samples and fail at the task. In order to remain competitive, they need to be extended with an attention mechanism (ACNP, AGNP).
5.3 Lotka-Volterra model
Neural process models excel in the so-called sim-to-real task where a model is trained using simulated data and then applied to real-world contexts. Previous studies have demonstrated this capability by training neural process models with simulated data generated from the stochastic Lotka-Volterra predator-prey equations and evaluating them with the famous hare-lynx dataset [32]. We run this benchmark evaluation using the simulator and experimental setup proposed by [3]. Here the CNP models are trained with simulated data and evaluated with both simulated and real data; the learning tasks represented in the training and evaluation data include interpolation, forecasting, and reconstruction. The evaluation results presented in Table 3 indicate that the best model depends on the task type, but overall our proposed relational CNP models with translational equivariance perform comparably to their convolutional and attentive counterparts, showing that our simpler approach does not hamper performance on real data. Full results with baseline CNPs are provided in Appendix F.
5.4 Reaction-Diffusion model
The Reaction-Diffusion (RD) model is a large class of state-space models originating from chemistry [45] with several applications in medicine [13] and biology [22]. We consider here a reduced model representing the evolution of cancerous cells [13], which interact with healthy cells through the
Table 3: Normalized log-likelihood scores in the Lotka-Volterra experiments (higher is better). The mean and (standard deviation) reported for each model are calculated based on 10 training outcomes evaluated with the same simulated (S) and real (R) learning tasks. The tasks include interpolation (INT), forecasting (FOR), and reconstruction (REC). Statistically significantly (see Appendix B.1) best results are bolded. RCNP and RGNP models perform on par with convolutional and attentive baselines.
INT (S)	FOR (S)	REC (S)	INT (R)	FOR (R)	REC (R)
RCNP	-3.57 (0.02)	-4.85 (0.00)	-4.20 (0.01)	-4.24 (0.02)	-4.83 (0.03)	-4.55 (0.05)
RGNP	-3.51 (0.01)	-4.27 (0.00)	-3.76 (0.00)	-4.31 (0.06)	-4.47 (0.03)	-4.39 (0.11)
ConvCNP	-3.47 (0.01)	-4.85 (0.00)	-4.06 (0.00)	-4.21 (0.04)	-5.01 (0.02)	-4.75 (0.05)
ConvGNP	-3.46 (0.00)	-4.30 (0.00)	-3.67 (0.01)	-4.19 (0.02)	-4.61 (0.03)	-4.62 (0.11)
ACNP	-4.04 (0.06)	-4.87 (0.01)	-4.36 (0.03)	-4.18 (0.05)	-4.79 (0.03)	-4.48 (0.02)
AGNP	-4.12 (0.17)	-4.35 (0.09)	-4.05 (0.19)	-4.33 (0.15)	-4.48 (0.06)	-4.29 (0.10)
production of acid. These three quantities (healthy cells, cancerous cells, acid concentration) are defined on a discretized space-time grid  $(2 + 1$  dimensions,  $d_{x} = 3)$ . We assume we only observe the difference in number between healthy and cancerous cells, which makes this model a hidden Markov model. Using realistic parameters inspired by [13], we simulated  $4\cdot 10^{3}$  full trajectories, from which we subsample observations to generate training data for the models; another set of  $10^{3}$  trajectories is used for testing. More details can be found in Appendix G.
We propose two tasks: first, a completion task, where target points at time  $t$  are inferred through the context at different spatial locations at time  $t - 1$ ,  $t$  and  $t + 1$ ; secondly, a forecasting task, where the target at time  $t$  is inferred from the context at  $t - 1$  and  $t - 2$ . These tasks, along with the form of the equation describing the model, induce translation invariance in both space and time, which requires the models to incorporate translational equivariance for all three dimensions.
Results. We compare our translational-equivariant RCNP models to ACNP, CNP, and their GNP variants in Table 4. Comparison with ConvCNP is not feasible, as this problem is three-dimensional. Our methods outperform the others on both the completion and forecasting tasks, showing the advantage of leveraging translational equivariance in this complex spatio-temporal modeling problem.
Table 4: Normalized log-likelihood scores in the Reaction-Diffusion problem for both tasks (higher is better). Mean and (standard deviation) from 10 training runs evaluated on a separate test dataset.
RCNP	RGNP	ACNP	AGNP	CNP	GNP
Completion	0.22 (0.33)	1.38 (0.62)	0.17 (0.03)	0.20 (0.03)	0.10 (0.01)	0.13 (0.03)
Forecasting	0.07 (0.18)	-0.18 (0.58)	-0.65 (0.31)	-1.58 (1.05)	-0.51 (0.20)	-0.50 (0.30)
5.5 Additional experiments
As further empirical tests of our method, we study our technique in the context of autoregressive CNPs in Appendix H.1, present a proof-of-concept of incorporating rotational symmetry in Appendix H.2, and examine the performance of RCNPs on image regression in Appendix H.3.
This work builds upon the foundation laid by CNPs [10] and other members of the CNP family, covered at length in Section 2. A significant body of work has focused on incorporating equivariances into neural network models [33, 4, 23, 41]. The concept of equivariance has been explored in Convolutional Neural Networks (CNNs; [24]) with translational equivariance, and more generally in Group Equivariant CNNs [5], where rotations and reflections are also considered. Work on DeepSets laid out the conditions for permutation invariance and equivariance [48]. Set Transformers [25] extend
this approach with an attention mechanism to learn higher-order interaction terms among instances of a set. Our work focuses on incorporating equivariances into prediction maps, and specifically CNPs.
Prior work on incorporating equivariances into CNPs requires a regular discrete lattice of the input space for their convolutional operations [15, 30]. EquivCNPs [19] build on work by [7] which operates on irregular point clouds, but they still require a constructed lattice over the input space. SteerCNPs [18] generalize ConvCNPs to other equivariances, but still suffer from the same scaling issues. These methods are therefore in practice limited to low-dimensional (one to two equivariant dimensions), whereas our proposal does not suffer from this constraint.
Our approach is also related to metric-based meta-learning, such as Prototypical Networks [36] and Relation Networks [37]. These methods learn an embedding space where classification can be performed by computing distances to prototype representations. While effective for few-shot classification tasks, they may not be suitable for more complex tasks or those requiring uncertainty quantification. GSSM [46] aims to learn relational biases via a graph structure on the context set, while we directly build exact equivariances into the CNP architecture.
Kernel methods and Gaussian processes (GPs) have long addressed issues of equivariance by customizing kernel designs to encode specific equivariances [14]. For instance, stationary kernels are used to capture globally consistent patterns [34], with applications in many areas, notably Bayesian Optimization [12]. However, despite recent computational advances [9], kernel methods and GPs still struggle with high-dimensional, complex data, with open challenges in deep kernel learning [47] and amortized kernel learning (or metalearning) [26, 35], motivating our proposal of RCNPs.
RCNPs crucially rely on a comparison function  $g(\mathbf{x}, \mathbf{x}^{\prime})$  to encode equivariances. The comparison functions we described (e.g., for isotropy and translational equivariance) already represent a large class of useful equivariances. Notably, key contributions to the neural process literature focus only on translation equivariance (e.g., ConvCNP [15], ConvGNP [30], FullConvGNP [2]). Extending our method to other equivariances will require the construction of new comparison functions.
The main limitation of the RCNP class is its increased computational complexity in terms of context and target set sizes (respectively,  $N$  and  $M$ ). The FullRCNP model can be cumbersome, with  $O(N^2 M)$  cost for training and deployment. However, we showed that the simple or 'diagonal' RCNP variant can fully implement translational invariance with a  $O(NM)$  cost. Still, this cost is larger than  $O(N + M)$  of basic CNPs. Given the typical metalearning setting of small-data regime (small context sets), the increased complexity is often acceptable, outweighed by the large performance improvement obtained by leveraging available equivariances. This is shown in our empirical validation, in which RCNs almost always outperformed their CNP counterparts.
In this paper, we introduced Relational Conditional Neural Processes (RCNPs), a new member of the neural process family which incorporates equivariances through relational encoding of the context and target sets. Our method applies to equivariances that can be induced via an appropriate comparison function; here we focused on translational equivariances (induced by the difference comparison) and equivariances to rigid transformations (induced by the distance comparison). How to express other equivariances via our relational approach is an interesting direction for future work.
We demonstrated with both theoretical results and extensive empirical validation that our method successfully introduces equivariances in the CNP model class, performing comparably to the translational-equivariant ConvCNP models in low dimension, but with a simpler construction that allows RCNPs to scale to larger equivariant input dimensions ( $d_x > 2$ ) and outperform other CNP models.
In summary, we showed that the RCNP model class provides a simple and effective way to implement translational and other equivariances into the CNP model family. Exploiting equivariances intrinsic to a problem can significantly improve performance. Open problems remain in extending the current approach to other equivariances which are not expressible via a comparison function, and making the existing relational approach more efficient and scalable to larger context datasets.
This work was supported by the Academy of Finland Flagship programme: Finnish Center for Artificial Intelligence FCAI. Samuel Kaski was supported by the UKRI Turing AI World-Leading Researcher Fellowship, [EP/W002973/1]. The authors wish to thank the Finnish Computing Competence Infrastructure (FCCI), Aalto Science-IT project, and CSC-IT Center for Science, Finland, for the computational and data storage resources provided.