Improving Equivariant Model Training via Constraint Relaxation

We thank the reviewer for the positive comments as well as raising many points for clarification. We try to address the points raised below.

Citation Format

We apologize for the incorrect formatting of citations and appreciate the reviewer for pointing it out. We will correct this in the final version of the paper.

Theoretical contributions

Unlike in other works on relaxed equivariance, here we specifically focus on the setting where the training data distribution and the model don't have a mis-match in terms of symmetry. As a result the performance improvements come from the fact that the proposed relaxation process can ease the optimization of equivariant networks. As we also noted with other reviewers, amongst equivariant NNs practitioners it is now a somewhat common observation that equivariant NNs can be harder to optimize than their non-equivariant counterparts [1][2][3]. However, the question itself remains unexplored. We take a step towards examining this question in more detail --- we believe that identifying processes that can provide empirical improvements in the optimization of such networks is a reasonable standalone contribution that can be valuable to the community.

However, we have been working on developing theoretical models that develop an optimization-generalization trade-off which could perhaps provide some rationale for why equivariant models might be harder to optimize. We leave explicating on this for a future paper, since the situation can get quite involved.

Significance of experimental results

We thank the reviewer for the feedback regarding the result shown in Figure 2 of the paper. We followed the suggestion and let the models train over the limit of 200 epochs and we show the results in Figure 1 of the attached rebuttal PDF. We see that indeed the additional epochs allow the model to converge, increasing their performance in some cases. It is important to note that while we extended the training epochs we kept the scheduling of theta the same, meaning that theta becomes equal to zero from epoch 200 and onwards. In the original submission we chose to follow the exact setup used from the baseline method (that used 200 epochs) to showcase how our method can provide improvements in performance without the need to tune again all of the training hyperparameters. However, we appreciate the comment of the reviewer and will include the updated figure in the final version of the paper.

Regarding the error bars to quantify the significance of the results for the Equiformer experiment: due to the high computational cost of training a single equiformer model, there is a significant cost involved in providing detailed error bars. The original paper (Liao and Smidt 2023) did not include error bars in their reported results. However, we are actively training more models and we can provide the resulting error bars in the final version of the paper.

Choice of $\theta$ scheduling.

We thank the reviewer for the comment. The main constraint for the scheduling of $\theta$ is that we want to be equal to zero at the end of training so that the final model is equivariant. As a result, the choice of linearly increasing theta doesn't satisfy that constraint. Additionally, we observed that a warm up period in the beginning of training with a small $\theta$ provides significant improvement in the training process. These two observations informed our choice of $\theta$ scheduling where we linearly increase it for the first half of the training and then we linearly decrease it. After the suggestion of the reviewer we will add these observations to the Appendix of the paper.

Sensitivity to the regularization coefficient.

In Figure 2 in the attached rebuttal PDF we show the performance of the VN-Pointnet model for different values of the regularization parameter. This experiment was part of a hyperparameter search using cross-validaton on 80%-20% split of the original training set. As a result the documented performance represents models trained only on the 80% of the train dataset and evaluated on the other 20%.

Additional Details

Regarding evaluation on unprojected relaxed model: We appreciate the suggestion of ther reviewer regarding adding comparison between the relaxed non-equivariant model and the projected equivariant model. Figure 3 (included in the attached rebuttal PDF) shows a comparison between a model trained with our proposed method and a relaxed model before and after we project it to the equivariant space. For the latter model the $\theta$ parameter is kept constant and the equivariant error is controled only through the regularization term. We can observe that, although the performance of the relaxed model before projection is close to the results achieved by our method, the control of the relaxation through our proposed annealing of $\theta$ benefits the performance. We will add this comparison in the final version of the paper.

Regarding the use of different weighting in the Lie derivative regularization: It is true that a more general approach would use different regularization weights for each of the layers of the network. However the use of different weights per layer heavily depends on the individual architectures and will introduce additional complexity to our method. Thus we chose to use the same weight for all the layers to make our method simpler and easier applicable to different tasks independently to the specific architectures used in the task.

References used in this response

[1] Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs Yi-Lun Liao, Tess Smidt, arXiv:2206.11990

[2] Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution Rui Wang, Elyssa Hofgard, Han Gao, Robin Walters, Tess E. Smidt, arXiv:2310.02299

[3] Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network Risi Kondor, Zhen Lin, Shubhendu Trivedi, arXiv:1806.09231

Thank you for the response. We appreciate your comment and engagement.

Regarding the suggestion of using a non-equivariant model and including a Lie derivative regularization:

