Ensembles provably learn equivariance through data augmentation

Authors: Oskar Nordenfors, Axel Flinth
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
Keywords: equivariance, invariance, ensemble models, data augmentation, sgd

Abstract:
Recently, it was proved that group equivariance emerges in ensembles of neural networks as the result of full augmentation in the limit of infinitely wide neural networks (neural tangent kernel limit). In this paper, we extend this result significantly. We provide a proof that this emergence does not depend on the neural tangent kernel limit at all. We also consider stochastic settings, and furthermore general architectures. For the latter, we provide a simple sufficient condition on the relation between the architecture and the action of the group for our results to hold. We validate our findings through simple numeric experiments.

Decision: reject
Reviewer scores: 3, 6, 6

Reviews and discussion:
<thread>
creation date:
1730497417502

thread_type:
official_review

<thread_object>
creator: reviewer_nmuk

<content>
summary:
The paper presents a theoretical analysis showing that data augmentation can lead to equivariance in deep ensembles. The paper's main result is that under several assumptions (e.g. on initialization, architecture, etc.), deep ensembles trained with data augmentation are equivariant in mean, even when individual models are generally not. A similar result was previously presented, but the paper extends these previous results, which were primarily focused on infinitely wide NNs trained with gradient descent under full augmentation, to ensembles of finite-width trained with SGD and random augmentation.
The paper is mainly theoretical and validates the theoretical results through limited and small-scale empirical experiments.

strengths:
1. The paper is well-structured and easy to follow.
1. The paper extends previous results to more reasonable and applicable settings. This is a significant extension.

weaknesses:
I like the paper and believe it has a sufficient contribution and interesting results. However, there are several limitations stated below:

1. While the assumptions for the theoretical analysis are more applicable compared to previous works, they still hold only for infinite-size ensembles. Any analysis (including empirical) on the error bounds for finite ensembles would be beneficial.
1. While the results are important, the novelty is somewhat moderate in the sense that the emergent equivariance property of ensembles was previously proposed and the fact that the theoretical analysis heavily relies on previous works [1].
1. From the empirical evidence, it is unclear if some of the assumptions (like symmetric initialization) are indeed necessary. The authors discuss this, but I believe it can be extended further.
1. Empirical evaluation is limited. It would be beneficial to extend it to more settings, even by small modifications like considering cyclic groups C_k of different orders (k), different architectures, model sizes, etc.
1. It would be beneficial to see the impact of ensemble size on the metrics in Table 1, like adding a line plot for ensemble size vs. OSP. The authors show results for different sizes, but summarizing them in one clear view would make it easier to follow.
1. The paper could benefit from a clearer and more explicit discussion of the limitations of the results.
1. Minor:
    - Line 37: “... a definitive question to the question…”.

Reference

[1] Flinth & Ohlsson, Optimization Dynamics of Equivariant and Augmented Neural Networks, 2023.

questions:
1. Why does the OSP not increase at initialization when ensemble size increases?
1. From the figures, it seems like the results could improve with more epochs (also for baselines). Could you please provide results with a larger number of epochs?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730565560048

thread_type:
official_review

<thread_object>
creator: reviewer_ljyp

<content>
summary:
The paper expands the results of Gerken & Kessel that show that data augmentation produces equivariant ensembles of models using NTK, by looking at finite network sizes. They then show empirically that their theoretical results indeed hold in practice (up to sampling errors).

strengths:
- It generalizes the results in Gerken & Kessel 
- The topic of invariance/equivariance is important so these results would be of interest to people in that community

weaknesses:
My main issue is with the writing: 
- The results presented in the main text are quite trivial, that if you start with an invariant distribution and use an invariant flow you end up with an invariant distribution. The more interesting results are in the appendix (appendix B and C)
- You writing $\mathcal{L} = A_\mathcal{L} + T\mathcal{L}$ with $T\mathcal{L}$ the tangent space is very confusing, as tangent space is defined for a manifold and we are talking about a linear space. It needlessly complicates things as there is no need to involve differential geometry when we are working on linear spaces.

questions:
The results in Table 1 aren't that clear to me. In the asymmetric case where you have a symmetric initialization, shouldn't you get results that are similar to the symmetric case? Yet there is a large gap
</content>
</thread_object>

</thread>
<thread>
creation date:
1730600353980

thread_type:
official_review

<thread_object>
creator: reviewer_yfbu

<content>
summary:
This paper shows that an ensemble of models when trained with data augmentation leads to emergence of equivariance properties naturally. The results generalize over past known results based on NTKs. The theory assumes some basic assumptions on the architecture and shows that, when the initialization of the weights in an architecture has some symmetry, then, the expected architecture of the ensemble is equivariant. Experimental results with various ensembles validates the results for the C4 group of symmetries.

strengths:
- The work show the emergence of equivariant in ensemble models
- The work generalizes previous works where the proof relied on NTKs
- Experiments with large ensemble of models show the emergence of equivariance

weaknesses:
I have several concerns over the usefulness of the theory and the experimental results.

Usefulness of theory:
- What is the use of the theory in model design or practical use cases? Since equivariant models seems to give perfect equivariance and data augmentation techniques give approximate equivariance. So, I am wondering what is the use of ensemble technique for symmetries, especially, given that we need over 1000 models to get good equivariant results.
- What are the advantages of the proposed technique compared to existing symmetrization and canonicalization methods [1-4] that can convert non-equivariant models into equivariant ones using techniques somewhat similar to ensemble methods but with additional transformations that looks similar to augmentation.

