Research map / Probabilistic and scientific modeling / Molecules and proteins
Learning (Approximately) Equivariant Networks via Constrained Optimization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Adaptive Constrained Equivariance (ACE), a primal-dual constrained optimization method that starts a network non-equivariant and automatically tightens or relaxes equivariance via learned modulation coefficients, improving accuracy and sample efficiency on N-body, QM9, ModelNet40 and MoCap.
Problem
Training equivariant neural networks is hard: strict symmetry yields complex loss landscapes and hurts fit when data is only partially symmetric, while existing relaxation methods rely on hand-tuned penalties or schedules.
Approach
Cast equivariant training as a constrained problem where non-equivariant branches are gated by homotopy parameters gamma constrained to zero (equality) or bounded by learned slacks (resilient inequality); dual variables automatically set the annealing and the amount of relaxation per layer.
Key results
- On N-Body, the ACE-trained SEGNN reaches test MSE 3.8e-3 versus 5.6e-3 for vanilla SEGNN, and matches the 9000-sample vanilla performance with 5000 samples (~44% less data).
- On ModelNet40, ACE VN-DGCNN improves class and instance accuracy by +2.78 and +1.77 percentage points over normal training.
- On QM9, ACE lowers SchNet MAE on all 12 targets (e.g., U0 from 9.64 to 3.40), while ACE SEGNN gives comparable results with mostly small gains.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Train SEGNN with ACE equality constraints and use its equivariant projection | SEGNN · N-Body | test MSE: 5.6{ $\pm$0.25} -> \textbf{3.8{ $\pm$0.17}} | Helped |
| Train SchNet with ACE equality constraints | SchNet · QM9 μ | MAE: 0.0511{ $\pm$0.0011} -> \textbf{0.0393{ $\pm$0.0007}} | Helped |
| Train SEGNN with ACE equality constraints | SEGNN · QM9 μ | MAE: 0.0371{ $\pm$0.0002} -> \textbf{0.0362{ $\pm$0.0007}} | Helped |
| Train SEGNN with ACE equality constraints | SEGNN · QM9 ε_lumo | MAE: \textbf{0.0244{ $\pm$0.0003}} -> 0.0257{ $\pm$0.0006} | Hurt |
| Train ACE equivariant VN-DGCNN | VN-DGCNN · ModelNet40 | class accuracy: +2.78 percentage points | Helped |
| Train ACE equivariant VN-DGCNN | VN-DGCNN · ModelNet40 | instance accuracy: +1.77 percentage points | Helped |
| Use ACE equality-constrained projection instead of the strictly equivariant baseline under point dropout | VN-DGCNN · ModelNet40 with 0-85% point dropout | test accuracy | Helped |
| Allow a non-equivariant branch with ACE resilient inequality constraints | EGNO · CMU MoCap Run and Walk | test MSE | Helped |
| Initialize γ to 1 rather than 0 (ablation) | SEGNN with ACE · N-Body | test MSE: $5.39 \pm 0.33$ -> $4.97 \pm 0.01$ | Helped |
| Initialize γ to 0 rather than 1 (ablation) | SEGNN with ACE · N-Body | test MSE: $4.97 \pm 0.01$ -> $5.39 \pm 0.33$ | Hurt |
| Use ACE rather than removing ACE (ablation) | SEGNN · N-Body | test MSE: $6.64 \pm 0.14$ -> $4.97 \pm 0.01$ | Helped |
| Set dual learning rate η_d to 8×10^-4 (ablation) | SEGNN with ACE · N-Body | test MSE: $6.64 \pm 0.14$ -> $4.97 \pm 0.01$ | Helped |
Surprising findings
- Strictly equivariant SchNet on QM9 severely hinders convergence without relaxation, and ACE reduces its MAE substantially on every target.
- On CMU MoCap, the gamma coefficients oscillate around zero without converging, indicating strict equivariance is too restrictive for that dataset.
Limitations
- Adds extra non-equivariant MLP layers during training, increasing parameter count and modestly slowing training.
- Experiments focus on point clouds, molecules and motion trajectories; 2D image and graph settings are only a toy C4 experiment.
- Hyperparameters (primal and dual learning rates) still require a grid search per dataset.
Experimental setup
Datasets: N-Body QM9 ModelNet40 CMU Motion Capture
Benchmarks: N-Body QM9 ModelNet40 CMU Motion Capture
Metrics: MSE MAE class accuracy instance accuracy
Baselines: SEGNN SEGNN Rel. REMUL EGNO SchNet VN-DGCNN EGNN ClofNet SE(3)-Transformer Radial Field
Builds on: SEGNN resilient constrained learning REMUL Residual Pathway Priors VN-DGCNN
Closest papers by meaning
- Improving Equivariant Model Training via Constraint Relaxation (NeurIPS 2024 · similarity 0.93)
- A Tale of Two Symmetries: Exploring the Loss Landscape of Equivariant Models (NeurIPS 2025 · similarity 0.92)
- Improving Equivariant Networks with Probabilistic Symmetry Breaking (ICLR 2025 · similarity 0.90)
- On the hardness of learning under symmetries (ICLR 2024 · similarity 0.90)
- Approximately Equivariant Neural Processes (NeurIPS 2024 · similarity 0.90)
- On Universality Classes of Equivariant Networks (NeurIPS 2025 · similarity 0.90)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.