Learning (Approximately) Equivariant Networks via Constrained Optimization

We would like to thank Reviewer bUHD for their comments and for the positive assessment of our work. In the following, we address the Reviewer’s concerns:

W1 - Learning rate $\eta_d$ for the dual variables: Both Reviewer bUHD and Reviewer bnrd point out that, while our method indeed doesn't require hand-crafted schedules for $\gamma$ or tuning of regularization hyperparameters, the dual problem optimization involves a new learning rate $\eta_d$. Empirically, we have observed that having a dual learning rate $\eta_d \approx \eta_p$ works well in most situations. Our recommendation from a practical perspective is to find a good learning rate $\eta_p$ and then set $\eta_d = \eta_p$, or to set $\eta_d$ to a slightly lower value (e.g. if $\eta_p = 9e^{-4}$ then set $\eta_d=8e^{-4}$). We usually had a harder time searching for $\eta_p$ than $\eta_d$, as can be seen in Supplementary Table 4. Using the default learning rates for $\eta_p$ (e.g., from the official implementation of a model) will generally produce good results, but we observed that in some cases they were suboptimal, as stated on Line 951 in the Supplementary.

We have run a sensitivity analysis on the N-Body dataset (1000 samples). Please note that the table contains a baseline experiment ($\eta_d$=N/A means ACE is not used). The initial values for $\eta_p, \eta_d$ were obtained via the initial learning rate grid search (Table 4 of the Supplementary). We set $\eta_p=9e^{-4}$ for all experiments and only evaluate with the strictly equivariant projection $f_{eq}$.

As can be seen, different learning rates $\eta_d$ can affect final performance. However, in all experiments, using ACE produced better results than the base model (N/A), with the best results being obtained when $\eta_p \approx \eta_d$. 

We agree with the reviewer that this is an important point. We believe that this discussion adds further insight to our method and is very useful from a practitioner's perspective. We will include the discussion together with additional plots in the next iteration of our paper.

W2 - Assumption 1 is strong: The main reason why we are making Assumption 1 is that the magnitude of the output of the non-equivariant layer $f_\theta^{neq,i}$ might otherwise dominate $\gamma_i$.

We would argue that the assumption is not strong, and can indeed be easily satisfied in practice - we assume that the layers are M-Lipschitz and bounded operators with $||f_\theta^{neq,i}||\leq B||x||$. This can be easily enforced by reparametrizing the weights, e.g., by using Spectral Normalization [1] (L180-181) so that $B\leq 1$, even when using ReLU non-linearities. Moreover, we do use ReLU non-linearities in practice when $f_\theta^{neq}$ is a MLP.

From a practical standpoint, when seeking to obtain exact equivariance (Algorithm 1), we did not need to use Spectral Normalization --- when doing model selection, we checkpoint based on the validation performance of the equivariant projection $f_\theta^{eq}$, ignoring the non-equivariant component. However, we have observed that in the case of Resilient Constrained Learning (Algorithm 2), Spectral Normalization is needed for the equivariance error of $f_\theta^{eq} + f_\theta^{neq}$ to decrease (Figure 4); this is indeed due to B not being bounded if the weights are not reparametrized.

We will add a more explicit discussion regarding the use of Spectral Normalization in the final version of our paper.

Once again, we would like to kindly thank the Reviewer for taking their time to review our paper and provide constructive feedback. We believe that these points have allowed us to improve both the practical aspects and the theoretical clarity of our method. We are happy to discuss any further questions or concerns in the next stage.

[1]: Miyato et al, "Spectral Normalization for Generative Adversarial Networks", ICLR 2018

We would like to thank Reviewer k8VA for their thoughtful and constructive feedback. In the following, we address each of their concerns in detail.

Q1 - ACE with 2D convolutions:  

We thank the Reviewer for this suggestion. As also noted by Reviewer hCpi, incorporating experiments with 2D equivariant convolutions would be a valuable addition to our paper.

We are currently preparing experiments inspired by the Relaxed Group Convolutions introduced in Wang et al., 2024 [27], where the authors visualize the weights of a relaxed $C_4$-equivariant convolution under different symmetry groups, such as $C_2$. In a similar spirit, we plan to train a $C_4$-equivariant network using ACE and analyze how the non-equivariant component $f_\theta^{neq}$ changes across different symmetry settings.

