[Paper Review] Capsule Routing via Variational Bayes

Fabio De Sousa Ribeiro, Georgios Leontidis+1 more|arXiv (Cornell University)|May 27, 2019

Advanced Neural Network Applications25 references66 citations

TL;DR

This paper introduces Variational Bayes capsule routing (VB routing) for CapsNets, formulating routing as a mixture of transforming Gaussians and enabling a Capsule-VAE; it achieves state-of-the-art on smallNORB with ~50% fewer capsules and strong results on other datasets.

ABSTRACT

Capsule networks are a recently proposed type of neural network shown to outperform alternatives in challenging 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. In this paper, we propose a new capsule routing algorithm derived from Variational Bayes for fitting a mixture of transforming gaussians, and show it is possible transform our capsule network into a Capsule-VAE. Our Bayesian approach addresses some of the inherent weaknesses of MLE based models such as the variance-collapse by modelling uncertainty over capsule pose parameters. We outperform the state-of-the-art on smallNORB using 50% fewer capsules than previously reported, achieve competitive performances on CIFAR-10, Fashion-MNIST, SVHN, and demonstrate significant improvement in MNIST to affNIST generalisation over previous works.

Motivation & Objective

• Motivate and improve capsule networks by replacing traditional routing with a probabilistic Bayesian routing scheme.
• Develop a variational Bayes formulation for routing transforming Gaussian capsules to reduce variance collapse and improve stability.
• Demonstrate that the approach achieves competitive or state-of-the-art results on smallNORB, CIFAR-10, Fashion-MNIST, SVHN, and MNIST-affNIST generalization.
• Provide practical guidelines for initialization, normalization, and training to assist practitioners.

Proposed method

• Model capsules as a mixture of transforming Gaussians with per-capsule transform matrices W_ij.
• Use variational inference with a factorised q(z, pi, mu, Lambda) to infer routing assignments and capsule parameters.
• Represent votes V_j|i = M_i W_ij and form higher-layer pose M_j via Bayesian aggregation with responsibilities R_ij.
• Place conjugate priors (Dirichlet for pi, Gaussian-Wishart for mu, Lambda) and derive closed-form VB updates (Algorithm 1).
• Activate parent capsules a_j via a logistic function that uses prior evidence and posterior entropy to measure vote agreement (Eq. 14).
• Optionally convert the CapsNet into a Capsule-VAE by sampling from q*(mu_j, Lambda_j) to form latent codes (Equations 15–16).

Experimental results

Research questions

• RQ1 Does variational Bayes routing improve training stability and performance over EM and Dynamic routing in CapsNets?
• RQ2 Can a VB-based capsule routing framework reduce the variance collapse issue inherent in MLE-based mixture models and allow flexible sparsity via priors?
• RQ3 How does the VB routing approach influence performance and generalization across standard benchmarks (smallNORB, CIFAR-10, Fashion-MNIST, SVHN, MNIST-affNIST) and viewpoint generalization?

Key findings

• Achieves 1.55% test error on smallNORB (1.6%±0.06 over 5 runs), surpassing previous state-of-the-art 1.8%.
• Uses ~50% fewer capsules than prior methods to reach competitive results.
• Fashion-MNIST best test error 5.15% (5.2%±0.07).
• SVHN best test error 3.87% (3.9%±0.06).
• CIFAR-10 best test error 11.14% (11.2%±0.09); deeper variants reach 7.8% with more layers.
• Demonstrates MNIST to affNIST generalisation improvements over prior works.
• Shows VB routing provides stable training and competitive or superior performance across multiple datasets with fewer parameters.