Efficient-CapsNet: Capsule Network with Self-Attention Routing

In plain words

Capsule networks are a way to build image models that group neurons into small vectors instead of single numbers. A vector's length is read as how likely an object part is to be present, and its direction carries details such as position and orientation. The hard part is deciding which lower-level vectors feed which higher-level ones. The original method, dynamic routing, works this out by repeating an agreement check, which costs time and parameters.

This paper proposes a single pass instead. The lower-level vectors are compared with each other, and a softmax over that comparison decides how strongly each one feeds each higher-level vector. The network is kept very small, about 161 thousand parameters, and is tested on three image sets: handwritten digits, toy objects and overlapping digit pairs. The authors report accuracy comparable to much larger capsule networks on these sets.

TL;DR

Efficient-CapsNet is a compact capsule network (about 161K parameters) for MNIST, smallNORB and MultiMNIST. Its routing is non-iterative: a self-attention style agreement score, a softmax over output capsules, and learned log priors replace repeated dynamic routing §3.1 eq. 3-6. Author-reported errors are 0.26% (MNIST, single model) and 2.54% (smallNORB) tab. 2 tab. 4. Check the Pitfalls section before reusing the numbers: the official code differs from the paper in several places, and model selection uses the test set.

Problem

Convolutional networks cope with viewpoint changes mostly by adding feature maps, data augmentation and more parameters, which makes them redundant for large datasets §1. Capsule networks are meant to generalise to new viewpoints with far fewer parameters, but the authors note that little work had examined their efficiency §1. The original CapsNet, which uses iterative dynamic routing, has 6800K parameters in the comparison table tab. 1. The question addressed here is how small a capsule network can be while keeping competitive accuracy §1 §4.

Method

The network has three parts: a convolutional stem, a primary-capsule layer, and a digit-capsule layer with routing §3 fig. 2.

1. Convolutional stem §3.1. Four valid convolutions, each followed by batch normalisation and a ReLU activation, map the input to a higher-dimensional tensor. The first kernel is 5 and the rest are 3, with 32, 64, 64 and 128 feature maps §3.1 eq. 1 §4.1. The paper states stride 1 for all layers, but the code uses stride 2 in the last layer (see Paper vs code).
2. Primary capsules §3.1 fig. 3. One depthwise convolution with linear activation and a kernel as large as the stem's output map turns the tensor into 16 capsules of dimension 8 for MNIST §4.1. Each capsule goes through the squash function eq. 2.
3. Self-attention routing §3.1 (the subsection is unnumbered in the source) fig. 4. The lower capsules are multiplied by a learned weight tensor to give one prediction per (lower, higher) pair eq. 3. The predictions of each higher capsule are compared pairwise, using a dot product scaled by a square root eq. 5. These scores are summed over the lower capsules and normalised with a softmax over the higher capsules eq. 6. Learned log priors are then added to the coupling coefficients, and each higher capsule is the coupling-weighted sum of predictions, passed through squash eq. 4 eq. 2.
4. Loss §3.2. A margin loss with separate weights for present and absent classes is summed over classes eq. 7. A decoder reconstructs the input from the output capsules, and its mean squared error is added with weight r §3.2 §4.1.

Equation numbers are counted from the LaTeX source, not from a compiled PDF.

Pseudocode for the routing step:

# Self-attention routing for one capsule layer  [§3.1] [eq. 3-6]
# Input:  lower capsules (B, N_low, d_low)
# Output: higher capsules (B, N_high, d_high)

def self_attention_routing(lower, weights, log_priors):
    # Predict each higher capsule from each lower one  [§3.1] [eq. 3]
    # weights has shape (N_low, N_high, d_low, d_high)
    predictions = predict(lower, weights)

    # Pairwise agreement among the predictions of each higher capsule  [eq. 5]
    # The divisor is sqrt(d_low) in the paper; see Pitfalls
    agreement = gram(predictions) / sqrt(d_low)

    # Sum each lower capsule's agreement over all lower capsules  [eq. 6]
    summed = sum(agreement, over_last_axis)

    # Softmax over the higher capsules, one distribution per lower  [eq. 6]
    coupling = softmax(summed, axis=higher)

    # Add learned log priors to the coupling coefficients  [§3.1] [eq. 4]
    coupling = coupling + log_priors

    # Coupling-weighted sum of predictions, then squash  [eq. 4] [eq. 2]
    outputs = sum(coupling * predictions, over=lower)
    return squash(outputs)

Squash and loss:

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

# Training loss for one image  [§3.2] [eq. 7] [§4.1]
# lengths and targets have one entry per class; targets are one-hot
def training_loss(lengths, targets, reconstruction, image):
    # Penalise short lengths for the true class  [eq. 7]
    present = targets * max(0, 0.9 - lengths) ** 2
    # Penalise long lengths for absent classes, with weight 0.5  [eq. 7]
    absent = 0.5 * (1 - targets) * max(0, lengths - 0.1) ** 2
    margin = sum(present + absent)
    # Reconstruction term, weighted by r = 0.392  [§3.2] [§4.1]
    error = mean((reconstruction - image) ** 2)
    return margin + 0.392 * error

