Equivariant Neural Tangent Kernels

Expectation and “off-manifold”

We are thankful for pointing out this confusion. The mentioned ensembles are a collection of independently initialized NNs and the average is understood over this family. For considerations about single networks, see heading “Ensembles” in the rebuttal for reviewer hpPR. The presented theorems do not make any particular assumptions over the data distribution.

The term “off-manifold” refers to the manifold hypothesis of the data distribution. Our presented results hold for arbitrary inputs and hence, also hold off-manifold. We will clarify this in the updated manuscript.

Average vs. max pooling

Indeed our recursion relation in Theorem 4.3 only applies to average pooling. Since the expectation value and the maximum operation do not commute this is already true for MLPs and not a specific property of GCNNs.

The average pooling which appears in Theorem 5.2 on the other hand is part of our result: Training an MLP with data augmentation corresponds at infinite width to a GCNN with a group average layer.

Network depth in Theorems 5.2 and 5.3

Indeed, both Theorem 5.2 and 5.3 hold for arbitrary number of layers. For example, the equivalence between the fully connected network defined in (33) and the corresponding GCNN defined in (34) holds for any number of layers $L$. We will add a clarifying remark.

Assumption of NTK-relation in Theorem 5.1

Theorem 5.1 yields a condition on the NTKs to relate augmented to non-augmented networks. The group average on the right hand side of (32) is a consequence of the data augmentation. Theorems 5.2 and 5.3 then provide particular architectures that satisfy this condition.

Extension to equivariance

