Capsule Routing via Variational Bayes

A capsule network is a recently proposed type of neural network, which has been shown to outperform alternatives in difficult shape recognition tasks. In capsule networks, scalar neurons are replaced with capsule vectors or matrices, whose entries represent different properties of objects. The relationships between objects and their parts are learned via trainable viewpoint-invariant transformation matrices, and the presence of a given object is decided by the level of agreement among votes from its parts. This interaction occurs between capsule layers and is a process called routing-by-agreement. Although promising, capsule networks remain underexplored by the community, and in this paper we present a new capsule routing algorithm derived from Variational Bayes for fitting a mixture of transforming gaussians. Our Bayesian approach addresses some of the inherent weaknesses of EM routing such as the variance-collapse by modelling uncertainty over the capsule parameters in addition to the routing assignment probabilities. We achieve competitive performance on 4 public domain datasets, and outperform the state-of-the-art performance on smallNORB using approximately 50% fewer capsules than previously reported.