Relaxed Equivariance via Multitask Learning

TL;DR

This paper introduces REMUL, a multitask learning approach that enables unconstrained neural networks to learn approximate equivariance, balancing symmetry constraints with computational efficiency for applications in geometric data modeling.

Contribution

REMUL provides a novel training procedure that learns approximate equivariance through multitask learning, reducing computational costs while maintaining competitive performance.

Findings

Achieves up to 10× faster inference compared to equivariant models.

Enables control over the degree of equivariance during training.

Maintains competitive accuracy while relaxing strict symmetry constraints.

Abstract

Incorporating equivariance as an inductive bias into deep learning architectures to take advantage of the data symmetry has been successful in multiple applications, such as chemistry and dynamical systems. In particular, roto-translations are crucial for effectively modeling geometric graphs and molecules, where understanding the 3D structures enhances generalization. However, strictly equivariant models often pose challenges due to their higher computational complexity. In this paper, we introduce REMUL, a training procedure that learns \emph{approximate} equivariance for unconstrained networks via multitask learning. By formulating equivariance as a tunable objective alongside the primary task loss, REMUL offers a principled way to control the degree of approximate symmetry, relaxing the rigid constraints of traditional equivariant architectures. We show that unconstrained models (which do not build equivariance into the architecture) can learn approximate symmetries by minimizing an additional simple equivariance loss. This enables quantitative control over the trade-off between enforcing equivariance constraints and optimizing for task-specific performance. Our method achieves competitive performance compared to equivariant baselines while being significantly faster (up to 10$\times$ at inference and 2.5$\times$ at training), offering a practical and adaptable approach to leveraging symmetry in unconstrained architectures.