Capsule Networks

Summary

Capsule Networks are a deep learning architecture introduced by Geoffrey Hinton and his team in 2017. Unlike classical convolutional neural networks (CNNs) that produce scalar activations, Capsule Networks manipulate activation vectors called capsules. Each capsule encodes not only the presence of a visual entity, but also its spatial properties — orientation, scale, position, and deformation — within a single continuous vector.

This conceptual innovation solves a fundamental weakness of CNNs: their inability to explicitly model hierarchical spatial relationships between detected features. Through the dynamic routing by agreement mechanism, Capsule Networks establish dynamic connections between capsules at different levels, allowing the network to learn a robust and coherent hierarchical representation of the visual scene.

Capsule Networks thus offer exceptional robustness to geometric transformations (rotations, translations, scale changes) without requiring the massive training data augmentation needed by conventional approaches.

Mathematical Principle

A capsule is a group of neurons whose output is a vector rather than a scalar. The length (norm) of this vector represents the probability of presence of a particular entity, while its orientation encodes the pose parameters of that entity.

The Squash Function

To ensure that the norm of the output vector is between 0 and 1 (interpretable as a probability), the squash function is applied:

v_j = (||s_j||² / (1 + ||s_j||²)) · (s_j / ||s_j||)

where s_j is the weighted sum of input predictions for capsule j, and v_j is the output vector after squashing. This function has two essential properties:

• For small-norm vectors, the result is nearly zero — the capsule is inactive.
• For large-norm vectors, the result tends toward the unit vector in the same direction — the capsule is fully active.

This non-linearity preserves the direction of the vector (which contains the pose information) while normalizing its length (which encodes the probability of presence).

Prediction and Coupling

Each lower-level capsule i generates a prediction for each higher-level capsule j:

û_j|i = W_ij · u_i

where u_i is the output vector of capsule i and W_ij is a learned weight matrix that transforms the pose of capsule i into the reference frame of capsule j.

Dynamic Routing by Agreement

The core mechanism of Capsule Networks is routing by agreement. The idea is elegant: a higher-level capsule should only receive information from lower-level capsules whose predictions are consistent with each other.

The iterative process works as follows:

1. Initialization: coupling coefficients b_ij are initialized to zero.
2. Weight computation: c_ij = softmax(b_ij) — coefficients are normalized for each capsule i across all capsules j.
3. Aggregation: s_j = Σ_i c_ij · û_j|i — each capsule receives a weighted combination of predictions.
4. Squash: v_j = squash(s_j) — application of the squash function.
5. Agreement update: b_ij ← b_ij + û_j|i · v_j — coupling is increased when the prediction and the final output are aligned.

This process is repeated a fixed number of iterations (typically 3). Result: capsules whose predictions are consistent with the actual output see their influence increase, while discordant predictions are progressively ignored. This is an emergent attention mechanism, fully differentiable and learned end-to-end.

Geometric Intuition

To understand why Capsule Networks represent a major conceptual advance, one must examine the intrinsic weakness of classical convolutional networks.

A traditional CNN applies pooling operations (max-pooling or average-pooling) that reduce the spatial resolution of information. This reduction achieves invariance: an object detected in a corner of the image or at the center will actively produce the same response. However, this invariance comes at an enormous cost — the loss of spatial relationships between detected features.

Imagine a human face. A classical CNN will learn to detect eyes, nose, and mouth as independent entities. If all these features are present in the image, the network will conclude it is a face. But consider the following absurdity: a face where the eyes are placed at the bottom and the mouth at the top. A classical CNN will probably still recognize a face, because the individual elements are present! It does not verify the spatial arrangement.

Capsule Networks solve this fundamental problem elegantly. Instead of producing a simple scalar indicating “an eye is present here,” each capsule produces a pose vector that encodes the position, orientation, scale, and deformation of the detected entity. Higher-level capsules receive these pose vectors and, through routing by agreement, verify whether the predictions of lower-level capsules are spatially consistent.

This is the philosophical difference between:

