Learning (Approximately) Equivariant Networks via Constrained Optimization

Equivariant neural networks are designed to respect symmetries through their architecture, boosting generalization and sample efficiency when those symmetries are present in the data distribution. Real-world data, however, often departs from perfect symmetry because of noise, structural variation, measurement bias, or other symmetry-breaking effects. Strictly equivariant models may struggle to fit the data, while unconstrained models lack a principled way to leverage partial symmetries. Even when the data is fully symmetric, enforcing equivariance can hurt training by limiting the model to a restricted region of the parameter space. Guided by homotopy principles, where an optimization problem is solved by gradually transforming a simpler problem into a complex one, we introduce Adaptive Constrained Equivariance (ACE), a constrained optimization approach that starts with a flexible, non-equivariant model and gradually reduces its deviation from equivariance. This gradual tightening smooths training early on and settles the model at a data-driven equilibrium, balancing between equivariance and non-equivariance. Across multiple architectures and tasks, our method consistently improves performance metrics, sample efficiency, and robustness to input perturbations compared with strictly equivariant models and heuristic equivariance relaxations.

💡 Research Summary

This paper introduces Adaptive Constrained Equivariance (ACE), a novel constrained optimization framework designed to address the fundamental challenges in training equivariant neural networks. Equivariant networks explicitly encode known symmetries (e.g., rotations, translations) into their architecture, leading to improved generalization and sample efficiency. However, real-world data often exhibits imperfect symmetry due to noise, measurement bias, or inherent symmetry-breaking phenomena. Strictly enforcing equivariance can then hinder model expressivity and create complex loss landscapes that impede training. Conversely, unconstrained models fail to leverage beneficial symmetries.

Existing relaxation methods, such as adding equivariance penalties to the loss (REMUL) or gradually annealing architectural perturbations, rely heavily on manually tuned hyperparameters (penalty weights, schedules) and lack a principled mechanism to handle partial symmetry in data.

ACE reformulates equivariant network training as a constrained optimization problem: minimize the primary task loss (e.g., classification error) subject to the constraint that the model’s equivariance relaxation parameters γ are zero. By forming the Lagrangian dual of this problem, the constraints are enforced via adaptive dual variables (λ). The proposed algorithm (Algorithm 1) performs gradient descent on the network parameters (θ) and the relaxation parameters (γ), coupled with gradient ascent on the dual variables (λ).

The process begins with a flexible, non-equivariant model (γ initialized away from zero). During training, the update for γ is influenced by both the primary task gradient and the current dual variable λ. The dual variable λ itself accumulates over time based on the magnitude of γ, effectively acting as an automatic, data-driven penalty. If the symmetry is beneficial for the task, the pressure from λ drives γ towards zero, steering the model towards equivariance. If the symmetry is detrimental or partially broken, γ can remain at a non-zero value that minimizes the primary loss, allowing the model to adapt. This creates a homotopy path from a non-equivariant to an (approximately) equivariant model, guided entirely by the optimization dynamics without any pre-defined schedule.

For handling partially equivariant data, the authors extend ACE to inequality constraints (Algorithm 2), where |γ_i| ≤ u_i, with u_i being a learnable upper bound. This allows the model to automatically discover and accommodate the degree of symmetry breaking present in the dataset.

Theoretical analysis (Theorem 4.1) provides bounds on the approximation error incurred by setting small, learned γ values to zero, justifying the practical use of the final “projected” equivariant model.

Comprehensive experiments across diverse domains—image classification (CIFAR-10, Rotated MNIST), graph property prediction, and molecular dynamics simulation—using various architectures (CNNs, GNNs, EGNNs) demonstrate the advantages of ACE. It consistently outperforms strictly equivariant baselines in final accuracy, convergence speed, and sample efficiency. It also surpasses heuristic relaxation methods like REMUL in performance and stability. Furthermore, ACE models exhibit superior robustness to input perturbations and noise. Crucially, in experiments with artificially injected symmetry-breaking noise, ACE automatically adjusts its equivariance level to maintain performance, while strict equivariant models suffer significant degradation.

In summary, ACE provides a principled, general-purpose framework for training equivariant networks. It automates the tricky balance between leveraging useful symmetries and retaining the flexibility to model symmetry-breaking aspects of data, leading to models with better performance, efficiency, and robustness without the need for manual penalty or schedule tuning.