Furthermore, since $f_\theta^{neq}$ can be interpreted as a correction term to $f_\theta^{eq}$ in the presence of symmetry mismatch, we also plan to visualize and compare the contributions of $f_\theta^{eq}(x)$ and $\gamma f_\theta^{neq}(x)$ to the final prediction. These visualizations will help illustrate how ACE adapts when exact symmetry is not present.

Q2 - Experimental Section is too crowded:  

The Reviewer has a good point. The final version of the paper will include one additional content page that we will use  to refactor the experimental section for clarity and structure. Specifically, we will reorganize the section into subsections, expand on the descriptions of each experimental setup, and provide more detailed analysis of the results.

Q3 - Obtaining Partially vs. Exactly Equivariant Functions; Exact Symmetries Across Experiments:  

We appreciate the Reviewer’s observation and will revise the paper to clarify the distinction between exact and partial equivariance. In particular, we will specify for each experiment which symmetry assumptions are enforced exactly and which are only approximately satisfied.

Q4 - Initializing $\gamma$ at Different Values:  

All experiments use $\gamma_{init}=1$ (we did not see the need to run any hyperparameter search for $\gamma$). Below, we show the results of an experiment on the N-Body dataset that shows the effect of using $\gamma_{init}=0$:

As can be seen, both cases yield better results than not using ACE ($\eta_d=N/A, \gamma_{init}=N/A$). However, better performance is obtained when starting in an unconstrained scenario ($\gamma_{init}=1$).

These results indicate that a zero initialization may lead to faster convergence, but it may also bias the model toward solutions that closely follow the enforced symmetries, potentially limiting exploration in the early stages of training. In contrast, initializing $\gamma_{init}=1$ results in slower convergence, but can lead to better final performance. Due to this year's NeurIPS policy restricting the use of images in the rebuttal, we are unable to include optimization dynamics plots at this stage. However, we will include these plots along with a detailed discussion in the Appendix of the final version.

Limitations Section:  

For the final version, we will move the limitations discussion to a dedicated section before the conclusions.

Once again, we would like to kindly thank the reviewer for taking their time to review our work and provide us with constructive feedback. We believe that addressing these points has helped us significantly improve the clarity of our paper and allowed us to add some important discussions to it.

[27]: Wang et al., "Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution", ICML 2024

We would like to thank the Reviewer for reviewing our work and for their positive assessment!

In the following, we address their concern regarding the implementation of the primal-dual algorithm:

For both the strictly equivariant version (Algorithm 1) and the partially equivariant Resilient Constrained Learning version (Algorithm 2), only minimal modifications are required. In practice, implementing a new optimizer is not necessary. In our codebase, we use the standard PyTorch implementations of SGD or Adam (though any optimizer can be used for either the primal or dual updates) and define separate optimizers for the primal and dual variables. The same approach can be easily adapted to other frameworks such as TensorFlow and JAX.

In the following, we provide a practical PyTorch example showing how a standard equivariant architecture with a conventional optimization routine can be adapted to use ACE:

1. Modifications to the model - adding the $\gamma$ and the non-equivariant parameters:

class Model(nn.Module):
    def __init__(...):
        (...)
        self.gammas = nn.ParameterList([nn.Parameter(torch.tensor(1.)) for _ in range(n_layers)])
        self.f_neq = nn.ModuleList([NoneqModule(...) for _ in range(n_layers)])
        self.f_eq = nn.ModuleList([EqModule(...) for _ in range(n_layers)])

    def forward(...):
        (...)
        for i in range(len(self.gammas)):
            x = self.f_eq[i](x) + gamma[i] * self.f_neq[i](x)
        (...)

2. Modifications to the training routine:

gammas = model.get_gammas()
dual_parameters = []
lambdas = [torch.zeros(1, dtype=torch.float, requires_grad=False).squeeze() for _ in gammas]
dual_parameters.extend(lambdas)

# If we want to use Algorithm 2, add:
# us = [nn.Parameter(torch.zeros(1, requires_grad=True).squeeze()) for _ in range(len(lambdas))]
# dual_parameters.extend(us)

optimizer_primal = optim.Adam/SGD/etc(model.parameters())
optimizer_dual = optim.Adam/SGD/etc(dual_parameters)

# Now, in the training routine:
(...)
loss = loss_fn(preds, y)
dual_loss = 0
slacks = [gamma for gamma in gammas]

for dual_var, slack in zip(lambdas, slacks):
    dual_loss += dual_var * slack

