Improving Equivariant Model Training via Constraint Relaxation

Summary

The paper proposes to relax the equivariance constraints in equivariance networks by adding a non-equivariant residual in the network weights. The methodology can be applied to different equivariant architectures, e.g. vector neurons, SE(3)-equivariant GNNs, Equiformers. Besides these strictly equivariant networks, the proposed framework can also extend to the approximately equivariant networks by modulating the strength of the unconstrained network weights. Experiments suggest moderate performance increases on a variety of tasks and equivariant architectures when applied the proposed framework.

Strengths

The paper addresses an important problem: overcoming the optimization challenges for equivariant neural networks. The methodology is clearly motivated and easy to follow. The proposed framework of adding an unconstrained component in the network weights seems general enough to be implemented on various (though not all) equivariant architectures, which the authors have demonstrated with quite a few examples.

Weaknesses

* Formatting: the paper is visually difficult to read because of the incorrect use of the parentheses in citations. * Theoretical contribution: as far as I understand, the motivation for relaxing the equivariance constraints is (i) the equivariant networks can be difficult to optimize, or (ii) the data have imperfect symmetry. The paper does not include any theoretical evidence of how the proposed framework of adding unconstrained weights can be helpful in these scenarios. * Specifically, for (i), the paper did not point out what could be the specific challenges during the optimization of equivariant networks, compared to unconstrained optimization. * And for (ii), previous works have shown the error bound of approximately equivariant networks on imperfectly symmetric data (Wang et al 2022) and proposed how to find the symmetry-breaking factors (Wang et al 2023). Compared to these works, I feel this paper didn't provide enough analysis e.g. of equivariance error of the proposed network, or how to interpret the learning result and possibly identify the imperfect symmetry in data. Also, many of the results in the paper are already well-known, e.g. the equation about the Lie derivative. Several works have already proposed to use the Lie derivative as a regularization to encourage equivariance, e.g. Otto et al 2023. * Significance of experimental results: in the experiments, the proposed method only increases performance by a little. Also, in Figure 2, it seems that some models have not converged after 200 epochs. It would be better to train for more epochs and include the full results. In Table 1, the difference between Equiformer and your method is very small. It's hard to verify the significance of this result without looking at the error bar.

Questions

1. The scheduling choice of $\theta$ in Section 3.2 seems arbitrary. Have you tried other scheduling choices, e.g. initializing with zero and constantly increasing it? What is the intuition for the current choice? 2. How sensitive is the model to the regularization coefficient $\lambda_\text{reg}$? 3. L193: "During inference, we evaluate only on the equivariant part of the model." It would be interesting to also see the result for the full model (i.e. including W). I wonder how equivariant and how large W would be under the current regularization. I'm asking because approximately equivariant networks have proved to outperform strictly equivariant networks on certain tasks with imperfect symmetries (Wang et al 2022). As you have both the non-equivariant network and the equivariant one, I'm curious about the comparison between them. 4. Is it possible to have different weighing coefficients for each network layer in the Lie derivative regularization and projection error regularization? As these errors can accumulate after passing through multiple layers, I think it is intuitively reasonable to have larger weights for the first few layers. ### References * Wang, Rui, Robin Walters, and Rose Yu. "Approximately equivariant networks for imperfectly symmetric dynamics." International Conference on Machine Learning. PMLR, 2022. * Wang, Rui, Robin Walters, and Tess E. Smidt. "Relaxed Octahedral Group Convolution for Learning Symmetry Breaking in 3D Physical Systems." arXiv preprint arXiv:2310.02299 (2023). * Otto, Samuel E., et al. "A unified framework to enforce, discover, and promote symmetry in machine learning." arXiv preprint arXiv:2311.00212 (2023).

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors have discussed the limitations.

Summary

The paper proposes relaxing the equivariance constraint on an equivariant network during training. This is done by adding free weights to equivariant linear layers but setting the free weights to zero after training. Further, two regularizations are introduced to stabilize the training: a Lie derivative term encouraging the free weights to be close to equivariant and a term encouraging the influence of the free weights to be low compared to the equivariant weights. The approach is evaluated on several equivariant tasks, showing improved performance compared to the baseline of non-relaxed optimization.

Strengths

1. The paper provides a solid contribution to the understudied topic of optimizing equivariant networks. As far as I am aware, the idea has not been studied in the literature before. 2. The presented experiments show a small but consistent benefit using the proposed approach.