• Recognizing a face by its individual elements — classical CNN approach (presence of pieces, without verifying the assembly).
• Recognizing a face by the spatial arrangement of its elements — Capsule approach (presence AND consistency of the spatial configuration).

This approach is analogous to human perception. Our visual system does not recognize an object simply by the list of its components — it evaluates the configuration of those components. A triangle remains a triangle even if you rotate it, because our brain encodes the angular relationships between the vertices, not the absolute positions of each point. Capsule Networks attempt to reproduce this fundamental principle of perception.

Python Implementation with PyTorch

Here is a complete implementation of Capsule Networks, inspired by the original architecture of Sabour, Frosst, and Hinton (2017), adapted for training on MNIST.

Squash Function

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms

class Squash(nn.Module):
    """Squash function for Capsule Networks.

    v_j = (||s_j||² / (1 + ||s_j||²)) * (s_j / ||s_j||)

    This function normalizes the vector length between 0 and 1
    while preserving its direction.
    """
    def __init__(self, dim=-1, epsilon=1e-7):
        super().__init__()
        self.dim = dim
        self.epsilon = epsilon

    def forward(self, s):
        sq_norm = (s ** 2).sum(dim=self.dim, keepdim=True)
        scale = sq_norm / (1.0 + sq_norm)
        norm = torch.sqrt(sq_norm + self.epsilon)
        v = scale * (s / norm)
        return v

Dynamic Routing Layer

class DynamicRouting(nn.Module):
    """Dynamic routing by agreement mechanism.

    Implements the iterative routing algorithm where coupling
    coefficients b_ij are updated by: b_ij ← b_ij + û_j|i · v_j
    """
    def __init__(self, num_iterations=3):
        super().__init__()
        self.num_iterations = num_iterations

    def forward(self, u_hat, num_input_caps, num_output_caps):
        """
        Args:
            u_hat: prediction tensor [batch, num_input, num_output, dim_out]
            num_input_caps: number of input capsules
            num_output_caps: number of output capsules
        Returns:
            v: output capsule vectors [batch, num_output, dim_out]
        """
        batch_size = u_hat.size(0)
        dim_out = u_hat.size(-1)

        # Initialize coupling logits
        b = torch.zeros(batch_size, num_input_caps, num_output_caps,
                       device=u_hat.device)

        for iteration in range(self.num_iterations):
            # Compute coupling coefficients via softmax
            c = F.softmax(b, dim=2)

            # Weighted aggregation of predictions
            c_expand = c.unsqueeze(-1)
            s = (c_expand * u_hat).sum(dim=1)

            # Apply squash function
            sq_norm = (s ** 2).sum(dim=-1, keepdim=True)
            scale = sq_norm / (1.0 + sq_norm)
            norm = torch.sqrt(sq_norm + 1e-7)
            v = scale * (s / norm)

            # Update agreements (except on last iteration)
            if iteration < self.num_iterations - 1:
                b += (u_hat * v.unsqueeze(1)).sum(dim=-1)

        return v

Complete CapsNet Architecture

