Improving Equivariant Model Training via Constraint Relaxation

TL;DR

The paper tackles the optimization difficulties of equivariant neural networks by relaxing hard equivariance constraints during training through an additive unconstrained term weighted by a controllable parameter $\theta$, then projecting back to the equivariant space at inference. It introduces Lie derivative regularization to quantify and constrain deviation from equivariance and includes a mechanism to control projection error via a diminishing $\theta$ schedule. Empirical results across point-cloud classification, molecular dynamics, Nbody simulations, and approximate equivariant models show consistent generalization gains and reduced training variance, validating the approach across diverse architectures (e.g., Vector Neurons, SEGNN, Equiformer). The work provides a practical, scalable method to improve optimization and performance of symmetry-aware networks, with potential extensions to broader symmetry groups and theoretical analysis forthcoming.

Abstract

Equivariant neural networks have been widely used in a variety of applications due to their ability to generalize well in tasks where the underlying data symmetries are known. Despite their successes, such networks can be difficult to optimize and require careful hyperparameter tuning to train successfully. In this work, we propose a novel framework for improving the optimization of such models by relaxing the hard equivariance constraint during training: We relax the equivariance constraint of the network's intermediate layers by introducing an additional non-equivariant term that we progressively constrain until we arrive at an equivariant solution. By controlling the magnitude of the activation of the additional relaxation term, we allow the model to optimize over a larger hypothesis space containing approximate equivariant networks and converge back to an equivariant solution at the end of training. We provide experimental results on different state-of-the-art network architectures, demonstrating how this training framework can result in equivariant models with improved generalization performance. Our code is available at https://github.com/StefanosPert/Equivariant_Optimization_CR

Figures (6)

• Figure 1: Standard training of equivariant NNs is constrained to a limited parameter space which can result in a challenging training process. We propose to relax these equivariant constraints during training, allowing optimization over a broader space of approximately equivariant models. During testing, we project the trained model back to the constrained space---arriving at an equivariant model with enhanced performance compared to equivalent models trained with the standard process.
• Figure 2: Test accuracy on ModelNet40 classification, during training of equivariant PointNet and DGCNN using a baseline training process and different versions of our method. The accuracy is computed for the equivariant models, i.e. for the models after they are projected in the equivariant space.
• Figure 3: (a) Norm of the total Lie derivative of the relaxed PointNet model trained with and without the Lie derivative regularization term. For the computation of the Lie derivative we use the method proposed in gruver2023the. (b) Value of the Lie derivative regularization term for each individual layer of the relaxed PointNet model while we train using our framework and with Lie derivative regularization weight set to $\lambda_\mathrm{reg}=0.01$
• Figure 4: Mean Average Error on the Nbody particle simulation for (a) different model sizes, (b): different dataset sizes.
• Figure 5: ModelNet40 classification accuracy on the validation set using our proposed method with different values of $\lambda_{reg}$. The base model used was the VN-PointNet. The model was trained on a split of the training set containing 80% of the training data. The other 20% of the data were held out as the validation set used to evaluate the model.