Optimizing a non-equivariant model with a Lie derivative regularization, even in the case where scheduling of $\lambda_{reg}$ is performed, doesn't guarantee that the final model will be exactly equivariant i.e. the Lie derivative will be zero for all possible inputs. Thus we believe that if we want to learn a model that is exactly equivariant, which is the focus of this work, a projection operation that projects the model to the equivariant space is required. While there are multiple works on training relaxed equivariant networks, they do not consider this projection step. In our work, we described a projection operation that is simple-- so it can be easily incorporated into a typical training process-- and we provided experimental evidence showcasing how it benefits training of equivariant networks.

Regarding the contribution of this work

We recognize that a theoretical analysis (or motivation) of this phenomenon will be an important contribution, and it is a research direction we are interested in pursuing in the future. In general, we think that a first principles theoretical approach to come up with a better optimization could be beneficial to the community since the right language for this problem (in the equivariant setting) is also missing. Nevertheless, we believe that providing a simple training procedure that improves the performance of a large range of equivariant networks is an important standalone contribution to the community. As we discussed in our rebuttal responses, previous works mainly showcase the benefits of learning relaxed equivariant networks and do not consider the case that we focus on, where we want to return to the space of exactly equivariant networks. As a result, we believe that our work still provides a novel perspective that shows that even when a practitioner requires an exact equivariant network, it can still be beneficial to relax the equivariant constraint during training and project back to the space of equivariant networks during inference. Due to the lack of attention, the problem space of improving optimization of equivariant networks is unexplored and our work can be seen as a step to start exploring the space.

As for the other suggestion by the reviewer, we are happy to run experiments using that and include it in the appendix if the reviewer thinks that will make it more comprehensive.

Thank you for the positive assessment of our work. For the questions you raise, we attempt to address them below. Please let us know if we can provide further clarifications.

Theoretical motivation

Within the equivariant NNs community, especially amongst practitioners using such networks in scientific applications, it seems to be a fairly common observation that such networks can be harder to optimize as compared to their non-equivariant counterparts [1][2][3]. We wanted to take a step towards exploring this question and its attendant space, which we believe to be a good enough contribution to the community. We have started to form some theoretical justifications for how/why such approaches should work or be modified. We hope to explicate on a theory for future work. Generally, working out an optimization-generalization trade-off is hard (say compared to working a tradeoff between generalization and approximation). Here we would like to work it out for equivariant networks versus non-equivariant ones.

Regarding the additional Optimization Cost

Thank you for the question. As you correctly point out the cost of optimization of our method will be higher than a baseline equivariant network due to the additional parameters. Nevertheless, due to the parallel nature of the added unconstrained component, the main overhead of the method is on the additional memory required to store the added parameters. This number of additional parameters makes the model of similar size to the typical unconstrained non-equivariant models, which is expected since during training we optimize over a space that is larger to the constrained space of equivariant models. We would like to note that contrary to other methods on approximate equivariant models, that require to access additional parameters both in training and inference, our method requires additional computational resources only during training. As a result, after training is completed our proposed method will not have an effect on inference time. We provide a comparison between the time cost of our proposed training process and the baseline training of different equivariant models. Also to showcase the memory overhead we provide a comparison between the number of learnable parameters between a model trained with our method and an equivariant model trained with a standard training process:

We will add a discussion about the optimization cost of our method in the Appendix of the final version of the paper.

Regarding Suggestion in Question 1

Thanks for the interesting question. An important part of our proposed method is the ability to control the level of the equivariant relaxation explicitly, in addition to the implicit control that comes as a result of the regularization term. With our current parametrization, we can control the level of relaxation by controlling the value of $\theta$. The importance of this control and of the annealing of theta can be seen in Figure 2 (of the submitted paper), where in the case where only the Lie derivative regularization is applied with the value of $\theta$ being constant, the performance of the method decreases.

We think that with the group averaging approach, it could be harder to have explicit control on the level of relaxation of the equivariant constraint during training. As a result when projecting back to the equivariant case it is possible that the performance gap between the relaxed and projected model might be large, even with the lie derivative regularizer.

Additionally, the operation of group averaging can be performed up to approximation in the case of continuous groups, since we need to sample the group elements to approximate the convolution integral.

Performance of the Network without the Projection

In Figure 3 in the rebuttal PDF we added a comparison between the performance of a model trained with our method and the performance of a relaxed equivariant model before and after we project it to the equivariant space. For the relaxed model that we compare with, the $\theta$ parameter is kept constant and the equivariant error is controlled only by the regularizer. We can observe that our method outperforms the relaxed equivariant model with constant $\theta$ even before we project the latter in the equivariant space.

