Learning Layer-wise Equivariances Automatically using Gradients

Summary

Edit: Rating updated from 6 to 7 after rebuttal. The goal of the paper is to learn an interpolation between non-equivariant and equivariant models. The authors introduce different convolutional and non-convolutional linear layers, optionally being sparsified via factorizations or a smooth spatial basis. The basic idea is to define the model as a linear combination of translation-equivariant and non-equivariant layers, and to optimize their relative contribution in order to select whether the final model is equivariant or not. This is achieved by posing different Gaussian priors on their parameters, whose widths constitute hyperparameters to be optimized via Bayesian model selection. This model selection differs from prior work, which tuned the hyperparameters explicitly. Bayesian model selection requires the evaluation of a marginal likelihood term, which is infeasible. This is addressed via the Laplace approximation of the parameter posterior proposed by Immer et al. (2022). As this would still require the computation of a Hessian matrix which scales quadratically in the number of parameters, a Kronecker factorization approximation is used. Experiments on image classification datasets show that the interpolated models outperform non-equivariant models and, on tasks with broken equivariance, also strictly equivariant models. The latter are, in turn, performing slightly better on strictly equivariant learning tasks. An ablation to a pure MAP optimization baseline shows that the model selection approach results in significant gains. Results for further groups beyond translations are discussed in the appendix and briefly evaluated in the experimental section.

Strengths

The paper combines ideas from prior work and demonstrates a superior performance compared to these baselines. It relies on a Bayesian approach and shows how it can be made feasible via approximations despite the analytical intractability and high dimensionalities of parameter spaces. Automatically learning hyperparameters is certainly an improvement over manual tuning. The presented empirical evidence shows that the model selection does indeed select equivariant or non-equivariant layers depending on the symmetries of the learning task.

Weaknesses

The main downside of the approach is that it does not really learn equivariant models from scratch, but rather learns to select between pre-specified models with different levels of equivariance. However, this downside is shared by a line of prior work, which the current approach improves upon. I am also not entirely convinced that learning such selections is practically relevant since the appropriate levels of equivariance groups are usually known a-priori or become evident when comparing an equivariant model against a non-equivariant baseline. The experiments suggest that using strictly equivariant models is working better in cases where the desired equivariance group is indeed known a-priori. Instead of describing general equivariant mappings, the main paper considers only conventional convolutions, i.e. translation equivariance. More general results are discussed in the supplementary material and briefly evaluated in the last experimental section. It would have been nice to have a more general formulation in the main paper. It would also have been interesting to see how the method scales to multiple symmetry groups at once. For instance, one could consider all 2^3 (i.e. exponentially many) combinations of using translations, rotations and reflections. The explanation of the method could also be more clear. I had to re-read section 3 and the exact nature of the hyperparameters, introduced in section 3, remained vague till section 4.2. It would have been easier to just mention model selection of hyperparameters in the intro and moving the current section 3 between sections 4 and 5.

Questions

The authors state that "adding symmetries does not directly improve training losses that rely on training fit", however, many papers on equivariant models show their improved convergence rate and final loss when being fitted. Could this statement be clarified and supported more explicitly with evidence? The notation of the limit in equation 9 does not seem to make sense. I guess that the authors just want to say that the implication follows given $\sigma^2=0$? I found the notation of the feature map's domain as $\mathbb{Z}^3$ somewhat confusing and would rather write $[0,C]\times \mathbb{Z}^2$ (spatially supported on $[0,X]\times[0,Y]$). More questions/suggestions are found in the "weakness" section.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

Limitations

Limitations are not explicitly discussed. I do not have ethical concerns

Summary

While (group) convolutions encode strict symmetries into neural network architectures, this paper presents a method for representing flexible symmetry constraints and learning the degree of symmetry automatically (through marginal likelihood objectives). Their method builds on residual pathways to represent each NN layer as the sum of fully connected and convolutional layers, where the initialization variance hyperparameter of the FC weights controls the degree of equivariance and is automatically optimized through the marginal likelihood. They introduce a number of techniques to make this process tractable in practice: * The FC layer is factored along spatial dimensions and possibly sparsified * To efficiently compute a laplace approximation of the marginal likelihood, they use a KFAC approximation of the Hessian. The factorization/sparsification of the FC layers admit additional simplifications to the KFAC Hessian computation. The paper includes experiments comparing this method with plain FC or CNN architectures on synthetic and natural image datasets (MNIST and CIFAR-10), and show that it can learn to adjust the degree of equivariance to achieve good performance on these tasks. Analyzing the optimized hyperparameters shows that the method prefers to make earlier layers more equivariant and later layers less equivariant, in agreement with common architecture design principles. Finally, they show the ability to select between multiple symmetry groups (Conv and 90degree rotation GConv) on CIFAR10.

