[Paper Review] A Context-aware Capsule Network for Multi-label Classification

Sameera Ramasinghe, C.D. Athuralya+1 more|arXiv (Cornell University)|Oct 15, 2018

Domain Adaptation and Few-Shot Learning13 references4 citations

TL;DR

This paper proposes a context-aware capsule network that enhances multi-label image classification by introducing trainable routing weight initialization, a conditional random field (CRF) to model semantic relationships between primary capsule activations, and a Cholesky-based correlation module to learn dataset-wide priority schemes. The model achieves a 56.71 mAP on ADE20K, a 14.33-point improvement over the original CapsNet.

ABSTRACT

Recently proposed Capsule Network is a brain inspired architecture that brings a new paradigm to deep learning by modelling input domain variations through vector based representations. Despite being a seminal contribution, CapsNet does not explicitly model structured relationships between the detected entities and among the capsule features for related inputs. Motivated by the working of cortical network in human visual system, we seek to resolve CapsNet limitations by proposing several intuitive modifications to the CapsNet architecture. We introduce, (1) a novel routing weight initialization technique, (2) an improved CapsNet design that exploits semantic relationships between the primary capsule activations using a densely connected Conditional Random Field and (3) a Cholesky transformation based correlation module to learn a general priority scheme. Our proposed design allows CapsNet to scale better to more complex problems, such as the multi-label classification task, where semantically related categories co-exist with various interdependencies. We present theoretical bases for our extensions and demonstrate significant improvements on ADE20K scene dataset.

Motivation & Objective

• To address the limitations of CapsNet in multi-label classification, where structured relationships between object classes and capsule features are not explicitly modeled.
• To improve training convergence and performance by replacing equal routing weight initialization with a trainable, back-propagation-compatible initialization scheme.
• To exploit semantic dependencies between neighboring primary capsule activations using an end-to-end trainable Conditional Random Field (CRF).
• To learn a general, data-wide priority scheme between primary and decision capsules instead of per-sample routing, using a Cholesky transformation-based correlation module.
• To scale CapsNet to complex, real-world datasets like ADE20K with rich multi-label annotations and complex scene compositions.

Proposed method

• Proposes a trainable routing weight initialization method that uses convolutional kernels on low-level feature maps to generate dynamic α parameters, enabling back-propagation through initial routing weights.
• Introduces a densely connected CRF module that models conditional dependencies between primary capsule predictions, encouraging context-aware predictions by leveraging spatial and semantic relationships.
• Designs a Cholesky-based correlation module that computes a general priority scheme across the entire dataset by learning correlations between primary capsule outputs, improving global consistency.
• Employs a recursive function fρ to combine primary capsule predictions using learned α parameters, with differentiable operations to ensure end-to-end training.
• Uses Equation 10 to recursively aggregate primary capsule predictions with learned attention weights, where α is derived from convolved feature maps to adaptively weight contributions.
• Applies the proposed architecture to ADE20K, using mean average precision (mAP) with a 0.5 threshold as the evaluation metric across 150 object classes.

Experimental results

Research questions

• RQ1 Can trainable initial routing weights improve convergence speed and final performance in CapsNet for multi-label classification?
• RQ2 To what extent can a CRF-based module enhance prediction accuracy by modeling semantic relationships between neighboring primary capsule activations?
• RQ3 Can a dataset-wide correlation module that learns a general priority scheme between primary and decision capsules outperform per-sample routing in complex multi-label tasks?
• RQ4 How do the proposed components—trainable routing, CRF, and correlation module—complement each other in improving overall mAP on ADE20K?
• RQ5 Can CapsNet be effectively scaled to large, complex, multi-label datasets like ADE20K, where object co-occurrence and spatial dependencies are prevalent?

Key findings

• The proposed trainable routing weight initialization enables faster convergence, stabilizing validation mAP by the 9th epoch compared to the 15th epoch with equal initialization.
• The full model with all components (RW + CRF + CORR) achieves a mAP of 56.71 on ADE20K, a 14.33-point improvement over the original CapsNet’s 42.38 mAP.
• The CRF module alone contributes a 10.12-point mAP gain over the original CapsNet, demonstrating its effectiveness in capturing contextual dependencies.
• The correlation module adds an additional 4.21-point mAP improvement when stacked on top of the CRF, showing its complementary role in learning global priority structures.
• The ablation study confirms that each component—trainable routing, CRF, and correlation module—provides distinct and cumulative improvements, validating their individual and synergistic contributions.