Typos

Thank you for pointing this out. We will correct it in the paper.

References used in this response

[1] Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs Yi-Lun Liao, Tess Smidt, arXiv:2206.11990

[2] Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution Rui Wang, Elyssa Hofgard, Han Gao, Robin Walters, Tess E. Smidt, arXiv:2310.02299

[3] Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network Risi Kondor, Zhen Lin, Shubhendu Trivedi, arXiv:1806.09231

We thank the reviewer for the positive assessment of our paper and for the comments. Below we address all of the points the reviewer has raised.

Comparison with works on Equivariant Adaptation/Fine-tuning of pre-trained models

Thank you for raising this point. While it is true that both our method and the works of of Mondal et al. (2023), Basu et al. (2023b) have a similar motivation i.e. that the optimization of equivariant models can be more challenging than their non-equivariant counterparts, our paper addresses a different research question from their work. Specifically, our paper focuses on improving the optimization of equivariant architectures themselves, while the works of Mondal et al. (2023) and Basu et al. (2023b) focus on techniques that bypass the need for training equivariant models. Specifically, they focus on creating equivariant functions by using pre-trained non-equivariant models e.g. by canonicalization (as in Mondal et al.).

We believe that the difference in focus between our work and that of Mondal et al. (2023) and Basu et al. (2023b), does not quite permit a straightforward and fair comparison. Nevertheless, we include a comparison on a task of point cloud classification (ModelNet40) below.

Here it is important to note that in the case of Basu et al. (2023b), equivariance is achieved by averaging the results over multiple transformed inputs. As a result during inference, the model is required to perform multiple forward passes, which slows down the method's inference time.

We will be happy to include this comparison in the final version (in the main paper or the appendix) if the reviewer believes that it will benefit the overall narrative of the paper.

Additional Details on Computational and Memory costs of the proposed method

We provide a table of the computational overhead of our method compared to the baseline models. Additionally, to illustrate the overhead on the memory constraint we provide a comparison on the number of learnable parameters between our method and the baseline equivariant model. We would like to note that this overhead is only during the training of the models since during inference we remove the additional relaxation layers.

Trade-off for optimization stability

Could the reviewer clarify further? We might be misunderstanding and would be happy to discuss. In general, however, we think that the optimization stability of our method depends on the projection error term which can affect a tradeoff. For perfectly equivariant models in which we don't do any relaxation, we would expect them to be somewhat harder to optimize than a setup with relaxed weights that don't have high projection error. In cases where the projection error is too high, the quality of optimization will go down. Note that in our proposed method we can control the projection error by annealing the value of the $\theta$ parameter of equation 2.

Examples where the proposed method is not applicable

The approach will generally work for a variety of groups of real-world interest. The approach as described will work for compact Lie groups such as $S(1), SO(n), O(n)$, and their finite subgroups like $Z/NZ$, as well as their quotients and products. The approach can be made to work for certain non-compact Lie groups which are reductive for which we can write expressions for the bracket, such as the Lorentz group. It might be challenging for the approach to work for permutation groups.

We thank the reviewer for sharing constructive feedback. We agree that we can improve attribution to related work to emphasize the differences from our work. It is certainly true that there are many works on relaxed equivariance, some of which also employ regularization terms that penalize relaxation errors. However, these works tend to assume that the data itself has some (possibly) relaxed symmetry. The regularization term can then be seen as trying to match the symmetry encoded in the model to the symmetry of the data itself. In our case, the main difference---although it may seem subtle---is that we aren't looking to correct for model misspecification, although we include experiments for approximately equivariant NNs too. We try to make the case that even if we assume that the model is correctly specified, relaxing the equivariance constraint during optimization and then projecting back to the equivariant space can itself help in improving performance. Within the equivariant NNs community it is a common observation that equivariant NNs are harder to optimize than their non-equivariant counterparts [1][2][3]. We take a step towards examining this question in more detail as we believe that this particular hasn't seen exploration specifically.

With the above context in mind, we expand on each of your points below.

Paper Contribution compared to related work

As we said, while there are multiple works on different forms of relaxed equivariance, that also introduce regularization by minimizing the equivariance error, the main difference with our work is that they always remain in the space of relaxed equivariance models. Furthermore, the regularization solves for model mis-specification. We believe that the contribution of our work lies in the projection of the models back to the equivariant space which, contrary to previous works, guarantees that the solution has zero (or fixed) equivariant error. Thus, we are not aiming to address model mis-specification. However, we will improve the attribution to existing work to emphasize this difference, and to highlight the observation that optimizing equivariant networks can be harder than their non-equivariant counterparts.

