SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We simulate trajectories from Eq. S8 on a discretized space-time grid using the SpatialPy simulator [6]. For the purpose of our experiment, we selected realistic parameter ranges that would also produce quick evolution. The selected parameter ranges are therefore close to [13], but for example the diffusion rates were increased. The parameters for the simulations are sampled according to the following distributions:  $k_{u} = 10$ ,  $k_{v} = 100$ ,  $r_{u} \sim \mathcal{U}(0.0027, 0.0033)$ ,  $r_{v} \sim \mathcal{U}(0.027, 0.033)$ ,  $r_{w} \sim \mathcal{U}(1.8, 2.2)$ ,  $D_{v} \sim U(0.0009, 0.0011)$ ,  $D_{w} \sim \mathcal{U}(0.009, 0.011)$ ,  $d_{w} = 0$ ,  $d_{u} \sim \mathcal{U}(0.45, 0.55)$ . These parameters lead to a relatively slowly-growing number of cancerous cells, with a fast decay in healthy cells due to acid spread, as depicted in Figure S3. The process starts with the healthy cells count  $u$  constant across the whole space, cancerous cells count  $v$  zero except for one cancerous cell at a uniformly random position, and acid count  $w$  zero across the whole space.
Figure S3: Example simulation of the Reaction-Diffusion model at times  $t = 0,1,2$ . Each row represents a different quantity in the model (from top to bottom, healthy cells, cancerous cells, and acid). At time  $t = 0$ , the population of healthy cells is 10 at each point of space, and a single point is occupied by additional cancerous cells. After this time, cancerous cells start spreading and slowly multiplying, while acid spreads faster, killing healthy cells along its way.
We generated a total  $5 \cdot 10^{3}$  trajectories on a  $11 \times 11$  grid at times  $t = 0, \dots, 5$ . We used the first  $4 \cdot 10^{3}$  trajectories to generate datasets for model training and the other  $10^{3}$  to generate datasets for evaluation. The datasets were generated by sampling from the trajectories as follows. First, we extracted the population counts at three consecutive time points from a trajectory and time point selected at random. We then sampled 50-200 context points and 100-200 target points from this subsample to create a dataset representing either a completion or forecasting task. For the completion task, the context and target points were sampled from the subsample at random, and for the forecasting task, all context points were sampled from the first two time points, and all target points from the third time point were included in the subsample. All datasets included in the training and validation data represented the completion task.
Finally, the input and output variables at each context and target point were determined as follows. In the experiments reported in the main text, we considered the case of partial observations, where only the difference  $u - v$  is observed. This makes the simulator model a hidden Markov model and thus intractable. Predicting future states of the quantity  $u - v$  without running inference and without having access to  $w$  is a challenging problem, similar to several application cases of interest. In this appendix, we additionally consider a scenario where we include the acid count in the input variables and take the healthy cells count as the output variable.
G.3 Additional results
A promising scientific application of neural processes is inferring latent variables given other measured quantities. To explore this capability, we attempt to infer the quantity of healthy cells  $(u(\mathbf{z},t))$  given as input the acidity concentration  $(w(\mathbf{z},t))$  for some  $\mathbf{z}, t$ . For this study, we augmented the inputs to  $(\mathbf{z},t,w(\mathbf{z},t))$ . Notably, the regression problem is not equivariant in the last input dimension, so we encoded  $w(\mathbf{z},t)$  in a non-relational manner also in the RCNP and RGNP models. The results are presented in Table S10. The difference between the completion and forecasting tasks is small in this experiment, and the best results are observed with RGNP in both tasks. AGNP performance is surprisingly poor and will require further investigation.
Table S10: Normalized log-likelihood scores in the additional experiments of 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.35 (0.12)	0.74 (0.18)	0.24 (0.25)	-2.20 (3.67)	0.44 (0.17)	0.32 (0.26)
Forecasting	0.29 (0.11)	0.63 (0.15)	0.38 (0.21)	-1.7 (3.0)	0.46 (0.13)	0.41 (0.20)
G.4 Computation time
On GPU, for all methods the training runs of this experiment were particularly fast, each one lasting no more than 5 minutes. We carried out a total of 10 training runs for each method, with a total of 120 training runs. We also performed around 60 exploratory training runs (with longer training time, larger networks, etc.). The total training time comes to less than 15 hours.
In addition to the experimental evaluation in the main paper and their extended discussion in the sections above, we explored several further experimental setups to demonstrate the wide applicability of our proposal. These are (i) an application of the autoregressive structure proposed by [3] to our proposed approach (Section H.1), (ii) a demonstration of how to incorporate rotational equivariance into the model (Section H.2), and (iii) an exploration of the method's performance on two popular image data regression tasks (Section H.3).
H.1 Autoregressive CNPs
To further demonstrate the effectiveness of our RCNP architectures, we ran a set of experiments adopting an autoregressive approach to both our RCNP and other CNP models. Autoregressive
CNPs (AR-CNPs) are a new family of CNP models which were recently proposed by [3]. Unlike traditional CNPs, which generate predictions independently for each target point, or GNPs, which produce a joint multivariate normal distribution over the target set, AR-CNPs define a joint predictive distribution autoregressively, leveraging the chain rule of probability. Without changing the CNP architectures and the training procedure, AR-CNPs feed earlier output predictions back into the model autoregressively to predict new points. This adjustment enables the AR-CNPs to capture highly dependent, non-Gaussian predictive distributions, resulting in better performance compared to traditional CNPs.
In this experiment, we compare AR-RCNP, AR-CNP, and AR-ConvCNP on synthetic regression data. The AR approach is applied at the evaluation stage to generate the target predictions autoregressively. When applying AR, the training method for the models remains unchanged. Thus, for this experiment, we use the same models that were trained with synthetic data. Details about the data generation and model configurations can be found in Section D. We chose the translational-equivariant version of RCNP for each task.
The full results are reported in Table S11. As we maintain consistent experimental settings with previous synthetic experiments (Section D), we can directly compare the performance of models without the AR approach (Table S3-S7) to those enhanced with the AR approach. We observe that the AR approach improves RCNP performance in the Gaussian experiments in all dimensions and in the non-Gaussian sawtooth and mixture experiments when  $d_x > 1$ . When  $d_x = 1$ , AR-RCNP results are comparable to RCNP results in the mixture task and worse in the sawtooth task. We also observe that AR-CNP results are worse than CNP results in most sawtooth and mixture experiments. Overall we can see that AR-RCNPs are better than AR-CNPs across all dimensions and tasks under both INT and OOID conditions, and achieve close performance compared with AR-ConvCNPs under  $d_x = 1$  in the Gaussian experiments.
Table S11: Comparison of the interpolation (INT) and out-of-input-distribution (OOID) performance of AR-RCNP model with other AR baselines on various synthetic regression tasks with varying input dimensions. We show the mean and (standard deviation) obtained from 10 runs with different seeds. "F" denotes failed attempts that yielded very bad results. Our AR-RCNP method performs competitively in low dimension, and scales to higher input dimensions ( $d_x > 2$ ), where AR-ConvCNP is not applicable.
EQ KL divergence(↓)	Matérn-5/2 KL divergence(↓)	Weakly-periodic KL divergence(↓)	Sawtooth log-likelihood(↑)	Mixture log-likelihood(↑)
INT	OOID	INT	OOID	INT	OOID	INT	OOID	INT	OOID
1	AR-RCNP	0.02 (0.00)	0.02 (0.00)	0.02 (0.00)	0.02 (0.00)	0.04 (0.00)	0.04 (0.00)	1.93 (0.05)	1.97 (0.04)	0.20 (0.27)	0.20 (0.27)
AR-ConvCNP	0.01 (0.00)	0.01 (0.00)	0.00 (0.00)	0.00 (0.00)	0.01 (0.00)	0.01 (0.00)	4.11 (0.09)	4.11 (0.09)	0.83 (0.03)	0.82 (0.03)
AR-CNP	0.09 (0.01)	4.33 (3.44)	0.12 (0.01)	F	0.21 (0.01)	5.05 (4.52)	-1.77 (1.20)	F	-0.25 (0.09)	F
2	AR-RCNP	0.14 (0.01)	0.14 (0.01)	0.13 (0.00)	0.13 (0.00)	0.11 (0.00)	0.11 (0.00)	1.55 (0.01)	1.55 (0.01)	0.20 (0.00)	0.20 (0.00)
AR-CNP	0.29 (0.00)	3.05 (1.26)	0.32 (0.03)	2.24 (1.20)	0.33 (0.00)	2.03 (0.75)	0.25 (0.33)	-0.26 (0.09)	-0.39 (0.01)	-3.72 (1.68)
5	AR-RCNP	0.23 (0.00)	0.23 (0.00)	0.15 (0.00)	0.15 (0.00)	0.16 (0.00)	0.16 (0.00)	0.87 (0.01)	0.87 (0.01)	-0.10 (0.00)	-0.10 (0.00)
AR-CNP	0.49 (0.00)	1.22 (0.05)	0.39 (0.00)	0.97 (0.05)	0.38 (0.00)	1.47 (0.35)	-0.00 (0.13)	F	-0.66 (0.08)	-3.51 (2.32)
H.1.1 Computation time
The AR procedure is not used in model training, so the training procedure aligns with that of standard CNP models. The experiments reported in this section used AR with models trained for the synthetic regression experiments in Section 5.1. If the models had been trained from scratch, we estimate that the total computational cost for these experiments would have been approximately 2000 GPU-hours. AR is applied in the evaluation phase, where we predict target points autoregressively. This process is considerably lengthier than standard evaluation, with the duration influenced by both the target set size and data dimension. Nonetheless, since only a single pass is required during evaluation, the overall computational time remains comparable to standard CNP models.
H.2 Rotation equivariance
The experiments presented in this section use a two-dimensional toy example to investigate modeling datasets that are rotation equivariant but not translation equivariant. We created regression tasks based
on synthetic data sampled from a GP model with an exponentiated quadratic covariance function and mean function  $m(\mathbf{x}) = \|\mathbf{A}\mathbf{R}\mathbf{x}\|_2^2$ , where  $\mathbf{A} = \mathrm{diag}(\mathbf{a})$  is a fixed diagonal matrix with unequal entries and  $\mathbf{R}$  is a random rotation matrix. We considered an isotropic-kernel model version with the standard EQ covariance function (Equation S3) and an anisotropic-kernel model version with the covariance function  $k(\mathbf{x},\mathbf{x}') = \exp(-||\mathbf{B}\mathbf{R}\mathbf{x} - \mathbf{B}\mathbf{R}\mathbf{x}'||_2^2)$ , where  $\mathbf{B} = \mathrm{diag}(\mathbf{b})$  is a fixed diagonal matrix with unequal entries and  $\mathbf{R}$  is the same random rotation matrix as in the mean function. Both models use the anisotropic GP mean function defined previously.
We generated the datasets representing regression tasks by sampling context and target points from the synthetic data generated with isotropic-kernel or anisotropic-kernel model version as follows. The number of context points sampled varied uniformly between 1 and 70, while the number of target points was fixed at 70. All training and validation datasets were sampled from the range  $\mathcal{X} = [-4,0] \times [-4,0]$ . To evaluate the models in an interpolation (INT) task, we generated evaluation datasets by sampling context and target points from the training range. To evaluate the models in an out-of-input distribution (OOID) task, we generated evaluation datasets with context and target points sampled in the range  $[0,4] \times [0,4]$ .
We chose a comparison function  $g$  that is rotation invariant but not translation invariant defined as<sup>2</sup>
$$
g (\mathbf {x}, \mathbf {x} ^ {\prime}) = (\| \mathbf {x} - \mathbf {x} ^ {\prime} \| _ {2}, \| \mathbf {x} \| _ {2}, \| \mathbf {x} ^ {\prime} \| _ {2}).
$$
Since 'simple' RCNPs are only context preserving for translation-equivariance and not for other equivariances (Proposition 4.9), we combine the rotation-equivariant comparison function with a FullRCNP model.
We compared FullRCNP with the rotation-equivariant comparison function, FullRCNP (rot), with CNP and ConvCNP models in both the isotropic-kernel and anisotropic-kernel test condition. We also used isotropic-kernel data to train and evaluate a translation-equivariant ('stationary') RCNP version, RCNP (sta), that uses the difference comparison function. The model architectures and training details were the same as in the two-dimensional synthetic regression experiments (Section D.1).
The results reported in Table S12 show that, as expected, the isotropic-kernel test condition is easier than the anisotropic-kernel test condition where the random rotations change the covariance structure. The best results in both conditions are observed with ConvCNP when the models are evaluated in the interpolation task and with FullRCNP (rot) when the models are evaluated in the OOID task. Since the input range in the OOID task is rotated compared to the training range, FullRCNP (rot) results are on the same level in both tasks while the other models are not able to make reasonable predictions in the OOID task. We also observe that the FullRCNP (rot) results in the interpolation task are better than CNP or RCNP (sta) results in the isotropic-kernel test condition and similar to CNP in the anisotropic-kernel test condition.
Table S12: Normalized log-likelihood scores in the experiments using synthetic data generated with isotropic and anisotropic GP kernels and an anisotropic GP mean function. The mean and (standard deviation) reported for each model are calculated based on 10 training outcomes evaluated with separate datasets representing interpolation (INT) and out-of-input-distribution (OOID) tasks. "F" denotes failed attempts with log-likelihood scores below  $-60$ .
Isotropic GP kernel	Anisotropic GP kernel
INT	OOID	INT	OOID
FullRCNP (rot)	-0.89 (0.04)	-0.90 (0.04)	-6.37 (0.34)	-6.38 (0.31)
RCNP (sta)	-1.16 (0.03)	F	—	—
ConvCNP	-0.60 (0.02)	F	-2.83 (0.12)	F
CNP	-1.42 (0.03)	F	-6.43 (1.13)	F
H.2.1 Computation time
The models were trained on a single GPU, and each training run took around 6 hours with FullRCNP, 2 hours with CNP, 4 hours with ConvCNP and 3 hours with RCNP. The results provided in this
2This comparison function is invariant to proper and improper rotations around the origin (i.e., including mirroring).
appendix required 70 training runs and consumed around 250 GPU hours. Additional experiments related to this example consumed around 2500 GPU hours.
H.3 Image datasets
The experiments reported in the main paper demonstrate how RCNP can leverage prior information about equivariances in regression tasks including tasks with more than 1-2 input dimensions. In this section, we provide details and results from additional experiments that used MNIST [24] and CelebA [27] image data to create regression tasks. The neural process models were used in these experiments to predict the mean and variance over pixel values across the two-dimensional image given some pixels as observed context. The experiments investigate how RCNP compares to other conditional neural process models when the assumption about task equivariance is incorrect.
H.3.1 Setup
The experiments reported in this section compared RCNP, CNP, and ConvCNP models using interpolation tasks generated based on  $16 \times 16$  pixel images. The pixels were viewed as data points with  $d_x = 2$  input variables that indicate the pixel location on image canvas and  $d_y$  output variables that encode the pixel value. The output dimension depended on the image type. Pixels in grayscale images were viewed as data points with one output variable ( $d_y = 1$ ) that takes values between 0 and 1 while pixels in RGB images were viewed as data points with three output variables ( $d_y = 3$ ) that each take values between 0 and 1 to encode one RGB channel. We evaluated the selected models using standard MNIST and CelebA images (Section H.3.2) and examined modeling translation equivariance with toroidal MNIST images (Section H.3.3).
The model architectures used in the image experiments match the architectures used in the synthetic regression experiments (Section D.1) extended with an output transformation to bound the predicted mean values between 0 and 1. While we used images with multiple output channels, there was no need to use a multi-output architecture such as discussed in Section F.1. This is because we assumed that all output channels are observed when a pixel is included in the context set and unobserved otherwise, meaning that the output channels in image data are not associated with separate context sets. RCNP models were used with the standard stationary comparison function unless otherwise mentioned.
The models were trained for 200 epochs with  $2^{14}$  datasets in each epoch and learning rate  $3 \cdot 10^{-4}$ , and tested with validation and evaluation sets that included  $2^{12}$  datasets each. The log-likelihood score that is used as training objective is usually calculated based on the predicted mean and variance, but when the models were trained on MNIST data, we started training with the predicted variance fixed to a small value. This was done to prevent the neural process models from learning to predict the constant background values with low variance and ignoring the image details. We trained with the fixed variance between epochs 1-150 and with the predicted variance as usual between epochs 151-200.
The datasets used to train and evaluate the neural process models were generated as follows. Each dataset was generated based on an image that was sampled from the selected image set and downscaled to  $16 \times 16$  pixels. The downscaled images were either used as such or transformed with random translations as explained in Section H.3.3. The context sets generated based on the image data included 2-128 pixels sampled at random while the target sets included all 256 pixels. The evaluation sets used in each experiment were generated with the context set size fixed to 2, 20, 40, and 100 pixels.
H.3.2 Experiment 1: Centered images (no translation)
The experiments reported in this section used MNIST and CelebA images scaled to the size  $16 \times 16$  (Figure S4). The normalized log-likelihood scores calculated based on each evaluation set and across training runs for each neural process model are reported in Table S13. The best results are observed with ConvCNP, while the order between RCNP and CNP depends on the context set size. CNP results are better than RCNP results when predictions are generated based on  $N = 2$  context points while RCNP results are better when more context data is available.
We believe CNP works better than RCNP when  $N = 2$  because the absolute pixel locations are generally informative about the possible pixel values in the MNIST and CelebA images. This means
(a)
Figure S4: Example (a) MNIST and (b) CelebA  $16 \times 16$  images.
Table S13: Normalized log-likelihood scores in the image experiments using (a)  $16 \times 16$  MNIST images or (b)  $16 \times 16$  CelebA images (higher is better). The mean and (standard deviation) reported for each model are calculated based on 10 training outcomes evaluated in interpolation task with context sizes 2, 10, 20, and 100. RCNP results are the lowest when the context set is small but improve when the context size increases.
2	20	40	100
(a)	RCNP	5.52 (0.30)	7.13 (0.10)	8.08 (0.13)
ConvCNP	6.43 (0.02)	7.35 (0.04)	8.27 (0.04)
CNP	6.29 (0.12)	7.04 (0.03)	7.47 (0.03)
(b)	RCNP	0.18 (0.00)	1.19 (0.01)	1.79 (0.01)
ConvCNP	0.33 (0.00)	1.36 (0.01)	2.03 (0.01)
CNP	0.31 (0.00)	0.90 (0.01)	1.09 (0.01)
that the regression tasks generated based on these images are not translation equivariant and the assumption encoded in the RCNP comparison function is incorrect. Since RCNP does not encode the absolute target location, it needs context data to make predictions that CNP and ConvCNP can make based on the target location alone (Figure S5).
While RCNP needs context data to work around the incorrect assumption about translation equivariance, RCNP results are better than CNP results when more context data is available. We believe this is because the context set representation in CNP does not depend on the current target location while RCNP can learn to preserve context information that is relevant to the current target. For example, while CNP looses information about the exact observed values when the context set is encoded, RCNP generally learns to reproduce the observed context. RCNP may also learn to emphasize context points that are close to the current target and may learn local features that are translation equivariant.
H.3.3 Experiment 2: Translated images
The interpolation tasks in the previous image experiments do not exhibit translation equivariance since the input range is fixed and the image content is centered on the canvas. To run an image experiment with translation equivariance, we converted the  $16 \times 16$  MNIST images into toruses and applied random translations to move image content around. The translations do not change the relative differences between image pixels on the torus, but when the torus is viewed as an image, both the absolute pixel locations and relative differences between pixel locations are transformed. This means that to capture the translation equivariance, we needed to calculate differences between pixel locations on the torus.
The experiments reported in this section compare models trained using either centered image data (Figure S4 (a)) or image data with random translations on the torus (Figure S6). We ran experiments with both RCNP (sta) that calculates the difference between pixel locations on the image canvas and RCNP (tor) that calculates the difference between pixel locations on the torus. The models trained with centered image data were evaluated with both centered and translated data while the models
(a)
Figure S5: Example context sets and predicted mean values in the image experiment using (a) MNIST and (b) Celeba  $16 \times 16$  data.
Figure S6: MNIST  $16 \times 16$  example images with random translations on a torus.
trained with translated data were evaluated with translated data. The normalized log-likelihood scores are reported in Table S14.
Table S14: Normalized log-likelihood scores in the image experiments using  $16 \times 16$  MNIST images with random translations on a torus (higher is better). We compare (a) models trained on centered image data evaluated on centered image data, (b) models trained on centered image data evaluated on image data with random translations, and (c) models trained on image data with random translations evaluated on image data with random translations. The mean and (standard deviation) reported for each model are calculated based on 10 training outcomes evaluated in interpolation task with context sizes 2, 10, 20, and 100. "F" denotes very bad results with negative log-likelihood scores. RCNP (tor) trained on centered image data is able to make predictions about images with random translations because the translation mechanism is modeled in the comparison function.
2	20	40	100
(a)	RCNP (sta)	5.52 (0.30)	7.13 (0.10)	8.08 (0.13)
RCNP (tor)	5.51 (0.01)	6.63 (0.04)	7.59 (0.03)
ConvCNP	6.43 (0.02)	7.35 (0.04)	8.27 (0.04)
CNP	6.29 (0.12)	7.04 (0.03)	7.47 (0.03)
(b)	RCNP (sta)	F	F	F
RCNP (tor)	5.62 (0.01)	6.68 (0.04)	7.61 (0.06)
ConvCNP	F	F	F
CNP	F	F	F
(c)	RCNP (sta)	5.51 (0.21)	6.66 (0.11)	7.67 (0.09)
RCNP (tor)	5.56 (0.10)	6.72 (0.06)	7.70 (0.09)
ConvCNP	5.61 (0.01)	6.72 (0.03)	7.78 (0.04)
CNP	5.55 (0.03)	6.29 (0.03)	6.80 (0.02)