for i, slack in enumerate(slacks): # Explicit, but can be obtained via autodiff
    lambdas[i].grad = -slack

# If we use Algorithm 2, change to:
# loss = loss + 1/2 * torch.norm(us) ** 2
# dual_loss = 0
# slacks = [torch.norm(gamma) - u for gamma, u in zip(gammas, us)]
# for dual_var, slack, u in zip(lambdas, slacks, us):
#    dual_loss += (dual_var * slack) - (u * slack)
# for i, slack, u in enumerate(zip(slacks, us)):
#    lambdas[i].grad = -slack - u

lagrangian = loss + dual_loss
lagrangian.backward()
optimizer.step()
optimizer_dual.step()

# If we use Algorithm 2, add:  
# for lambda in lambdas:
#     lambda.clamp_(min=0)

We will include this example on implementing ACE in the Supplementary material of the final version of the paper, along with links to the official implementation.

We would once again thank the Reviewer for their feedback and are open to address any further questions or concerns that they might have!

We're very happy to hear that Reviewer hCpi found our paper interesting and compelling!

In the following, we will address the questions raised by the Reviewer:

Q1 - Bounded Translation clarification: We did not find a reference to “translation” or “bounded” on line 214 and therefore assume that the Reviewer refers to line 204. There, we highlight that in most cases, the transformations that we are interested in are isometries, i.e., norm-preserving transformations (B = 1), the most common of which are rotations and translations (involving either padding or wrapping). The results of Theorem 4.2, however, hold for arbitrary B-bounded transformations. We will add additional clarifications regarding this. 

We might, however, have misunderstood the question of the Reviewer; if our answer is not satisfactory, we kindly ask the Reviewer to ask us for further clarification.

Q2 - Symmetry assumed in MoCap: The symmetry assumed for MoCap is the Special Euclidean group $SE(3)$ --- so translations $\mu \in \mathbb{R}^3$ and rotation matrices $R\in SO(3)$. Intuitively, as mentioned in the original EGNO paper [1], "As a practical example, to model dynamics by learning a function to predict $\mathcal{G}^{(t+1)}$ from $\mathcal{G}^{(t)}$, we hope that any rotations and translations applied to the current state coordinate $x(t)$ should be accordingly applied on the predicted next one $x(t+1)$, while the velocities $v$ should also rotate in the same way but remain unaffected by translations". We will add further clarifications in the final version of our paper.

Q3 - Simple example with an easy symmetry: We thank the Reviewer for this suggestion. As also noted by Reviewer k8VA, incorporating experiments with 2D equivariant convolutions would be a valuable addition to our paper.

We are currently preparing experiments inspired by the Relaxed Group Convolutions introduced in Wang et al., 2024 [2], where the authors visualize the weights of a relaxed $C_4$-equivariant convolution under different symmetry groups, such as $C_2$. In a similar spirit, we plan to train a $C_4$-equivariant network using ACE and analyze how the non-equivariant component $f_\theta^{neq}$ changes across different symmetry settings.

Furthermore, since $f_\theta^{neq}$ can be interpreted as a correction term to $f_\theta^{eq}$ in the presence of symmetry mismatch, we also plan to visualize and compare the contributions of $f_\theta^{eq}(x)$ and $\gamma f_\theta^{neq}(x)$ to the final prediction. These visualizations will help illustrate how ACE adapts when exact symmetry is not present.

Once again, we would like to kindly thank the Reviewer for their overwhelmingly positive assessment of our paper and for their feedback. If they have any further questions, we would be happy to address them!

[1]: Xu et al., "Equivariant Graph Neural Operator for Modeling 3D Dynamics", ICML 2024

[2]: Wang et al., "Discovering Symmetry Breaking in Physical Systems with Relaxed Group Convolution", ICML 2024

We would like to thank Reviewer bnrd for their thoughtful and constructive review. In what follows we clarify the Reviewer's concerns.

W1 - Learning rate $\eta_d$ for the dual variables: Both Reviewer bnrd and Reviewer bUHD point out that, while our method indeed doesn't require hand-crafted schedules for $\gamma$ or tuning of regularization hyperparameters, the dual problem optimization involves a new learning rate $\eta_d$. Empirically, we have observed that having a dual learning rate $\eta_d \approx \eta_p$ works well in most situations. Our recommendation from a practical perspective is to find a good learning rate $\eta_p$ and then set $\eta_d = \eta_p$, or to set $\eta_d$ to a slightly lower value (e.g. if $\eta_p = 9e^{-4}$ then set $\eta_d=8e^{-4}$). We usually had a harder time searching for $\eta_p$ than $\eta_d$, as can be seen in Supplementary Table 4. Using the default learning rates for $\eta_p$ (e.g., from the official implementation of a model) will generally produce good results, but we observed that in some cases they were suboptimal, as stated on Line 951 in the Supplementary.

