Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance
This work addresses the need for flexible and efficient equivariant learning methods in machine learning, offering a novel approach that reduces sample complexity and leverages pretrained models, though it is incremental in advancing beyond existing equivariant architectures.
The authors tackled the problem of learning functions with group symmetries without requiring specialized equivariant architectures by introducing a probabilistic symmetrization framework that uses an arbitrary base model and a small equivariant network to enforce equivariance. They demonstrated competitive results against tailored equivariant architectures across various symmetry groups, such as permutation and Euclidean groups, and showed enhanced learning in symmetric modalities like graphs when pretrained from non-symmetric ones like vision.
We present a novel framework to overcome the limitations of equivariant architectures in learning functions with group symmetries. In contrary to equivariant architectures, we use an arbitrary base model such as an MLP or a transformer and symmetrize it to be equivariant to the given group by employing a small equivariant network that parameterizes the probabilistic distribution underlying the symmetrization. The distribution is end-to-end trained with the base model which can maximize performance while reducing sample complexity of symmetrization. We show that this approach ensures not only equivariance to given group but also universal approximation capability in expectation. We implement our method on various base models, including patch-based transformers that can be initialized from pretrained vision transformers, and test them for a wide range of symmetry groups including permutation and Euclidean groups and their combinations. Empirical tests show competitive results against tailored equivariant architectures, suggesting the potential for learning equivariant functions for diverse groups using a non-equivariant universal base architecture. We further show evidence of enhanced learning in symmetric modalities, like graphs, when pretrained from non-symmetric modalities, like vision. Code is available at https://github.com/jw9730/lps.