Relaxed Equivariance via Multitask Learning

2024-10-23

arXiv.org

AI Summary

Explicitly enforcing rotational/translational equivariance in geometric deep learning incurs substantial computational overhead. Method: This paper proposes a lightweight, architecture-agnostic approximate equivariance learning framework that reframes equivariance as a tunable multi-task objective—introducing a continuously relaxed equivariance loss applied directly to non-equivariant base models (e.g., MLPs or GNNs). Contribution/Results: The approach achieves performance on par with state-of-the-art equivariant models on molecular property prediction and geometric graph tasks, while reducing training time by 2.5× and inference latency by 10×. Crucially, it avoids architectural modifications, enabling plug-and-play integration into existing pipelines. By decoupling equivariance enforcement from structural constraints, the method significantly improves both computational efficiency and practical deployability without sacrificing predictive accuracy.

Technology Category

• Machine Learning: Learning with Manifolds
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

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
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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, equivariant models often pose challenges due to their high computational complexity. In this paper, we introduce REMUL, a training procedure for approximating equivariance with multitask learning. We show that unconstrained models (which do not build equivariance into the architecture) can learn approximate symmetries by minimizing an additional simple equivariance loss. By formulating equivariance as a new learning objective, we can control the level of approximate equivariance in the model. Our method achieves competitive performance compared to equivariant baselines while being $10 imes$ faster at inference and $2.5 imes$ at training.

Problem

• Efficiency
• Symmetry
• Equivariant Features

Innovation

• Multi-task Learning
• Approximate Equivariance
• Efficiency Improvement