Learning Equivariant Models by Discovering Symmetries with Learnable Augmentations
This addresses the need for interpretable and robust equivariant models in geometric domains without requiring a priori symmetry knowledge, though it builds incrementally on existing soft equivariance methods.
The paper tackles the problem of learning equivariant models without prior knowledge of data symmetries by proposing SEMoLA, which discovers symmetries via learnable augmentations and encodes them into unconstrained models, achieving high prediction performance across various datasets.
Recently, a trend has emerged that favors shifting away from designing constrained equivariant architectures for data in geometric domains and instead (1) modifying the training protocol, e.g., with a specific loss and data augmentations (soft equivariance), or (2) ignoring equivariance and inferring it only implicitly. However, both options have limitations, e.g., soft equivariance still requires a priori knowledge about the underlying symmetries, while implicitly learning equivariance from data lacks interpretability. To address these limitations, we propose SEMoLA, an end-to-end approach that jointly (1) discovers a priori unknown symmetries in the data via learnable data augmentations, and uses them to (2) encode the respective approximate equivariance into arbitrary unconstrained models. Hence, it enables learning equivariant models that do not need prior knowledge about symmetries, offer interpretability, and maintain robustness to distribution shifts. Empirically, we demonstrate the ability of SEMoLA to robustly discover relevant symmetries while achieving high prediction performance across various datasets, encompassing multiple data modalities and underlying symmetry groups.