Emergence of Equivariance in Deep Ensembles

Strengths: The paper theoretically shows that the equivariance emerges by model ensemble without hand-crafted architecture design. This direction, obtaining equivariance while we can freely choose networks, is important in practical usage. Although the paper is theory-flavored, it is easy to read and follow.

We thank the reviewer for the positive feedback.

Weaknesses: The main finding (emergence of equivariance in deep ensembles) is not very surprising. Data augmentation imposes a bias on a model toward invariance/equivariance, and for me, it's natural to see the averaged model archives that property. I mean, if we have an infinite number of data instances and the model capacity is large enough, the model trained for many steps would be equivariant. The neural tangent approach is of course different from this asymptotic approach, but the main idea should be the same.

It is indeed intuitive that the ensemble becomes equivariant on the training- and test data when trained with data augmentation. However, our results go significantly beyond this intuition in several important directions:

1. Our theoretical analysis shows that for large-width ensembles, the equivariance is not only approximate, as expected from data augmentation, but indeed exact.
2. For training with data augmentation, one would expect (approximate) equivariance to hold only on the data manifold, but not away from it since the ensemble is only trained on the data manifold. This would be true even in the asymptotic case of infinite training data and -time. However, both our neural tangent kernel based arguments and our experiments show that the ensemble becomes a truly equivariant function which is also equivariant away from the data manifold.
3. We show that the ensemble is equivariant for all training times. This is in contrast to the above intuition which would suggest that the ensemble becomes more and more equivariant throughout training as it learns to fit the augmented training data.

The experiments have room for improvement.

1. Instead of equivariance, invariance is evaluated.

In the new version of the manuscript, we have extended our experiments to also include an equivariant task, see Appendix C and in particular Figure 10. We have trained ensembles of fully-connected networks to predict the cross product in $\mathbb{R}^3$, a task which is equivariant with respect to rotations in SO(3). Our experiments show that the predictions of even small ensembles are significantly more equivariant than the predictions of their ensemble members. This is the case even though we only augment with a small finite subset of the continuous symmetry group and holds on- as well as off manifold.

2. Only a cyclic group is considered so it is not clear what kind of consequence can we get for more complex groups such as SO(2), SO(3), or SE(3).

In order to address this point, we have extended our experiments on FashionMNIST. As detailed in Section B.2 in the appendix of the revised manuscript, we have measured invariance with respect to the full SO(2) group of ensembles trained on finite subgroups $C_4$, $C_8$ and $C_6$. Our experiments summarized in Figure 5 show that even for moderately-sized subgroups, the invariance of the ensemble is very high, in accordance to Lemma 6. Furthermore, the ensemble output is much more invariant than individual ensemble members, as is the case for invariance with respect to finite subgroups.

Note that our new experiments on the cross product concern the continuous symmetry group SO(3).

Strengths: Although I am not very familiar with neural tangent kernels, the authors presented the work in a way that was easy to follow. Theorem 4 in particular seems like a very strong result. The authors further consider practical and very relevant limitations such as the finite ensemble, continuous groups, and finite width and prove error bounds. The experiments support the theory.

We thank the reviewer for the positive feedback.

Weaknesses: The type of data augmentation considered seems perhaps a little strong. By using all elements of the group orbit, it naturally lends itself to rewriting the group action as permutations, which seems to be critical in the proof. However, many common data augmentation strategies involve loss of information (e.g. random crops, random non-circular shifts, etc.). If the authors could provide any insights or foreseeable limitations of this work for other data augmentation types, that would be very helpful.

Although many commonly used data augmentation strategies involve a loss of information and are therefore not strictly speaking symmetry transformations, they can usually be interpreted as approximating some group transformation. E.g. random crops approximate scalings and non-circular shifts approximate translations. In particular for network architectures which use local transformations of the input features like convolutions or windowed self-attention, the deviations from a strict symmetry transformation are often restricted to some regions of the input, e.g. the edges. On these grounds, we expect the observed effects of emerging equivariance for ensembles trained with data augmentation to hold approximately in the case of lossy augmentation strategies.

The approximation in going from an augmentation involving information loss to an augmentation which is an exact symmetry transformation is comparable to augmenting with a finite subgroup of a continuous symmetry group and evaluating equivariance for the full symmetry group. For this case, we provide a theoretical bound on the equivariance error in Lemma 6. As our new experiments (see Appendix B.2, and in particular Figure 5, of the revised manuscript) on the FashionMNIST dataset in this setting demonstrate, shows the ensemble small equivariance errors on the full symmetry group, even when trained on relatively small finite subgroups. Therefore, we expect a similar emergence of true equivariance when going from a lossy data augmentation scheme to one which involves exact symmetry transformations.

Strengths:

1. Well-written and well-organized Section 5: the proof sketch is nice.
2. Interesting empirical results that support the theory. The ensembles become more equivariant as width or number of models increases.

We thank the reviewer for the positive feedback.

Weaknesses:

1. Do these results prescribe any particular practical methods, or does it give particular insights on models trained in practice? There does not seem to be much discussion on this. For instance, do people often train ensembles on equivariant data, and how does this compare to single models?

