Improving Equivariant Model Training via Constraint Relaxation

TL;DR

This paper introduces a training framework for equivariant neural networks that relaxes the equivariance constraint during training, enabling better optimization and improved generalization, then gradually enforces strict equivariance.

Contribution

The authors propose a novel relaxation-based training method for equivariant networks that enhances optimization and generalization performance.

Findings

Improved generalization on state-of-the-art architectures

Effective relaxation of equivariance constraints during training

Convergence to strict equivariance at the end of training

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