Our Further Clarifications to Major Misunderstandings
Thank you for your detailed response, to which we respond below.
**1a. Set Transformer (Major Misunderstanding)**
We apologize for a truly unfortunate typo in our rebuttal (not our paper), which likely emerged while trimming our response. We meant to say "We use **Deep Sets** to achieve invariance to permutations of the context set (see line 128)." The citation in line 128 is the DeepSets paper. **We do not use Set Transformer.** This can be verified by checking our code submitted in the Supplement.
Set Transformers differ from DeepSets in that they add an attention mechanism which *learns* interactions between elements of the set. **Our method does not use attention.** Instead, we encode an *exact equivariance* via the comparison function $g$. Notably, this is an orthogonal point. $g$ can likely be applied with Set Transformers (instead of DeepSets) to augment our approach with attention, while at the same time encoding exact equivariances from the start. However, departing from the DeepSet architecture (as mentioned also by you later) would confuse our contribution and muddle the comparison with prior CNP work on exact equivariances, based on (Conv)DeepSets. Thank you for raising this point, which we will add to Related Works and Discussion, with references to the Set Transformer (Lee et al., 2019) and Transformer Neural Processes (Nguyen & Grover, 2022). Apologies for not having included these references before – while they do not deal with exact equivariances, they should definitely be discussed.
**1b. Learning Equivariances vs. Implementing Exact Equivariances (Major Misunderstanding)**
Overall, the crux seems to be a difference in views between "learning (approximate) equivariances" and "building exact equivariances in the network architecture". Our method does the latter.
We stand by our statement that no existing method is able to incorporate *exact* equivariances in the CNP architecture from the start (e.g., like CNNs incorporate exact translational equivariance), in a way which scales well with input dimension. We believe that there is a fundamental conceptual difference between learning approximate equivariances and incorporating specific exact biases directly in the network architecture. Incidentally, the two approaches are not opposed; we could likely incorporate some exact equivariances and let the network learn others.
**2. Analysis**
**We do not use Set Transformers.** Our architecture is based on DeepSets and a fixed comparison function $g$. We aimed to keep a similar complexity between different networks in terms of e.g. number of parameters. We are happy to answer more specific questions about the analysis.
The KL is computed, when possible, between the predictive map of the CNP $p(Y^*| X^*; X, Y)$ and the ground-truth posterior predictive map from the GP. We will include a clear definition, apologies that it was not explained before. This metric was used by e.g. Bruinsma et al. (2023).
**3. Runtime Complexity in Testing**
First, please note that the complexity for the diagonal RCNP (which, as we proved both in theory and practice, is *exact* for translational equivariance) is only $O(MN)$. This is worth considering given that translational equivariance is such an important property and plenty of CNP work on exact equivariances focuses *exclusively* on that (e.g., ConvCNP, Gordon et al. 2020; ConvGNP, Markou et al., 2022; FullConvGNP, Bruinsma et al., 2020). In practice, diagonal RCNPs can be *faster* at runtime than ConvCNP, their direct counterpart (see Table S1).
Second, CNPs are amortized while GPs are not. So even if the asymptotic complexity of FullRCNP is $O(N^2M)$, due to lack of amortization, GPs will need to be trained, which is $O(N^3)$. In many applications, the model receives data sequentially (e.g., BayesOpt), which means that the cost of GP training is applied repeatedly.
**4. Definition of Equivariance is Correct**
Please note that we are defining equivariance of the *prediction map*. Since a prediction map fully depends on its representation $r$, *we can define equivariance of the prediction map based on properties of $r$* (i.e., invariance), there is no contradiction. We will clarify this in the paper.
Our condensed proof is:
The standard definition of equivariance for a generic mapping is
$$\tau F_Z = F_{\tau Z}$$
Starting from the definitions of the ConvCNP paper (Property 2, Gordon et al., 2019):
$$\tau(X,Y) = (\tau X, Y), \qquad \tau f(X) = f(\tau^{-1} X)$$
We apply the defs. above to a prediction map, seen as the function that takes as input the context set $(X, Y)$ and outputs a function that takes as input the target set $X^*$:
$$P_{\tau(X,Y)} = \tau P_{(X,Y)} \Leftrightarrow \forall X^*, \ p(\cdot \mid r((\tau X,Y), X^*))= p(\cdot \mid r((X,Y),\tau^{-1}X^*))$$
Applying $\tau^{-1}$ to each side:
$$P_{\tau(X,Y)} = \tau P_{(X,Y)} \Leftrightarrow \forall X^*, r((X,Y),X^*) = r((\tau X,Y),\tau X^*)$$
which is our definition (Eq. 3).