Wasserstein Embedding for Capsule Learning

Capsule networks (CapsNets) aim to parse images into a hierarchical component structure that consists of objects, parts, and their relations. Despite their potential, they are computationally expensive and pose a major drawback, which limits utilizing these networks efficiently on more complex datasets. The current CapsNet models only compare their performance with the capsule baselines and do not perform at the same level as deep CNN-based models on complicated tasks. This paper proposes an efficient way for learning capsules that detect atomic parts of an input image, through a group of SubCapsules, upon which an input vector is projected. Subsequently, we present the Wasserstein Embedding Module that first measures the dissimilarity between the input and components modeled by the SubCapsules, and then finds their degree of alignment based on the learned optimal transport. This strategy leverages new insights on defining alignment between the input and SubCapsules based on the similarity between their respective component distributions. Our proposed model, (i) is lightweight and allows to apply capsules for more complex vision tasks; (ii) performs better than or at par with CNN-based models on these challenging tasks. Our experimental results indicate that Wasserstein Embedding Capsules (WECapsules) perform more robustly on affine transformations, effectively scale up to larger datasets, and outperform the CNN and CapsNet models in several vision tasks.