We have run a sensitivity analysis on the N-Body dataset (1000 samples). Please note that the table contains a baseline experiment ($\eta_d$=N/A means ACE is not used). The initial values for $\eta_p, \eta_d$ were obtained via the initial learning rate grid search (Table 4 of the Supplementary). We set $\eta_p=9e^{-4}$ for all experiments and only evaluate with the strictly equivariant projection $f_{eq}$.

As can be seen, different learning rates $\eta_d$ can affect final performance. However, in all experiments, using ACE produced better results than the base model (N/A), with the best results being obtained when $\eta_p \approx \eta_d$. 

We agree with the reviewer that this is an important point. We believe that this discussion adds further insight to our method and is very useful from a practitioner's perspective. We will include the discussion together with additional plots in the next iteration of our paper.

W2 - Performance gap between official VN-DGCNN and our method: In our experiments, we adopt the setup from Pertigkiozoglou et al. [1], a related work that also focuses on training equivariant models. This setup uses a sub-sampled resolution of 300 points per point cloud, in contrast to the original VN-DGCNN paper [2], which uses 1024 points. Additionally, in our implementation, we only tuned the primal and dual learning rates $\eta_p, \eta_d$, without modifying other components or performing a broader hyperparameter search.

Due to these differences, both in input resolution and in the level of tuning, our results are not directly comparable to those reported in [2]. We will continue to explore the source of these differences and include the additional results and a discussion in the final version of our paper.

What is more, we thank the Reviewer for pointing out the typo in Figure 7.

W3 - Having a large norm for the non-equivariant component might increase the scale of its outputs: While it is indeed true that the magnitude of the output of the non-equivariant layer $f_\theta^{neq,i}$ might dominate the constraint $\gamma_i$, we assume that the layer is M-Lipschitz and a bounded operator with $||f_\theta^{neq,i}||\leq B||x||$ holding throughout training. This can be easily enforced by reparametrizing the weights, e.g., by using Spectral Normalization [3] (L180-181) so that $B\leq 1$.

From a practical standpoint, when seeking to obtain exact equivariance (Algorithm 1), we did not need to use Spectral Normalization --- when doing model selection, we checkpoint based on the validation performance of the equivariant projection $f_\theta^{eq}$, ignoring the non-equivariant component. However, we have observed that in the case of Resilient Constrained Learning (Algorithm 2), Spectral Normalization is needed for the equivariance error of $f_\theta^{eq} + f_\theta^{neq}$ to decrease (Figure 4); this is indeed due to B not being bounded if the weights are not reparametrized.

We will add a more explicit discussion regarding the use of Spectral Normalization in the final version of our paper.

W4 - Naming Inconsistencies: We thank the Reviewer for pointing out the naming inconsistencies. We will fix them for the final version.

Once again, we kindly thank the Reviewer for taking their time to review our paper. We believe that addressing their concerns has helped us significantly improve the clarity of our paper and allowed us to add some interesting discussions to it. We are open to further discussions in the next reviewing phase!

[1]: Pertigkiozoglou et al., "Improving Equivariant Model Training via Constraint Relaxation", NeurIPS 2024

[2]: Deng et al, "Vector Neurons: A General Framework for SO(3)-Equivariant Networks", ICCV 2021

[3]: Miyato et al, "Spectral Normalization for Generative Adversarial Networks", ICLR 2018

I thank the authors for their rebuttal and for addressing my concerns. I appreciate the inclusion of additional discussion regarding the choice of hyperparameters used in this work since it can significantly increase the ease of applying this method in different models and tasks. I thank the authors for clarifying my misunderstanding regarding the different setup between this work and [2]. Finally, regarding the use of Spectral Normalization, I believe that its use is important to be more extensively discussed in the paper, since (relating to my initial concern) it can also add additional hyperparameters that are important for the method. Since most of my concerns are addressed by the authors who are willing to add the additional discussions in the update version of the paper, I will increase my score to 5.