Experimental Results:
- Although the experimental does show that the architecture with symmetric support does give invariant output, but even the asymmetric architecture seems to be giving invariant output, questioning the usefulness of the theory. It is also discussed in the paper about the symmetric states being attractors potentially, but, it still makes the current theory not very useful.
- Experiments are only shown for C4 symmetries

[1] Basu, Sourya, et al. "Equi-tuning: Group equivariant fine-tuning of pretrained models." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 37. No. 6. 2023.

[2] Mondal, Arnab Kumar, et al. "Equivariant adaptation of large pretrained models." Advances in Neural Information Processing Systems 36 (2023): 50293-50309.

[3] Basu, Sourya, et al. "Efficient equivariant transfer learning from pretrained models." Advances in Neural Information Processing Systems 36 (2024).

[4] Kaba, Sékou-Oumar, et al. "Equivariance with learned canonicalization functions." International Conference on Machine Learning. PMLR, 2023.

questions:
Please see the weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732031576659

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Planned updates

comment:
We would like to thank all the reviewers for their work. Their reviews are all insightful, and contain many valuable suggestion. We have responded to their questions and comments in individual posts.

Let us here only advertize the two big updates we will make to the manuscript before the end of the discussion period.

* We will perform a new set of experiments for the C16 group.

* We will make a more serious evaluation of our models also for smaller ensemble sizes, providing some empirical results in this direction. Here is already a plot formed by measuring the metrics at epoch 10 for the different models (shown is a mean of 10 bootstrapped ensemble samples for each size and model) : [Plot](https://anonymous.4open.science/api/repo/ensemble_experiment-1B83/file/graphics/complete_plot.png?v=5df1ae3f) -- the general trend is that the difference in equivariance is detectable already for moderate ensemble sizes. In the updated version of the paper, we will provide data for 30 bootstraps, and perform some statistical tests. See also the comment to reviewer nmuK.

We will try to make the updates as soon as possible.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732549199100

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Results of updated experiments (I)

comment:
Dear reviewers,

we have had some technical issues, but have now finally managed to run our updated experiments. The results are interesting, and not as clear-cut as one could have wished for. Still, we think that they support the relevance of our theory rather than speak against it.

First, we have, as advertized, re-evaluated our previous experiments (i.e., for $C_4$) for 30 bootstraps instead of 10 bootstraps per sample size. There are no surprises here: the symmetric architecture still outperforms the asymmetric ones, and does so with statistical significance ($p<.001$) from $250$ ensemble members onwards (with respect to the divergence metric, even from $75$ ensemble members). Here is an updated plot: [plot_C4](https://anonymous.4open.science/api/repo/ensemble_experiment-1B83/file/graphics/C4_nearest.png?v=b54e4155)

We have also run the same experiment for the bigger group $C_{16}$. Let us first note that when using this group, we stray from the setting in the paper. The group is no longer acting directly on the support of the images - due to interpolation effects. Hence, the lifted representation $\rho$ on the linear layers $A_i$ no longer perfectly corresponds to rotation of the filters $\varphi_i$ (Example 3.2 is no longer valid). In fact, again due to interpolation, Assumption 1 and 3 are also not given and the spaces $\mathcal{L}$ are no longer invariant.

With this said, here is the plot for our experiments: [plot_C16](https://anonymous.4open.science/api/repo/ensemble_experiment-1B83/file/graphics/C16_nearest.png?v=d3772c9f). 

We see that while the symmetric filters still produce more equivariant ensembles than the asymmetric ones on the in-distribution MNIST test data, they are actually not better, and even worse with respect to the divergence metric, on the CIFAR10 data. The most striking difference to the $C_4$ experiments are however that the all of the models are significantly less equivariant on the CIFAR10 data. This was not what was expected. One realizes that this might have to do with the way we have performed our augmentation: We have used the default 'nearest' interpolation option in torchvision to perform the augmentation, and also making sure that the background of the images are uniform. These transformations are not a representation of the group $C_{16}$ -- if we in particular think about filters of size $3\times 3$, the small rotations in fact do nothing.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732737691728

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Updated version of the pdf

comment:
We have now uploaded an updated version of the pdf. All changes are marked in blue. 

In short, the added material is more or less what was posted in our last comment. We have also made some other minor changes, such as changing the terminology "tangent space" and correcting typos.

In a final version of the paper, we will redo the failed experiments for C16 with BILINEAR interpolation, and make an experiment with standard CNN:s for the NEAREST interpolation. We apologize for not being able to do so before the deadline for pdf updates.

We would like to thank the reviewers again their suggested changes. We think that they have improved the paper.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper studies how equivariance emerges in ensembles of neural networks trained with data augmentation. The authors extend prior theoretical results by showing that equivariance holds under weaker conditions that exist in prior work (e.g. without requiring the NTK limit).
The reviewers appreciated the clear writing and sound theoretical development. However, several concerns were raised about both the theory and experiments: The theoretical contribution it appears to be minor, and the empirical validation was quite limited in scope. The was also some apparent disconnect between the theory predictions and experimental results, which brought into question the relevance of the theory.
The authors engaged constructively with reviewers in the rebuttals. 
However, the concerns remain significant enough to prevent me from recommending acceptance.