Capsule Routing via Variational Bayes

Abstract

Capsule networks are a recently proposed type of neural network shown to\noutperform alternatives in challenging shape recognition tasks. In capsule\nnetworks, scalar neurons are replaced with capsule vectors or matrices, whose\nentries represent different properties of objects. The relationships between\nobjects and their parts are learned via trainable viewpoint-invariant\ntransformation matrices, and the presence of a given object is decided by the\nlevel of agreement among votes from its parts. This interaction occurs between\ncapsule layers and is a process called routing-by-agreement. In this paper, we\npropose a new capsule routing algorithm derived from Variational Bayes for\nfitting a mixture of transforming gaussians, and show it is possible transform\nour capsule network into a Capsule-VAE. Our Bayesian approach addresses some of\nthe inherent weaknesses of MLE based models such as the variance-collapse by\nmodelling uncertainty over capsule pose parameters. We outperform the\nstate-of-the-art on smallNORB using 50% fewer capsules than previously\nreported, achieve competitive performances on CIFAR-10, Fashion-MNIST, SVHN,\nand demonstrate significant improvement in MNIST to affNIST generalisation over\nprevious works.\n