Strengths

The paper introduces a number of technical innovations to tackle a very challenging problem, automatic symmetry discovery. Marginal likelihood optimization removes the need for validation datasets or handcrafted regularizers when learning symmetries (which were used by some prior work), but is difficult to compute efficiently at the scale of modern NNs. It also attempts to address a common problem of symmetry discovery work where, if you start with no symmetry assumptions (FC), you end up with an enormous number of parameters for high dimensional inputs which is not computationally feasible in practice. The authors show that factoring (along spatial dimensions) and sparsifying the FC layers can reduce the number of parameters and even admit more efficient marginal likelihood estimation.

Weaknesses

The empirical results are limited and do not demonstrated that this method will be broadly applicable, in my opinion. The real datasets considered (CIFAR-10 and MNIST) are relatively simple and existing techniques (e.g., strong data augmentation or SSL + resnets or ViT) likely obtain much stronger results on these benchmarks without needing symmetry learning. Although the purported advantage of symmetry learning is that it can hopefully do better than human chosen symmetry constraints (like humans choosing data augmentations), the empirical results here don't show that. Although in principle the method is applicable to any symmetry group, the experiments seem to focus almost exclusively on translation invariance. Only 6.3 studies the case with 2 possible symmetries: translation (conv) and rotation (GConv). Is there a concern with scalability to more possible symmetries? Ideally, we would like the method to be able to learn any relaxed symmetries more generally without having to restrict the search space to one or two options (translation and rotation). Relatedly, much of the paper and in particular the spatial factorization seem particularly suited to image inputs (or similar input modalities). Would spatial factorization still work well with other data modalities, like graphs representing molecules? The experiments also seem to be focused on image inputs only. Related work: the experiments and aims of this paper are reminiscent to [Elsayed et al, 2020](https://arxiv.org/abs/2002.02959), which also aimed to learn related spatial symmetry by operating on a spectrum between locally connected and convolutional layers. Using locally connected layers can also be viewed as "sparsifying" the fully connected layer to decrease computational costs. It would be interesting to discuss the differences and similarities here, either in text or in terms of empirical results (or both).

Questions