In our view, emergent equivariance of deep ensembles is mainly interesting in scenarios for which ensembles are also used for different reasons such as uncertainty prediction and robustness. It is an important contribution of our manuscript to prove that equivariance comes for free in these scenarios. We stress that deep ensembles are widely used and thus this is a finding of immediate practical value. For example, in protein structure prediction, many models use an ensemble to estimate the uncertainty of the prediction, see e.g. [1, 2]. There are even some works in this context that rely on non-equivariant architectures and ensure equivariance by averaging over an appropriate subset of the group orbit [3]. It is an exciting direction of future research to harness our insights to impose equivariance by ensemble. Note that we demonstrate in the revised Appendices B.2 and C that deep ensembles can also lead to emergent SO(2) and SO(3) equivariance, respectively.

[1] Ruffolo, Jeffrey A., et al. "Fast, accurate antibody structure prediction from deep learning on massive set of natural antibodies." Nature communications 14.1 (2023): 2389. https://www.nature.com/articles/s41467-023-38063-x

[2] Abanades, Brennan, et al. "ImmuneBuilder: Deep-Learning models for predicting the structures of immune proteins." Communications Biology 6.1 (2023): 575. https://www.nature.com/articles/s42003-023-04927-7

[3] Martinkus, Karolis, et al. "AbDiffuser: full-atom generation of in-vitro functioning antibodies." NeurIPS 2023, arXiv preprint arXiv:2308.05027 (2023).

2. Could use more details on the critical assumption on the input layer, see question 1 below.

See below

Questions:

1. Does your assumption on the networks depending on input through $w^k x$ on Page 5 really hold for CNNs? CNNs have their filter coefficients initialized via centered Gaussians, but the underlying matrix is not (because of weight sharing). Thus, there are orthogonal transformations on the input that may change the output (e.g. permute top left with top right pixel).

It is true that there are orthogonal transformations to a CNN which change its output. However, in the NTK, we are concerned with the expectation value over the (inner product of) derivatives of the output with respect to the trainable parameters. If we perform an orthogonal transformation of the input domain, the derivatives change, but their expectation value over initializations remains the same since the initialization distribution is the same for all filter components.

2. Intuitively, what does the deep ensemble output look like at initialization? I am trying to intuit why it is equivariant then.

The ensembles output at initialization is constant, see Eq. 6. As a result, the network is trivially equivariant for all symmetry groups. Data augmentation ensures that this equivariance is not broken by training although the output is no longer constant.

3. Could you give more explanation or intuition about C(x) in Lemma 6 in the main text?

There are some aspects of C that can indeed be intuited:

• The constant C vanishes at the beginning of training as the ensemble is trivially equivariant due to its constant output.
• The constant depends on an expectation over initializations involving the Libschitz constant as well as the expected gradient of the network over initializations. This is to be expected since violations of equivariance with respect to the continuous symmetry will strongly depend on how drastically the ensemble members can change between the points on the group orbit covered by the discretization of the continuous symmetry.

Strengths:

1. This paper presents a very interesting idea of the emergence of equivariance with data augmentation and model ensembles. 2. This paper is generally well-written. 3. The theoretical claims in the paper are sound.

We thank the reviewer for the positive feedback.

Weaknesses:

1. This paper lacks a proper comparison with other methods that can bring equivariance without any constraint on the architecture like [1, 2, 3, 4, 5]. When showing the out-of-distribution transformation results it'll be great to compare with those methods. The current results in the paper are more like ablations of the proposed augmentation and ensembling technique. It is not clear where it stands with other architecture-agnostic equivariance methods. (even if the proposed method does poorly compared to those it'll be good to have those results)

In light of your comments, we have substantially expanded the related works section discussing similarities and differences to the suggested references (see point immediately below).

Furthermore, we have added additional numerical experiments in Appendix B.2 comparing deep ensembles to the canoncialization method of [5]. Briefly summarized, we find the following: We use the exact same model architecture for both the predictor of the canonicalization and the members of the deep ensemble and train on FashionMNIST. Canonicalization leads to an exactly equivariant model whereas our method yields approximate equivariance due to finite number of ensemble members and finite width. As a result, a comparison in the orbit same prediction metric is trivial and clearly shows the benefits of the canonicalization approach. On the other hand, deep ensembles have the advantage that they naturally benefit from well-known advantages of deep ensembles, such as improved accuracy and robustness as well as natural uncertainty estimation. Interestingly, we find a comparable scaling of equivariance to the full SO(2) group when we compare deep ensembles trained with $C_k$ data augmentation with the corresponding canonicalized $C_k$ models, see Figure 5 and 6. Furthermore, we find that deep ensembles lead to a validation accuracy of 91% while the canonicalized model reaches 87%. Both comparisons however come with important caveats: canonicalization allows for full SO(2) equivariance by choosing an appropriate canonicalizer. Futhermore, we are comparing an ensemble of models with a single canonalized model (with the same architecture). An ensembling of canonalized models most likely closes the performance gap.

In our view, emergent equivariance of deep ensembles is particularly interesting in scenarios for which ensembles are also used for different reasons such as uncertainty prediction and robustness. It is an important contribution of our manuscript to prove that an approximate form of equivariance comes for free in these scenarios. We stress that deep ensembles are widely used and thus this is a finding of immediate practical value.

Finally, we want to highlight that our manuscript is a theory paper. Its main objective is to theoretically elucidate surprising emergent capabilities of deep ensembles by using neural tangent kernel theory. We did not aim to propose a highly competitive method which would likely involve applying several tricks of the trade which are harder to model theoretically, such as sampled data augmentation as well as non-iid ensemble member selection.