class CapsNet(nn.Module):
    """Complete Capsule Network for MNIST classification.

    Architecture inspired by Sabour, Frosst & Hinton (2017).
    Includes primary capsules, dynamic routing, and
    a reconstruction decoder for regularization.
    """
    def __init__(self, num_primary_caps=1152, num_digit_caps=10,
                 capsule_dim=16, routing_iterations=3, recon_hidden=512):
        super().__init__()
        self.num_primary_caps = num_primary_caps
        self.num_digit_caps = num_digit_caps
        self.capsule_dim = capsule_dim

        # Initial convolutional layer (basic feature extraction)
        self.conv1 = nn.Conv2d(1, 256, kernel_size=9, stride=1)
        self.relu = nn.ReLU()

        # PrimaryCaps layer — transforms scalar features
        # into capsule vectors (32 channels × 6×6 positions = 1152 capsules)
        self.primary_caps = nn.Conv2d(256, 32 * capsule_dim,
                                       kernel_size=9, stride=2)
        self.squash = Squash()
        self.routing = DynamicRouting(num_iterations=routing_iterations)

        # Weight matrices for predictions
        self.W = nn.Parameter(torch.randn(1, num_primary_caps,
                                           num_digit_caps, capsule_dim,
                                           capsule_dim) * 0.1)

        # Reconstruction decoder for regularization
        self.decoder = nn.Sequential(
            nn.Linear(num_digit_caps * capsule_dim, recon_hidden),
            nn.ReLU(inplace=True),
            nn.Linear(recon_hidden, recon_hidden * 2),
            nn.ReLU(inplace=True),
            nn.Linear(recon_hidden * 2, 28 * 28),
            nn.Sigmoid()
        )

    def forward(self, x):
        # Convolutional feature extraction
        x = self.relu(self.conv1(x))  # [batch, 256, 20, 20]

        # Generate primary capsules
        pc = self.primary_caps(x)  # [batch, 32*16, 6, 6]
        batch_size = pc.size(0)

        # Reshape into individual capsules
        pc = pc.view(batch_size, -1, self.capsule_dim)  # [batch, 1152, 16]
        pc = self.squash(pc)

        # Prediction: û = W · u
        u = pc.unsqueeze(2).unsqueeze(3)  # [batch, 1152, 1, 1, 16]
        u_hat = torch.matmul(self.W, u).squeeze(3)  # [batch, 1152, 10, 16]

        # Dynamic routing
        v = self.routing(u_hat, self.num_primary_caps, self.num_digit_caps)

        # Classification based on capsule norms
        classes = (v ** 2).sum(dim=-1)  # [batch, 10]

        # Reconstruction
        max_idx = classes.argmax(dim=1)
        mask = torch.zeros_like(v)
        mask[range(batch_size), max_idx, :] = v[range(batch_size), max_idx, :]
        reconstruction = self.decoder(mask.view(batch_size, -1))

        return classes, v, reconstruction

Margin Loss Function and Training

def margin_loss(classes, targets, lambda_loss, m_plus=0.9, m_minus=0.1):
    """Margin loss function for Capsule Networks.

    Encourages the target capsule to have a norm >= m_plus
    and non-target capsules to have a norm <= m_minus.
    """
    left = F.relu(m_plus - classes).pow(2)
    right = F.relu(classes - m_minus).pow(2)

    mask = F.one_hot(targets, num_classes=10).float()
    loss = mask * left + lambda_loss * (1 - mask) * right
    return loss.sum(dim=1).mean()


def train_epoch(model, dataloader, optimizer, device, lambda_loss=0.5,
                recon_weight=0.0005):
    """Train one epoch on MNIST."""
    model.train()
    total_loss = 0.0
    correct = 0
    total = 0

    for images, targets in dataloader:
        images, targets = images.to(device), targets.to(device)

        optimizer.zero_grad()
        classes, v, reconstruction = model(images)

        # Margin loss
        m_loss = margin_loss(classes, targets, lambda_loss)

        # Reconstruction loss (regularization)
        recon_loss = F.mse_loss(reconstruction,
                                images.view(images.size(0), -1))

        # Total loss
        loss = m_loss + recon_weight * recon_loss
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        predictions = classes.argmax(dim=1)
        correct += (predictions == targets).sum().item()
        total += targets.size(0)

    avg_loss = total_loss / len(dataloader)
    accuracy = correct / total * 100
    return avg_loss, accuracy


def evaluate_robustness(model, dataloader, device, angle_degrees=45):
    """Evaluates the model's robustness to rotations.

    Capsule Networks should maintain good performance
    even with rotated images, unlike classical CNNs.
    """
    model.eval()
    correct = 0
    total = 0

    with torch.no_grad():
        for images, targets in dataloader:
            images, targets = images.to(device), targets.to(device)

            # Rotate images
            rotated = torchvision.transforms.functional.rotate(
                images, angle_degrees
            )

            classes, _, _ = model(rotated)
            predictions = classes.argmax(dim=1)
            correct += (predictions == targets).sum().item()
            total += targets.size(0)

    return correct / total * 100

Key Hyperparameters

Influence of the Number of Routing Iterations