Weaknesses

1. There is no theoretical motivation for why the proposed approach should work. 2. The optimization will be heavier using the proposed approach since a large number of additional parameters are introduced. See also Question 2 below.

Questions

1. The proposed parameterization of equivariant layers during training is $f(x) = f_e(x) + \theta W x$ where $f_e$ is equivariant. $W$ is also encouraged to be close to equivariant through Lie derivative regularization. Would it be possible to remove $f_e$ and $\theta$ to parameterize $f(x)=W x$ as a non-equivariant layer that is regularized by the Lie derivative loss? The projection to the equivariant subspace at the end could be group averaging: $W\mapsto \int_{g\in G} \rho(g^{-1}) W \rho(g) \mathrm{d}g$. Is there something that speaks against such an approach? 2. What is the introduced overhead during training? In terms of time and memory costs. 3. What is the performance of the obtained trained network without projecting to the equivariant subnetwork? I.e. is the approximately equivariant network better than the equivariant? Minor: - Line 243 "table 2" -> "Figure 2"

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

-

Summary

The work proposes a method for improving generalization by relaxing hard equivariance and minimizing the equivariance error as additional regularizer during training.

Strengths

Symmetries play an important role in machine learning and deep learning specifically. There has been recent attention to relaxed forms of equivariance, making it a relevant topic. The paper is well-written.

Weaknesses

There is a lack of attribution to related work, which results in a false sense of novelty. Although the paper does cite many papers on (relaxed) equivariance, it does not always give proper attribution to their contributions. Especially, since several cited papers in the related work section have proposed forms of relaxed equivariance and even minimizing such “relaxation errors” as regularization objective - either in their main method or as a baseline. Yet, the paper claims that “ solutions don’t directly focus on the optimization process itself”. This gives a false impression that regularizing the amount of equivariance in the loss is novel, while it is not. As such, to me it is not entirely clear what the contribution of this paper is.

Questions

1. Projecting back to equivariance models during testing. The first contribution mentions “projecting back to the space of equivariance models during testing”. where is this method? To my understanding, the model is not actually projected, but the error in the projection is rather minimized (but not zero). I could have misunderstood this aspect? 2. Regularization term for deep neural networks. In line 165, a regularization term is proposed to encourage equivariance solutions. It seems this term is only for a single layer. For multiple layers, how are the relative importances between layers chosen? Uniform? Also, it is not clear to me how the overall regularization strength \lambda_reg should be chosen. Cross-validation? Similarly for \theta.

Rating

4

Confidence

4

Soundness

3

Presentation

3

Contribution

1

Limitations

The main contribution of the paper seems to be adding a regularization term that penalizes the ‘equivariance error’. Several works have proposed similar losses and minimizing such regularizers in the training objective. The paper lacks comparison or discussion between different choices of such regularizers, yet alone an empirical comparison. To me it is not clear what the main contribution of this paper is?

Summary

Starting from the consideration that equivariant neural networks, though effective for tasks with known data symmetries, are hard to optimize and require careful hyperparameter tuning, this study proposes a framework to improve their optimization process. This is done by temporarily relaxing the equivariance constraint during training. By adding and progressively reducing a non-equivariance term in intermediate layers, the model explores a broader hypothesis space before converging to an equivariant solution.

Strengths

The paper is clearly written, the topic is significant and of interest to the community.

Weaknesses

Experimental results, despite confirming the theoretical considerations, seem to not compare with competitor methods (the ones reporte in Related works section, e.g. Mondal et al. (2023), Basu et al. (2023b) etc.) or to describe and analyze the additional computational costs (time/memory) of the provided optimization procedure (also in this case, compared to such alternative approaches).

Questions

1) What are the additional computational/memory costs of the proposed method? Is there a trade-off between the achieved increased optimization stability and exploiting the proposed procedure? 2) Could the authors give some examples where the proposed method is not applicable, i.e. the symmetry group is not a matrix Lie group or a discrete finite group? And do they believe that such settings are common or not, i.e. is the proposed method general enough to be applied to real-world scenarios?

Rating

6

Confidence

2

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors adequately addressed the limitations.

-

I acknowledge that I have read the reviews and rebuttals. The authors have successfully answered the most important critique points, so I lean towards keeping my score.