Gauge Equivariant Capsule Networks

Gauge equivariant networks are emerging for equivariant representation learning on manifolds. However, existing gauge equivariant networks suffer from the dilemma between restricted SO(2) irreducible representations with equivariance guarantee to arbitrary continuous gauge transformations and enhanced C N regular representations with limited equivariance to finite discretized gauge transformations. In this paper, we propose a gauge equivariant capsule network (GECN) for manifolds that guarantees mathematically provable gauge equivariance in SO(2) without sacrificing the expressive capability. We develop a novel gauge equivariant capsule (GEC) that leverages the spatio-capsule dynamic routing algorithm to jointly exploit the capsule-wise and spatial dependencies in a local region of manifolds and achieve gauge equivariant representation based on real irreducible representations of SO(2). Furthermore, we propose a gauge equivariant capsule network (GECN) for manifolds by designing initial pose extractor, GEC-based residual block, and gauge equivariant pooling and unpooling based on GECs. We show that GECN is general to unify existing gauge equivariant networks and demonstrate in theory that it is equivariant to continuous gauge transformations in SO(2). Experimental results demonstrate that GECN achieves state-of-the-art performance in the tasks of deformable shape classification and segmentation and classification of spherical image and 3-D objects without principal curvature direction.