The number of routing iterations represents a fundamental trade-off between accuracy and computational cost. With a single iteration, routing is equivalent to a standard softmax — no feedback information between capsules. From three iterations, the mechanism generally converges to stable couplings. Beyond five iterations, accuracy gains become marginal while computation time increases linearly.

Capsule Dimension

A typical dimension of 8 to 16 is recommended for simple datasets like MNIST. For more complex visual tasks, higher dimensions (32 to 64) may be necessary to encode enough pose parameters. However, increasing the dimension also multiplies the size of the weight matrices W_ij, which directly impacts GPU memory.

Advantages and Limitations

Advantages

Preservation of spatial relationships: This is the fundamental advantage. Capsule Networks explicitly encode hierarchical relationships between visual entities. A face with misplaced eyes will not be recognized as a face, which aligns much more closely with human perception.

Robustness to transformations: Thanks to pose vectors, Capsule Networks are naturally robust to rotations, translations, and scale changes. Unlike CNNs that require massive data augmentation to learn this invariance, Capsule Networks generalize it through their vector representation.

Less training data required: The architecture’s inherent robustness to transformations requires fewer training examples to achieve performance comparable to CNNs, particularly on tasks where geometric transformations are frequent.

Explicit reconstruction: The reconstruction decoder forces capsules to encode semantically rich information. Visualizations of reconstructions show that capsules do indeed learn meaningful representations of visual entities.

Adaptive routing: Unlike the static connections of CNNs, dynamic routing adapts connections between layers based on the input content. Each image receives personalized routing.

Limitations

High computational cost: Iterative routing by agreement is significantly more expensive than a standard convolution. Each iteration requires matrix operations across all capsules, making training and inference considerably slower.

Difficulty scaling: Capsule Networks show excellent results on MNIST but struggle to surpass state-of-the-art CNNs on complex datasets like CIFAR-10 and ImageNet. The weight matrix W_ij grows quadratically with the number of capsules, making large-scale application nearly impractical with the original architecture.

Slow convergence: Dynamic routing introduces complex non-linearities that make optimization more difficult. Training often requires more epochs and careful weight initialization.

Sensitivity to hyperparameters: The number of routing iterations, capsule dimension, and margins m+ / m- are sensitive hyperparameters that require meticulous tuning for good performance.

Still emerging research: The original 2017 architecture has since been improved (Matrix Capsules, EM Routing), but no variant has reached the level of dominance of CNNs in the field of computer vision.

4 Concrete Use Cases

1. Medical Detection with Anatomical Relationships

In medical imaging (MRI, X-ray), the relative position of anatomical structures is crucial for diagnosis. A tumor located next to a specific organ does not have the same significance as a tumor of similar appearance located elsewhere. Capsule Networks, by explicitly encoding spatial relationships, offer a natural advantage for these tasks where spatial configuration carries essential diagnostic information.

2. Object Recognition with Viewpoint Variations

In augmented reality or mobile robotics applications, objects are observed from varying and changing angles. Capsule Networks, thanks to their pose vectors, can recognize the same object from different perspectives without requiring exhaustive training on every possible angle. This viewpoint generalization capability is particularly valuable in dynamic and unpredictable environments.

3. Fraud Detection and Graph Analysis

The routing-by-agreement concept applies beyond computer vision. In financial fraud detection, each transaction can be modeled as a capsule and a similar routing mechanism can be used to identify suspicious patterns based on consistency between connected entities. Fraudulent transactions form configurations that diverge from normal patterns detected by routing.

4. Complex Scene Analysis with Occlusion

In visual scenes where some objects are partially hidden (occlusion), Capsule Networks can infer the presence and pose of occluded objects through the learned spatial relationships between capsules. If capsules detect a torso, arms, and a partially visible head, routing by agreement can reinforce the hypothesis of a complete person even if some parts are hidden. This contextual inference capability is difficult to achieve with standard CNNs.

See Also

• What? Where? When? Mastering Python Programming: Complete Guide for Beginners
• Mastering Packet Sorting in Python: Complete Guide for Beginners