SO(3)-Equivariant Representation Learning in 2D Images

Imaging physical objects that are free to rotate and translate in 3D is challenging. While an object's pose and location do not change its nature, varying them presents problems for current vision models. Equivariant models account for these nuisance transformations, but current architectures only model either 2D transformations of 2D signals or 3D transformations of 3D signals. Here, we propose a novel convolutional layer consisting of 2D projections of 3D filters that models 3D equivariances of 2D signals—critical for capturing the full space of spatial transformations of objects in imaging domains such as cryo-EM. We additionally present methods for aggregating our rotation-specific outputs. We demonstrate significant improvement on several tasks, including particle picking and pose estimation.