Projecting back to equivariance models during test

Thank you for your comment. We believe there might be a misunderstanding regarding the projection back to the equivariant space during testing. In Section 3 we define the approximate equivariant linear layer as: $f(x)=f\_e(x)+\theta Wx$

where $f\_e(x)$ is an equivariant linear layer and $\theta Wx$ is an unconstrained term. As we mentioned in lines 142-143, we can project this linear layer to be exactly equivariant by setting $\theta=0$. In that case, only the equivariant part of the layer is activated $f(x)=f\_e(x)$ which guarantees that the layer and as a result the overall model is equivariant and has exactly zero equivariant error. During inference, we set $\theta=0$, resulting in an exact equivariant model. To improve clarity, in addition to the sentence in lines 142-143 we added an explicit note that the operation of setting $\theta=0$ is what we refer to as projection in the rest of the paper.

Regularization Term for Multiple Layers

As we show in the overall training objective (Section 3.3, line 189) the regularization terms, derived for each individual layer, are added in the loss function with the same weight $\lambda\_{reg}$. While it is possible to introduce different weights for the different layers, these weights will heavily depend on the specific network architecture and will increase the complexity of the method. In order to keep the method simple and easily applicable to different architectures we choose to simplify the solution and use the same weight for all layers.

Choice of hyperparameters

For the choice of the hyperparameter $\lambda\_{reg}$ we performed a typical grid search using cross validation with a 80%-20% split of the training set into training and validation set. We found this value to be relatively robust across tasks, so we performed it on the task of point cloud classification and used the value in all other tasks. We include a figure (Figure 2) in the attached rebuttal PDF which shows how the value of $\lambda\_{reg}$ affects the performance of the method while using VN-PointNet as the baseline model. We will add these details and the additional figure in the Appendix of the final version of the paper. We can also include results tuning them individually in the appendix if the reviewer thinks it will be helpful. However, we must note that the accuracy shown in the figure is for the validation set when the model is trained on the 80% of the training set. The results shown in the paper are when we train the model on the complete training set (after we have chosen the hyperparameter $\lambda\_{reg}$).

Regarding the parameter $\theta$: since we perform the scheduling as described in section 3.2, we are not required to do hyperparameter search. The main constraint for the choice of $\theta$ is that it needs to arrive at zero at the end of training. As shown in the ablation in section 4.1 if we do not perform $\theta$ annealing we observe deterioration in performance.

References used in this response

[1] Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs Yi-Lun Liao, Tess Smidt, arXiv:2206.11990

[2] Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution Rui Wang, Elyssa Hofgard, Han Gao, Robin Walters, Tess E. Smidt, arXiv:2310.02299

[3] Clebsch-Gordan Nets: a Fully Fourier Space Spherical Convolutional Neural Network Risi Kondor, Zhen Lin, Shubhendu Trivedi, arXiv:1806.09231

Thank you for the response. I appreciate the authors' efforts to provide additional explanations. To me it seems that the main contribution of the work is the projecting back step. Although most works do not consider such projection, I remain concerned about the lack of sufficient contribution to prior work (for instance relating the used relaxation). Further, some experiments demonstrate that some training runs result in better test loss, but am not sure whether the provided experiments give sufficient evidence for the made claims, which is that introducing some relaxation during training but removing it afterwards through a back projection is beneficial. Especially, since the paper lacks a theoretical analysis.

Regarding the claimed contributions. The used relaxation to me seems exactly equivalent to a residual pathway prior RPP [1]. Although this work is mentioned, its relation to the used relaxation is not. As such, I also argue this not to be a particular novel contribution on its own which is claimed. Similarly, the authors do cite a range of papers that introduce parameterizations of relaxed equivariance, but do not explain how these relaxations relate to the relaxations used in this paper. It is true that these work do often not explicitly consider projecting back, as also argued in the rebuttal, but it is not clear to me why this would prohibits a comparison between these relaxations and explain differences and similarities. Moreover, since projecting back seems to be the primary contribution of this work, why not consider projecting back other common relaxations of equivariance?

Given the above, I lean towards keeping my current score.

[1] Finzi, M., Benton, G., and Wilson, A. G. Residual pathway priors for soft equivariance constraints. 385 In Advances in Neural Information Processing Systems, volume 34, 2021.