LieAugmenter: Equivariant Learning by Discovering Symmetries with Learnable Augmentations

TL;DR

LieAugmenter is a novel framework that automatically discovers and leverages continuous symmetries in data through learnable augmentations, enhancing model performance and interpretability across various tasks.

Contribution

It introduces a Lie group-based learnable augmentation method that jointly trains symmetry discovery with prediction models, enabling task-adaptive and interpretable symmetry learning.

Findings

Outperforms baselines on image classification tasks.

Effective in predicting N-body dynamics and molecular properties.

Provides interpretable signatures for symmetry absence.

Abstract

Data augmentation is a powerful mechanism in equivariant machine learning, encouraging symmetry by training networks to produce consistent outputs under transformed inputs. Yet, effective augmentation typically requires the underlying symmetry to be specified a priori, which can limit generalization when symmetries are unknown or only approximately valid. To address this, we introduce LieAugmenter, an end-to-end framework that discovers task-relevant continuous symmetries through learnable augmentations. Specifically, the augmentation generator is parameterized using the theory of Lie groups and trained jointly with the prediction network using the augmented views. The learned augmentations are task-adaptive, enabling effective and interpretable symmetry discovery. We provide a theoretical analysis of identifiability and show that our method yields symmetry-respecting models for the identified groups. Empirically, LieAugmenter outperforms baselines on image classification, as well as on the prediction of N-body dynamics and molecular properties. In addition, it can also provide an interpretable signature for detecting the absence of symmetries.