Learning Symmetric Embeddings for Equivariant World Models
This work addresses a bottleneck in using equivariant models for data with complex symmetries, offering a method to enhance data efficiency and generalization in machine learning applications.
The paper tackles the problem of applying equivariant models to data with unknown transformation effects by proposing symmetric embedding networks (SENs) that map inputs to a feature space with known symmetries, resulting in improved accuracy and generalization compared to baselines.
Incorporating symmetries can lead to highly data-efficient and generalizable models by defining equivalence classes of data samples related by transformations. However, characterizing how transformations act on input data is often difficult, limiting the applicability of equivariant models. We propose learning symmetric embedding networks (SENs) that encode an input space (e.g. images), where we do not know the effect of transformations (e.g. rotations), to a feature space that transforms in a known manner under these operations. This network can be trained end-to-end with an equivariant task network to learn an explicitly symmetric representation. We validate this approach in the context of equivariant transition models with 3 distinct forms of symmetry. Our experiments demonstrate that SENs facilitate the application of equivariant networks to data with complex symmetry representations. Moreover, doing so can yield improvements in accuracy and generalization relative to both fully-equivariant and non-equivariant baselines.