[Paper Review] Neural Network Encapsulation

Hongyang Li, Xiaoyang Guo+3 more|arXiv (Cornell University)|Aug 11, 2018

Advanced Neural Network Applications24 references4 citations

TL;DR

This paper proposes EncapNet, a capsule network architecture that replaces the computationally heavy iterative routing in capsule networks with a fast, one-time, supervised master-aide routing mechanism. It further introduces a feedback agreement unit using optimal transport divergence to regularize feature learning, achieving state-of-the-art performance on CIFAR-10/100, SVHN, and a subset of ImageNet with error rates of 3.10%, 24.01%, and 40.05% respectively under strong data augmentation.

ABSTRACT

A capsule is a collection of neurons which represents different variants of a pattern in the network. The routing scheme ensures only certain capsules which resemble lower counterparts in the higher layer should be activated. However, the computational complexity becomes a bottleneck for scaling up to larger networks, as lower capsules need to correspond to each and every higher capsule. To resolve this limitation, we approximate the routing process with two branches: a master branch which collects primary information from its direct contact in the lower layer and an aide branch that replenishes master based on pattern variants encoded in other lower capsules. Compared with previous iterative and unsupervised routing scheme, these two branches are communicated in a fast, supervised and one-time pass fashion. The complexity and runtime of the model are therefore decreased by a large margin. Motivated by the routing to make higher capsule have agreement with lower capsule, we extend the mechanism as a compensation for the rapid loss of information in nearby layers. We devise a feedback agreement unit to send back higher capsules as feedback. It could be regarded as an additional regularization to the network. The feedback agreement is achieved by comparing the optimal transport divergence between two distributions (lower and higher capsules). Such an add-on witnesses a unanimous gain in both capsule and vanilla networks. Our proposed EncapNet performs favorably better against previous state-of-the-arts on CIFAR10/100, SVHN and a subset of ImageNet.

Motivation & Objective

• To address the high computational complexity of iterative routing in capsule networks, which scales poorly with network depth.
• To improve feature representation by enabling lower capsules to contribute more effectively to higher capsule activation through a dual-branch routing mechanism.
• To reduce information loss in deep networks by introducing a feedback regularization mechanism based on optimal transport divergence.
• To enhance both capsule and vanilla CNN performance through a lightweight, trainable feedback unit that enforces agreement between layers.

Proposed method

• Introduce a master-aide routing mechanism: the master branch uses direct lower-layer capsules, while the aide branch searches for pattern variants across channels to replenish information.
• Train the master and aide branches in a one-time, supervised pass, replacing the iterative, unsupervised routing of traditional capsule networks.
• Design a feedback agreement unit that sends higher-layer capsule outputs back to lower layers as supervision.
• Use optimal transport divergence (Sinkhorn loss) to measure distributional distance between lower and higher capsule features, minimizing this distance during training.
• Integrate the feedback unit as a differentiable regularization loss during training, which is removed at inference.
• Implement the method in a convolutional capsule layer (capConv) that supports end-to-end training with standard backpropagation.

Experimental results

Research questions

• RQ1 Can a one-pass, supervised routing mechanism approximate the performance of iterative routing while drastically reducing computational cost?
• RQ2 Does introducing a feedback signal from higher to lower capsules improve feature learning and generalization in capsule networks?
• RQ3 Can optimal transport divergence serve as an effective regularization metric for inter-layer agreement in deep networks?
• RQ4 Does the proposed feedback mechanism improve performance not only in capsule networks but also in standard CNNs?

Key findings

• EncapNet achieves a test error rate of 3.10% on CIFAR-10, outperforming previous state-of-the-art methods under strong data augmentation.
• On CIFAR-100, EncapNet reaches 24.01% error rate, demonstrating significant improvement over prior capsule-based models.
• On SVHN, EncapNet achieves 1.52% error rate, showing strong generalization on digit recognition tasks.
• On a subset of ImageNet (h-ImageNet), EncapNet_v3++ achieves 40.05% top-1 error rate, surpassing the runner-up WRN model at 42.51%.
• The feedback agreement unit improves performance across both capsule and vanilla CNNs, indicating its generalizability as a regularization technique.
• Ablation studies confirm that the feature extractor, optimal transport divergence, and master-aide interaction are all essential for performance, with ablation leading to significant error increases.