胶囊网络

相关源文件

• labml_nn/adaptive_computation/ponder_net/__init__.py
• labml_nn/capsule_networks/__init__.py
• labml_nn/capsule_networks/mnist.py
• labml_nn/experiments/cifar10.py
• labml_nn/sketch_rnn/__init__.py

本文档提供了基于Hinton等人论文《Dynamic Routing Between Capsules》的胶囊网络（CapsNet）的注释实现。胶囊网络是一种深度学习的架构方法，旨在解决卷积神经网络的一些基本局限性，特别是在保持空间层级和视点不变表示方面。

该实现包括了CapsNet的核心组件、动态路由算法，以及用于MNIST数字分类的实际应用。有关此存储库中其他视觉模型实现的更多信息，请参阅SketchRNN。

1. 胶囊网络的核心概念

传统神经网络使用来自神经元的标量输出，而胶囊网络使用向量。每个胶囊输出一个向量，其中

• 向量的长度代表实体存在的概率
• 向量的方向编码了实体的实例化参数

这种方法使胶囊网络能够比传统CNN更好地保留特征之间的空间关系，并实现更鲁棒的视点不变性。

来源：labml_nn/capsule_networks/__init__.py10-26 labml_nn/capsule_networks/mnist.py7-13

2. 实现架构

胶囊网络的实现由三个主要组件组成：Squash激活函数、动态路由算法和Margin损失函数。

来源：labml_nn/capsule_networks/__init__.py10-26 labml_nn/capsule_networks/mnist.py29-97

2.1 Squash激活函数

Squash函数是一种专为胶囊网络设计的非线性激活函数。它将向量归一化，使其长度介于0到1之间，同时保留其方向。

$$\mathbf{v}_j = \frac{{\lVert \mathbf{s}_j \rVert}^2}{1 + {\lVert \mathbf{s}_j \rVert}^2} \frac{\mathbf{s}_j}{\lVert \mathbf{s}_j \rVert}$$

在计算向量范数时，实现通过添加一个小的epsilon来处理数值稳定性。

来源：labml_nn/capsule_networks/__init__.py39-70

2.2 动态路由算法

路由算法是胶囊网络的一项核心创新，它决定了信息如何在胶囊层之间传递。它实现了一种迭代的、基于协议的路由机制。

路由器使用权重矩阵和动态调整的耦合系数，将一个层的胶囊映射到下一层。

1. 输入胶囊（u_i）通过权重矩阵（W_ij）进行变换，以生成预测（û_j|i）。
2. 最初，路由逻辑（b_ij）被设置为零。
3. 对于指定的迭代次数：
   • 耦合系数（c_ij）通过在逻辑上应用softmax来计算。
   • 预测的加权和构成了更高级别胶囊的输入。
   • 将Squash函数应用于加权和。
   • 根据预测和输出之间的一致性来更新路由逻辑。

来源：labml_nn/capsule_networks/__init__.py73-133

2.3 Margin损失函数

胶囊网络使用一种特殊的损失函数，该函数强制规定输出胶囊长度之间的间隔。

$$\mathcal{L}_k = T_k \max(0, m^{+} - \lVert\mathbf{v}_k\rVert)^2 + \lambda (1 - T_k) \max(0, \lVert\mathbf{v}_k\rVert - m^{-})^2$$

其中

• T_k 如果类别k存在则为1，否则为0。
• m+ 是正间隔（通常为0.9）。
• m- 是负间隔（通常为0.1）。
• λ 是一个向下加权的因子（通常为0.5），以防止胶囊的长度在训练早期收缩。

第一项对正确类别的胶囊长度小于m+的情况进行惩罚，而第二项对错误类别的胶囊长度大于m-的情况进行惩罚。

来源：labml_nn/capsule_networks/__init__.py136-186

3. MNIST 分类器实现

该存储库在MNISTCapsuleNetworkModel中包含了一个用于MNIST数字分类的胶囊网络的具体实现。

3.1 模型架构

MNIST CapsNet模型遵循以下结构：

1. 第一个卷积层：256个9x9的滤波器，ReLU激活。
2. 第二个卷积层：32x8 = 256个9x9的滤波器，步长为2。
3. 主要胶囊：将输出重塑为32x6x6个胶囊，每个胶囊有8个维度。
4. 数字胶囊：10个胶囊（每个数字一个），每个胶囊有16个维度。
5. 解码器网络：从数字胶囊重构输入图像（16x10 -> 512 -> 1024 -> 784）。

层	输出形状	参数
输入	[batch_size, 1, 28, 28]	-
Conv1	[batch_size, 256, 20, 20]	9x9 核，256 通道
Conv2	[batch_size, 256, 6, 6]	9x9 核，步长为 2
PrimaryCaps	[batch_size, 32x6x6, 8]	重塑和置换
DigitCaps	[batch_size, 10, 16]	动态路由
解码器	[batch_size, 784]	全连接层
输出	[batch_size, 28, 28]	重塑为图像

来源：labml_nn/capsule_networks/mnist.py29-97

3.2 路由过程

在MNIST实现中，路由发生在主要胶囊和数字胶囊之间。

路由算法将1152个（32x6x6）主要胶囊连接到10个数字胶囊，每个主要胶囊有8个维度，每个数字胶囊有16个维度。

来源：labml_nn/capsule_networks/mnist.py46-50

3.3 重构作为正则化

MNIST CapsNet实现包含一个解码器网络，该网络从数字胶囊中重构输入图像。这充当了一种正则化形式。

1. 仅使用与正确数字对应的胶囊（其他胶囊被屏蔽）。
2. 解码器由三个全连接层组成（16x10 -> 512 -> 1024 -> 784）。
3. 重构损失（MSE）乘以0.0005，并添加到Margin损失中。

来源：labml_nn/capsule_networks/mnist.py84-95

4. 训练过程

MNIST胶囊网络的训练过程包括：

1. 前向传播:

   • 通过卷积层变换输入。
   • 创建主要胶囊。
   • 路由到数字胶囊。
   • 计算重构。

2. 损失计算:

   • 用于分类的Margin损失。
   • 用于重构的MSE损失（乘以0.0005）。
   • 总损失 = Margin损失 + 重构损失

3. 预测:

   • 每个数字胶囊的长度代表该数字的概率。
   • 具有最长胶囊向量的数字是预测结果。

该实现使用Adam优化器，学习率为1e-3。

来源：labml_nn/capsule_networks/mnist.py118-153

5. 优点和局限性

5.1 胶囊网络的优点

• 更好地处理特征之间的空间层级。
• 与CNN相比，改进了视点不变性。
• 对于某些任务，需要更少的参数和训练数据即可达到同等性能。
• 更具可解释性的内部表示。

5.2 局限性

• 由于路由算法，计算量大。
• 比传统CNN更复杂。
• 对更大、更复杂的数据集的可扩展性有限。
• 仍然是一个活跃的研究领域，持续发展中。

来源：labml_nn/capsule_networks/__init__.py10-26

6. 使用示例

要运行MNIST胶囊网络示例，请执行以下操作：

来源：labml_nn/capsule_networks/mnist.py162-178

该实现提供了一种清晰、模块化的方法来理解和实验胶囊网络，胶囊网络仍然是深度学习研究中一项重要的架构创新。