To extend our results to equivariant networks, one would first need to extend Theorem 5.1 to equivariant augmentation. This is indeed possible by considering a network $\mathcal{N}$ which maps feature maps $f:X\rightarrow\mathbb{R}^n$ into feature maps $\mathcal{N}(f):G\rightarrow\mathbb{R}^{n'}$ by optimizing $\mathcal{N}(\rho_{\text{reg}}(g)f)$ against targets $\rho_{\text{reg}}(g)\hat{f}$ for all $g\in G$. The result is that the mean predictions agree if the NTKs of the two networks satisfy

In order to find networks whose NTKs satisfy this relation, new layers need to be introduced since in the infinite-widht limit, the NTK of MLPs becomes proportional to the unit matrix in the output channels, trivializing the feature map. Here, a separation of the channels into a part which is taken to infinity and a part which is kept finite, e.g.

which corresponds to a fully-connected layer at finite width, could be suitable. We haven’t repeated the calculations for this layer, but it would be an interesting extension of our results which we will mention in the manuscript.

Non-regular representations

Our analysis focuses on group convolutions. These can also be used for (scalar) point clouds by defining the input feature map to be $f(x)=\sum_i \delta(x-x_i)f_i$ for features $f_i$ at $x_i$. Then, the GCNN layers are equivariant with respect to transformations of the positions of the input features. However, this does not cover features transforming in vector- or tensor representations. For these, the setup would need to be extended. Although we have not performed this extension, we do not see conceptual roadblocks at the moment.

Relation to Finzi et al.

In contrast to this work, we specialize to the regular representation. This has the advantage that we can use the convolution theorem to compute the kernel recursions in Fourier space. Extending our framework to arbitrary representations would be a very interesting. A starting point would be to define suitable infinte-width limits of the relevant layers, see e.g. the section “Extension to equivariance” above.

Introduction to the NTK

We acknowledge that NTKs are a topic of substantial technical difficulty. We will revise the background section to make it more accessible stressing that the additional section in the appendix is optional background material. We will also explain the term “frozen” NTK that was adapted from other references (see e.g. Mario Geiger et al. [J. Stat. Mech. 2020]). It refers to the fact that the NTK becomes time-independent during training in the infinite width limit.

Further remarks

• We will add Taco Cohen’s PhD thesis to the literature review.
• Thank you for pointing out the typos which we will fix in a revised version.
• For full generality, we formulated a general layer-$\ell$ lifting layer, because one may include other non-group convolutional layers before as a preprocessing step.
• We will clarify the notation in (8). Indeed, the regular representations on both sides act on functions with different domains.

Ensembles

Thank you for this important question. Although Theorem 5.1 is formulated in terms of ensembles, the statement in fact also holds for individual models at finite width if the infinite-width NTKs in (32) are replaced by the empirical NTKs

as long as they are initialized in such a way that $\mathcal{N}(x)=\mathcal{N}^{\text{aug}}(x)=0$ $\forall x$ at initialization.

Theorem 5.2 however only holds for infinite-width NTKs. Therefore, the combined statement of the equivalence of data augmentation and GCNNs holds also for individual models at infinite width, if $\mathcal{N}^{\text{FC}}(x)=\mathcal{N}^{\text{GC}}=0$ $\forall x$ at initialization.

We will add a corresponding comment to the manuscript.

Sample efficiency

This is indeed an interesting line of further research for which our study lays the ground work. It is well known, see [arXiv:1912.13053], that the conditioning number of the NTK kernel is related to the convergence speed of training. One could therefore compare the conditioning number of the equivariant and standard NTKs to deduce insights into sample efficiency. This is however a task of significant technical difficulty (see Appendix B of the publication linked above which we’d need to generalize to the equivariant case) as it requires the analytical study of the NTK’s spectra as well as their corresponding phase structure which warrants follow-up work.

Approximate results at finite width

Our theoretical claims hold for infinitely wide networks and in the ensemble mean, i.e. for infinitely large ensembles (for comments on extensions to single networks see rebuttal to reviewer hpPR, heading “Ensembles”). In this case, they predict exact agreement between the mean prediction of the invariant GCNN (due to group pooling) and the augmented network throughout training for arbitrary inputs. In our experiments, we test the agreement of the mean predictions for finite ensembles of finite-width networks and find that the predictions align more for larger ensembles. As expected, alignment is imperfect due to finite width, thus we call these results “approximate”. We will clarify this in the revised version.

Extension to equivariant infinite width regression

An extension of our framework to equivariant tasks like quantum mechanical force field prediction is indeed an interesting point. Using the GCNN layers we analyze, it is straightforward to perform regression on targets that are signals on $\mathrm{SO}(3)$, i.e. $f_\text{target}: \mathrm{SO}(3) \to \mathbb{R}^{n_\text{out}}$ in an equivariant way. Here, equivariance is understood with respect to the regular representation, see eq. (5). An equivariant model that is suitable for regression on vector quantities like forces necessitates an output layer that is equivariant with respect to the defining representation of $\mathrm{SO}(3)$. For this, the presented framework would need to be extended by additional layer types and their corresponding kernel relations. (One approach could be to take the argmax of the signal on $\mathrm{SO}(3)$ at the output layer and apply the resulting rotation to a fixed reference vector.) Furthermore, the infinite-width limit of the non-equivariant networks needs to be taken with care in this case. We have outlined a possible way to achieve this in the reply to reviewer kzsw under “Extension to equivariance”. Similarly as mentioned in the discussion by Cohen et al. [ICLR 2018], an extension to e.g. steerable CNNs would also be suitable approach. We consider this an interesting further research direction and will add a discussion to the manuscript.

Complexity of the finite width experiments

We recognize the interest in experiments with more demanding tasks. Since our theorems draw a connection between data augmented and invariant GCNN ensembles, we need to choose an invariant task to test the analytic results numerically. We have now repeated the experiment on a subset of the NCT-CRC-HE-100K data set (https://zenodo.org/records/1214456) consisting of histological images. We have trained on down-sampled images with a resolution of 32x32 pixels and produced OOD samples in this input space. The results are shown in https://postimg.cc/hQzcRpkf, which were obtained in the same fashion as Figure 4 and 6. We use the same architecture as for the CIFAR10 experiment apart from twice as many channels in each hidden layer and we also adapted the learning rate. Each curve is computed 20 times, we plot the means for the metric over the runs and their standard deviations. The larger ensembles reach a validation accuracy of around 80%. The results further support our theoretical claims and we will add them to the manuscript.

Citations

1. On genuine invariance learning without weight-tying, Mokalev et al. [TAGML, ICML 23]

2. On the Ability of Deep Networks to Learn Symmetries from Data: A Neural Kernel Theory, Perin et al.

3. Data Augmentation vs. Equivariant Networks: A Theory of Generalization on Dynamics Forecasting, Wang et al. [ICMLworkshop 22]

4. Fast, Expressive Equivariant Networks through Weight-Sharing in Position-Orientation Space, Bekkers et al. [ICLR 24]

5. Spectrum Dependent Learning Curves in Kernel Regression and Wide Neural Networks, Borelon et al. [ICML 22]

6. Spectral bias and task-model alignment explain generalization in kernel regression and infinitely wide neural networks, Canatar et al. [Nat Comms 2021]

We thank the reviewer for suggesting a number of relevant references and will add them to a revised version of the manuscript.

Relation to the work by Perin et al. [arXiv:2412.11521]

Thank you for pointing out Perin et al. We will include this very recent and interesting contribution in the revised manuscript. As you point out, although they study a similar problem, they focus on a very specific toy problem for which they analyze the spectrum of the NTK in detail. The only equivariant architecture they consider are CNNs. In contrast, our framework captures the learning dynamics of arbitrary group convolutional neural networks.