Derived shapes (our arithmetic from the code's valid convolutions, not stated in the paper):

• MNIST, 28×28 input: 24, 22, 20, then 9 after the stride-2 layer; the primary kernel is 9×9 and gives 1×1×128, which reshapes to 16×8 efficient_capsnet_graph_mnist.py.
• smallNORB, 48×48×2 input: 21, 19, 17, then 8; the primary kernel is 8×8 efficient_capsnet_graph_smallnorb.py.
• MultiMNIST, 36×36 input: 32, 30, 14, then 6; the primary kernel is 5 with stride 2 and gives 1×1 efficient_capsnet_graph_multimnist.py.

Reconstruction decoder: two hidden fully connected layers of 512 and 1024 units, with a sigmoid output over the image size §4.2 efficient_capsnet_graph_mnist.py. The paper does not say which capsule is fed to the decoder during training; the code masks the true class (see Pitfalls).

Hyperparameters

Values are the paper's unless the Notes column says otherwise. "Not stated" means the paper does not give the value.

Results

Author-reported. Setup: MNIST (60000 train, 10000 test) §4.2; smallNORB (48600 stereo images, 5 classes) §4.3; MultiMNIST built from MNIST with 36×36 canvases §4.4; baselines are CapsNet, Matrix-CapsNet, AR-CapsNet, DA-CapsNet, HFCs and VB-Routing as listed in each table.

Headline numbers:

Table 1, computational cost, copied from the paper. Parameters are in thousands and operations are per batch in billions (author-reported; not reproduced here):

Table 2, MNIST test error (%). Asterisk and parentheses mark the authors' own runs of the baseline. Reconstruction column as in the paper:

Table 3, MNIST ensemble test error (%) tab. 3: Ciregan et al. 2012, 0.23; Wan et al. 2013, 0.21; RMDL 2018, 0.18; Byerly et al. 2020, 0.16; Efficient-CapsNet 2021, 0.16. The ensemble averages the networks whose test accuracy exceeds 0.9973 §4.2.

Table 4, smallNORB test error (%). Asterisk marks the authors' own runs:

Table 5, share of variance captured by the first PCA component of output capsules, per transformation tab. 5: random vectors 25.57 / 25.54 / 13.49; CapsNet 83.78 / 79.82 / 88.01; Efficient-CapsNet 89.69 / 87.28 / 88.75 (translation x, translation y, rotation). Figure 8 shows cumulative variance curves fig. 8.

Figures: the paper shows the architecture fig. 2, the primary-capsule construction fig. 3, the routing scheme fig. 4, MNIST reconstructions across methods fig. 5, misclassified MNIST digits with capsule lengths fig. 6, perturbation of the longest output capsule fig. 7 and the PCA curves fig. 8. Figure 1 illustrates translation invariance in CNNs fig. 1.

Paper vs code

Checked against the official repository at cd0b53379ea011c35ac90d41f3ffdc0ec4d6929b, code licensed Apache-2.0. Only the parts relevant to the three reported datasets were read; the notebooks were not run.

Discrepancies:

• Convolution stride. Paper: stride 1 for every convolution §3.1 eq. 1. Code: the last stem convolution has stride 2 in MNIST efficient_capsnet_graph_mnist.py and in MultiMNIST efficient_capsnet_graph_multimnist.py; the first conv in smallNORB has stride 2 with a 7×7 kernel efficient_capsnet_graph_smallnorb.py. The MNIST primary kernel is 9 efficient_capsnet_graph_mnist.py. With stride 1 the MNIST map would be 20×20, so the paper's rule for the kernel size would give a 20×20 kernel instead.
• Routing scale. Paper: divide by the square root of d^l eq. 5, where d^l is "the individual dimension of the l-th layer" §3.1. Code: divides by the square root of the digit layer's output dimension, 16 layers.py, with the layer built as FCCaps(10, 16) efficient_capsnet_graph_mnist.py. The input capsules of that layer have dimension 8, so the code does not use the input dimension.
• Routing otherwise agrees. The softmax runs over the higher capsules layers.py, the log priors are added after the softmax layers.py, and the output is a weighted sum passed through squash layers.py. The squash matches eq. 2 with a small epsilon layers.py.
• Margin loss agrees. Same parameters 0.9, 0.1, 0.5 tools.py and the same form as eq. 7 tools.py. The batch mean is taken after summing over classes.
• Reconstruction weight agrees. r = 0.392 config.json applied to the Keras mean squared error model.py.
• Decoder input. Paper: output capsules fed to the decoder, with no masking rule for the single-digit datasets §3.2. Code: during training the decoder receives the capsule of the true class; at test time it receives the predicted class efficient_capsnet_graph_mnist.py.
• Epochs, batch size and decay. Paper: MNIST 100 epochs, batch 16, decay 0.98; smallNORB 200 epochs, batch 16, decay 0.99; MultiMNIST about 100 epochs, batch 64, decay 0.97 §4.2 §4.3 §4.4. Code: one configuration for all datasets, 150 epochs, batch 16, decay 0.97 config.json.
• Learning-rate floor and plateau callback. Code applies a floor of 5e-5 to the decayed rate tools.py. A ReduceLROnPlateau callback is defined tools.py but is not returned in the callback list tools.py, so it is inactive. The paper mentions neither.
• Validation data. Paper: hyperparameters are chosen on "a small percentage" of validation data §4.2. Code: the test set is passed as validation data, and the checkpoint with the best validation accuracy is saved dataset.py tools.py.
• MNIST augmentation. Paper: "the same augmentation as Byerly et al." §4.2. Code: its own random rotation, shift, squish and erase operations on each training image pre_process_mnist.py. The paper lists none of them.
• MNIST scaling. Code divides by 256, not 255 pre_process_mnist.py. The paper does not state the scaling.
• smallNORB architecture. Paper: same structure as MNIST, with instance normalisation in place of batch normalisation, and ReLU §4.3 §4.1. Code: a 7×7 first conv with stride 2, LeakyReLU, instance normalisation after every conv, no batch normalisation, and a primary kernel of 8 efficient_capsnet_graph_smallnorb.py. The code also has a comment "there could be an error" beside the primary layer efficient_capsnet_graph_smallnorb.py.
• smallNORB normalisation. Code standardises each split with that split's own mean and standard deviation over the whole array, before resizing dataset.py. Training then applies brightness changes with delta up to ±2.0 and contrast 0.5 to 1.5 pre_process_smallnorb.py pre_process_smallnorb.py. The paper gives none of these values.
• MultiMNIST pairs. Paper: the two digits must be of different classes §4.4. Code, training: only the two source images must differ, so two digits of the same class can be paired pre_process_multimnist.py. Code, test: the second digit must have a different label pre_process_multimnist.py.
• MultiMNIST shift and batch. Paper: shift up to 4 px and batch 64 §4.4. Code: shift 6 and batch 16 config.json config.json. The padding of 4 px to 36×36 matches config.json.
• MultiMNIST architecture. Code uses two stride-2 convolutions and a stride-2 primary layer efficient_capsnet_graph_multimnist.py. The paper's general architecture description (stride 1) does not mention this.
• Pretrained weights. The repository ships trained weights for MNIST, smallNORB and MultiMNIST, plus the original CapsNet on MNIST . No script for the 30-trial statistics was located.

Pitfalls

• Test set used for selection. The checkpoint is chosen on the test set. The ensembles are also chosen on test accuracy: networks above 0.9973 for MNIST §4.2, the two lowest test errors for smallNORB §4.3, and networks above 0.9470 for MultiMNIST §4.4. The reported ensemble numbers are therefore optimistic. Report single-model numbers with a separate validation split if the comparison matters.
• Single-model MNIST is not the best in the paper's own table. CapsNet with reconstruction reports 0.25 ± 0.005 and HFCs 0.25 ± 0.0002 tab. 2; Efficient-CapsNet's 0.26 is close but higher. The "state-of-the-art" claim rests on the ensemble tab. 3.
• Ambiguous smallNORB wording. The text says the two-network result is a "test accuracy of 1.23" §4.3, while the surrounding numbers are errors. Read it as an error of 1.23% unless the paper is corrected.
• Routing scale. The paper's d^l is ambiguous for the digit layer (see Paper vs code). The code divides by the digit-capsule dimension layers.py. The divisor sets the scale of the logits before the softmax, so the choice changes the routing weights.
• Coupling sum. The coupling coefficients sum to one over higher capsules only before the log priors are added §3.1 layers.py. After the priors they no longer sum to one, so do not treat them as probabilities.
• No stop-gradient. The same predictions feed the agreement scores and the output sum. The code has no stop-gradient on either path layers.py. The paper does not discuss this.
• Reconstruction masking. The decoder sees the true class during training efficient_capsnet_graph_mnist.py. An implementation that masks by the predicted class during training will get different reconstructions and different gradients.
• Schedule and floor. The code's learning rate is 5e-4 × decay^epoch with a floor of 5e-5 tools.py. Epoch counts and decay differ by dataset in the paper but not in the code, so state which setting you use.
• Augmentation is part of the result. MNIST training uses rotation, shift, squish and erase pre_process_mnist.py; smallNORB uses brightness and contrast jitter on standardised images pre_process_smallnorb.py. These are not in the paper's text, so reproduced errors can differ from the paper's without any bug.
• Parameter counts differ across tables. The MNIST table gives 161K parameters, the smallNORB table 151K tab. 2 tab. 4, and the MultiMNIST text 154K §4.4. The abstract says about 160K. The paper does not explain the differences.
• Computational costs are author-reported. The OPS column in Table 1 was not reproduced.

Lineage

• Builds on: the original CapsNet with dynamic routing [sabour2017dynamic], which the paper replaces at the routing step §2 §3.1. No card in this wiki yet.
• Compared against: AR CapsNet [choi2019attention], Matrix-CapsNet with EM routing [hinton2018matrix], DA-CapsNet [huang2020capsnet], HFCs [byerly2020branching], VB-Routing [ribeiro2020capsule] and others listed in the Results section §2 §4.
• Superseded by: not known from content already in this wiki.