Learning (Approximately) Equivariant Networks — MLS Lab

Andrei Manolache · Luiz F. O. Chamon · Mathias Niepert

Training starts with room to explore, then tightens symmetry constraints. The exact formulation ends on the symmetry-respecting region; the approximate formulation allows controlled departures.

The method in a nutshell

Rotate a physical system, and its predicted motion should rotate too. This is equivariance: useful structure that can also make a network harder to train.

Adaptive Constrained Equivariance (ACE) adds a flexible, non-equivariant branch to each layer. Constrained optimization adapts its influence during training, giving the model room to explore before enforcing symmetry.

With exact constraints, these branches are driven toward zero and removed for an exactly equivariant model. A resilient variant instead learns how much deviation each layer needs when symmetry is only approximate.

The payoff: on N-body prediction, ACE-trained SEGNN matches the standard model using 5,000 instead of 9,000 training examples. See Figure 2.

Read the paper Code

Learning (Approximately) Equivariant Networks via Constrained Optimization. NeurIPS 2025. Publication record