A Multi-prototype Capsule Network for Image Recognition with High Intra-class Variations

Abstract

Capsule networks are deep neural networks that perform a part-to-whole association and instantiate the parameters of a “whole” (e.g., a class) by searching the agreement of “parts”. These networks are based on grouping neurons into units called capsules. The activity of “part” capsules is propagated to the “whole” capsules in the next layer by a transformation matrix, instead of a scalar weight, as is typical in conventional networks. One of the challenges of capsule networks is training on complex data with high intra-class variations. The high intra-class variation in real-world datasets cannot be appropriately captured by considering just one representation for each class. In this paper, we present a new multi-prototype capsule network architecture and propose a soft competitive learning algorithm that guides the network towards learning the diversity that exists within classes. We show that by having multiple capsules for each class and using a soft competitive cross-entropy loss, specialized learning in capsule networks can occur in which each capsule represents a particular part of the feature space of the corresponding class. Discovering the intra-class prototypes makes it easier for the capsule network to deal with complex datasets. Our extensive experiments on MNIST, CEDAR, MCYT, C-Cube, and SVHN datasets show the superiority of the proposed approach. Results show that the proposed model achieves an average of 3.36% and 3.73% classification accuracy improvement compared to the original capsule networks and high-dimensional capsules, respectively.