Some minor things: What is P in Eq 1, the number of parameters? Might've missed this. Why does the definition of fully connected preserve the spatial dimensions X and Y? In principle, I'd expect that a fully connected layer (or a conv layer) can change the spatial dimensions. Also, the closing bracket $[0,X$ is missing above Eq 2. Section 6.2: If I understand properly, isn't the optimal behavior just to ignore the fully connected pathway altogether and set its variance -> 0? [Post-rebuttal update]: I have read the author's response. They answer most of my questions. I largely maintain my original score and evaluation.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

2 fair

Limitations

The authors discuss limitations of their method and the general line of work.

Summary

This paper proposes a neural network architecture and gradient-based training algorithm for modeling approximately equivariant functions. The architecture builds upon residual pathway (Finzi et al., 2021), where each layer of network is parameterized as an additive combination of constrained equivariant path and unconstrained fully connected path. A general challenge for such architectures is that equivariance is not favored for fitting training data, so empirical loss minimization will likely result in unstructured solutions that do not generalize. While the residual pathway paper solves this by putting higher weight regularization on unconstrained path, the strength of regularization is a hyperparameter that requires search with validation data in principle. The main motivation in this paper is to learn the extent of equivariance in each layer only from training data in a single training run (online). For this, the authors adopt Bayesian model selection which allows gradient-based learning of hyperparameters from training data by maximizing marginal likelihood estimate, specifically chosen in this work to be Laplace approximation with structured Hessian approximation with KFAC. Given that, the major technical contributions of the paper are on (1) improving the parameter efficiency of residual pathway by introducing convolution on Lie groups as well as factorization and spatial sparsification based on standard exponential basis functions, and (2) specifying the extent of equivariance as hyperparameters controlling the priors placed on the parameters, and (3) deriving KFAC for the proposed parameterizations so that gradient-based learning of hyperparameters controlling the extent of equivariance is made possible through maximization of marginal likelihood estimate. The authors provide experiments mainly regarding discrete 2-dimensional translation symmetry, and demonstrate that the proposed algorithm can (1) learn partial symmetry when needed, (2) recover the standard architecture of convolutional stack postfixed by fully-connected layers, and (3) learn to select from multiple symmetry groups depending on task, solely from training data. Finzi et al., Residual Pathway Priors for Soft Equivariance Constraints (2021)

Strengths

S1. Overall, I think this is a solid work that contributes towards solving a challenging and important problem of learning symmetry from training data by bridging (approximate) equivariant architecture design and Bayesian deep learning. I find no critical issue with originality, quality, clarity, and significance of the work; the writing is overall clear, the design of the algorithm is clearly motivated and presented, and the experimental results seem to support the main claims and motivations.

Weaknesses

W1. In Table 1-3, in addition to test performance, it would be nice if I could see how the models (over)fit to training data (i.e., how they generalize) given that a major motivation of the work comes from the fact that equivariance is not encouraged when fitting training data but beneficial for generalization. W2. The algorithm is empirically demonstrated for discrete 2-dimensional translation and 90-degree rotation symmetries. While the authors argue that extension to more general groups is possible in principle, I find the empirical demonstration regarding the argument is limited compared to the residual pathway paper (Finzi et al., 2021) that provided comprehensive experiments regarding e.g., continuous orthogonal groups as well. Finzi et al., Residual Pathway Priors for Soft Equivariance Constraints (2021)

Questions

Q1. Considering the space of input and output feature maps, how is the group lifting (Appendix A) done in general when there is a set of M groups G1, …, GM (Appendix B)? Do you use direct product G = G1 x … x GM? This doesn’t seem straightforward, I might have missed something. Q2. Does the derivation of KFAC in Appendix C extend to other groups in a similar way to Appendix A? It seems so intuitively, but I would like to ask for a clarification from the authors. Q3. Do the layer indices 0-7 in Appendix D match each layer in Table 6? That is, can I regard the learned layer 6 in Appendix D to correspond to the layer that maps spatial dimension 4 x 4 to 2 x 2 in Table 6? Q4. In Table 6, it is written that the architecture is used for all experiments, and CONV layers are marked with kernel size (e.g., 3 x 3). How should one interpret this for S-CONV layers where, to my understanding, kernel size is not fixed?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

The authors have partially addressed limitations of the work in Section 7; I encourage further clarifications of limitations that needs to be addressed in future work, if there are any.

Summary

The work proposed an automatic way to learn equivariance in each layer by finding a balance between the equivariant layer and the unrestricted fully connected layer. Unlike the previous work on soft equivariance, the work proposed to learn the balance between them via Bayesian model selection using gradients. The work also proposes different parameter-reduction techniques and achieves better results in the conducted experiments.

Strengths

1. The work provides a technique for learning equivariance structure automatically through gradient 2. The work addresses the overparameterization of such soft equivariant models by factorization and sparsifications. 3. The paper is well-organized and clearly written.

Weaknesses

1. Evaluations: The empirical evaluations are conducted primarily on image data where the motivation for finding a balance between the equivariant and non-equivariant layers is not clear (except for the toy problem). Clearly, for image classifications, labels should not be affected by translation or rotation. For the toy problem, a comparison with Finzi et al. [2021a] is not provided. 2. The performance of regular equivariant architecture is not provided, which makes it difficult to measure the gain by the proposed method. Marc Finzi, Gregory Benton, and Andrew G Wilson. Residual pathway priors for soft equivariance constraints. Advances in Neural Information Processing Systems

Questions

1. I am curious about the memory and time (training) required for the proposed model compared to the regular equivariant architecture. 2. What may be the reason for MAP performing worse than the proposed Diff. Laplace method?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

N/A

Response to the Rebuttal

I appreciate the authors for their prompt response and for clarifying the confusion regarding the comparison with the regular equivariant model. I have now updated my evaluation.

Response to rebuttal

Thank you for the response. I recommend the authors to include the discussion regarding W2, Q1, Q2 in the final version of the paper. I have no further questions for now.