Improving Equivariant Model Training via Constraint Relaxation

2024-08-23

arXiv.org

AI Summary

Equivariant neural networks suffer from training difficulties, optimization instability, and hyperparameter sensitivity due to strict equivariance constraints. To address this, we propose a progressive constraint relaxation training framework: learnable non-equivariant compensation terms are introduced in intermediate layers, and their activation is gradually tightened via a dynamically decaying soft regularization term—expanding the optimization landscape and enhancing robustness early in training while preserving exact equivariance at convergence. This work is the first to model hard equivariance constraints as learnable, time-varying soft constraints, unifying architectural flexibility with theoretical rigor. Experiments across state-of-the-art equivariant architectures—including SE(3)-Transformer and E(n)-GNN—demonstrate consistent improvements: +1.8% average test accuracy, reduced training failure rates, and diminished sensitivity to hyperparameters such as learning rate.

Technology Category

• Computer Vision: Learning & Optimization for CV
• Constraint Satisfaction and Optimization: Constraint Learning and Acquisition
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

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

Problem

• Equivariant Neural Networks
• Training Difficulty
• Adjustment Complexity

Innovation

• Equivariant Neural Networks
• Progressive Constraint Introduction
• Enhanced Learning Efficiency