Capsule network

Overview

A capsule network replaces scalar neurons with vectors ("capsules"). The length of a higher-level capsule's vector is read as how likely an entity is present, and its direction encodes properties such as pose [arxiv:2101.12491 §1]. The open question for any capsule network is how each lower-level capsule is assigned to higher-level capsules. The original CapsNet answers this with iterative dynamic routing, which repeats an agreement check and costs compute and parameters 2101.12491 §1 2101.12491 §3.1. Every variant shares the same skeleton: a convolutional stem, a primary-capsule layer, a routing layer, a squash non-linearity, and a margin loss on capsule lengths 2101.12491 §3 2101.12491 §3.2.

Variants

Iterative dynamic routing (original CapsNet)

The original design routes by repeated agreement between predictions. It is the baseline that the other variants try to replace. The card of Efficient-CapsNet reports its size as 6800K parameters in the comparison table 2101.12491 tab. 1. No card for the original paper exists in this wiki yet, so its method is not described here in detail; the cited numbers are the Efficient-CapsNet authors' report of it.

Self-attention routing (Efficient-CapsNet)

Routing is done in one non-iterative step. Lower capsules are mapped to one prediction per higher capsule. Predictions of each higher capsule are compared pairwise with a scaled dot product, the scores are summed over lower capsules, and a softmax over higher capsules gives the coupling weights. Learned log priors are added to these weights before the weighted sum and squash 2101.12491 §3.1 2101.12491 eq. 3-6.

The network is small: four convolutions with batch normalisation and ReLU, then a depthwise primary-capsule layer and one digit-capsule layer, about 161K parameters on MNIST 2101.12491 §3.1 2101.12491 tab. 2. Card: Efficient-CapsNet: Capsule Network with Self-Attention Routing.

When to use which

Only one card currently covers this concept, so the table compares the original routing with the single self-attention variant, using the paper's own statements.

The paper's headline "state of the art" claim rests on ensembles, not on a single model, so choose the variant on single-model numbers only after checking the caveats in the card 2101.12491 tab. 3.

Implementation sketch

The shared core is the routing step plus squash. Self-attention routing differs from dynamic routing in the agreement and coupling steps; the rest is common. Pseudocode is in our own words.

# Squash: keep the direction, map the length into (0, 1)  [arxiv:2101.12491 eq. 2]
def squash(v):
    length = norm(v)
    return (1 - exp(-length)) * v / length

# One routing layer, self-attention variant  [arxiv:2101.12491 §3.1] [eq. 3-6]
def self_attention_routing(lower, weights, log_priors, d):
    # One prediction per (lower, higher) pair  [eq. 3]
    pred = predict(lower, weights)
    # Pairwise scaled dot products among predictions of each higher capsule  [eq. 5]
    # The divisor d is the paper's d^l; see Contradictions
    agree = gram(pred) / sqrt(d)
    # Sum over lower capsules, then softmax over higher capsules  [eq. 6]
    coupling = softmax(sum(agree, over=lower), axis=higher)
    # Add learned log priors to the coupling weights  [eq. 4]
    coupling = coupling + log_priors
    # Weighted sum of predictions, then squash  [eq. 4] [eq. 2]
    return squash(sum(coupling * pred, over=lower))

# Training loss for one example  [arxiv:2101.12491 §3.2] [eq. 7]
def loss(lengths, target, recon, image, r):
    # Margin terms: 0.9 for the true class, 0.1 for others, weight 0.5  [eq. 7]
    margin = sum(target * max(0, 0.9 - lengths)**2
                 + 0.5 * (1 - target) * max(0, lengths - 0.1)**2)
    # Reconstruction error added with weight r  [§3.2]
    return margin + r * mse(recon, image)

Notes on where implementations diverge:

• The decoder input during training is a design choice the paper does not make explicit; the card describes the official code's rule 2101.12491 §3.2.
• Input convolutions and primary-capsule kernel sizes depend on input size; shapes must be recomputed for each dataset 2101.12491 §4.1.

Contradictions

• Routing divisor. The paper's d^l is "the individual dimension of the l-th layer" 2101.12491 §3.1 2101.12491 eq. 5. The official code divides by the digit layer's output dimension (16), while the input capsules have dimension 8. Both are recorded on the card's Paper vs code section; the paper does not resolve it.
• Convolution stride. The paper states stride 1 for every convolution 2101.12491 §3.1 2101.12491 eq. 1. The official code uses stride 2 in some layers for MNIST, smallNORB and MultiMNIST (card: Paper vs code).
• Training schedule. The paper gives per-dataset epochs (100, 200, about 100), batch sizes (16, 16, 64) and decay rates 2101.12491 §4.2 2101.12491 §4.3 2101.12491 §4.4. The official code uses one setting for all datasets (150 epochs, batch 16, decay 0.97) (card: Paper vs code).
• Validation data. The paper selects hyperparameters on "a small percentage" of validation data 2101.12491 §4.2. The official code passes the test set as validation data and keeps the best checkpoint on it (card: Pitfalls).
• Single-model MNIST. The paper's table lists CapsNet with reconstruction at 0.25 ± 0.005 and HFCs at 0.25 ± 0.0002, with Efficient-CapsNet at 0.26 ± 0.0002 2101.12491 tab. 2. The text presents the method as state of the art, but that claim rests on the 30-network ensemble at 0.16 2101.12491 tab. 3.
• Parameter counts. The MNIST table gives 161K 2101.12491 tab. 2, the smallNORB table 151K 2101.12491 tab. 4, the MultiMNIST text 154K 2101.12491 §4.4, and the abstract "about 160K". The paper does not explain the differences.
• smallNORB wording. The text calls the two-network result a "test accuracy of 1.23" 2101.12491 §4.3, while the surrounding numbers are error rates. The card reads it as an error of 1.23%.

Papers

• Efficient-CapsNet: Capsule Network with Self-Attention Routing: replaces iterative dynamic routing with one self-attention routing step in a network of about 161K parameters, reporting MNIST, smallNORB and MultiMNIST errors.