[解读] Improving the Robustness of Capsule Networks to Image Affine Transformations

本文从理论和实验上表明了动态路由机制对胶囊网络的泛化性和对数据仿射变换的健壮性都没有起到太大的作用, 并且提出了两个小改进, 一个是对所有胶囊使用同一个矩阵进行线性变换, 另一个是取消了动态路由, 使用平均池化处理前一层的胶囊.

论文链接: https://arxiv.org/abs/1911.07968v3

相关工作

有很多工作致力于改进路由迭代算法. 例如文献 [27] 将已有的路由算法归纳为一个加权的核密度估计框架, 并且提出了两个使用不同优化策略的快速路由算法. 文献 [6] 提出了一种基于注意力机制的路由算法, 只需要进行快速的前向传递.

除此以外, 一致性系数 $a_{ij}$ 不仅可以通过内积来得到, 也可以通过基于正态分布假设 [10, 2] 或距离度量 [16] 来计算.

也有一些改进胶囊计算复杂度的研究工作, 例如文献 [25] 把路由机制改为一种优化问题, [17] 使用两个分支算法来估计路由系数.

本文的工作

首先从理论上分析路由算法在胶囊网络中所起到的作用. 算法流程可以分为以下几步:
$$\begin{array}{rl}{\mathbf{s}}_j^{(t)}& =\sum _i^N{c}_{ij}^{(t)}{\hat{\mathbf{u}}}_{j\mid i}\\ {\mathbf{v}}_j^{(t)}& =g\left({\mathbf{s}}_j^{(t)}\right)\\ c_{ij}^{(t+1)}& =\frac{\exp \left(b_{ij}+{\sum }_{r=1}^t{\mathbf{v}}_j^{(r)}{\hat{\mathbf{u}}}_{j\mid i}\right)}{{\sum }_k\exp \left(b_{ik}+{\sum }_{r=1}^t{\mathbf{v}}_k^{(r)}{\hat{\mathbf{u}}}_{k\mid i}\right)}.\end{array}$$
上标 $t$ 代表迭代次数, 假设总共有 $K$ 次迭代, 胶囊网络最后的输出可以表示为
$$\mathbf{y}=\left(\Vert {\mathbf{v}}_1^{(K)}\Vert ,\cdots ,\Vert {\mathbf{v}}_M^{(K)}\Vert \right).$$
即总共有 $M$ 个类别, 损失函数记为 $\mathcal{L}(\mathbf{y},\mathbf{t})$, 其中 $\mathbf{t}$ 为目标类别. 通过链式法则, 可以得到以下偏导数
$$\frac{\partial \mathcal{L}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}=\frac{\partial \mathcal{L}}{\partial {\mathbf{v}}_m^{(K)}}\frac{\partial {\mathbf{v}}_m^{(K)}}{\partial {\mathbf{s}}_m^{(K)}}{c}_{im}^{(K)}+\sum _{j=1}^M\frac{\partial \mathcal{L}}{\partial {\mathbf{v}}_j^{(K)}}\frac{\partial {\mathbf{v}}_j^{(K)}}{\partial {\mathbf{s}}_j^{(K)}}{\hat{\mathbf{u}}}_{j\mid i}\frac{\partial c_{ij}^{(K)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}.$$
可以发现, 上式中第二项是主要的计算负担, 对其中的 $\frac{\partial c_{ij}^{(K)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}$ 展开可得到:
$$\begin{array}{rl}\frac{\partial c_{ij}^{(K)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}=c_{ij}^{(K)}\left(1-c_{ij}^{(K)}\right)& \frac{\partial B_{ij}^{(K-1)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}\\ & -\sum _{k=1\&k\ne j}^M{c}_{ij}^{(K)}{c}_{ik}^{(K)}\frac{\partial B_{ik}^{(K-1)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}.\end{array}$$
进一步展开上式到第一次路由, 会发现只有一项梯度不为 0, 并且因为耦合系数 $c\in (0,1)$, 所以经过多次迭代, 最终将有 $\frac{\partial c_{ij}^{(K)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}\approx 0$.

最后, 作者证明了以下结论(详见附录):
$$\frac{\partial \mathcal{L}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}\approx C\cdot \frac{\partial \mathcal{L}}{\partial {\mathbf{v}}_m^{(K)}}\frac{\partial {\mathbf{v}}_m^{(K)}}{\partial {\mathbf{s}}_m^{(K)}}{c}_{im}^{(K)}.$$
其中 $C$ 是一个常数, 并且可以被整合进优化过程中, 如前面所说的, 这意味着 $\frac{\partial c_{ij}^{(K)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}$ 可以被忽略, 路由算法可以省略掉.

作者还从实验上表明了理论的结果.

上面是梯度反向传递时分析的结果, 下面作者从正向传递的角度考虑. 整个路由算法可以视为一个函数, 将初级胶囊映射为耦合系数, 即 $\hat{\mathbf{u}}\to {\mathbf{C}}^{(\hat{K})}=\left\{c_{ij}^{(K)}\right\}$. 不失一般性, 现在假设目标类别为第 $M$ 类, 最优的耦合系数可以通过以下优化问题得到:
$$\begin{array}{rl}{\mathbf{C}}^{\ast }=\underset{\left\{c_{ij}\right\}}{\max }f(\hat{\mathbf{u}})& =\underset{\left\{c_{ij}\right\}}{\max }(\sum _i^N{c}_{iM}{\hat{\mathbf{u}}}_{M\mid i}g\left(\sum _i{c}_{iM}{\hat{\mathbf{u}}}_{M\mid i}\right)\\ & -\sum _j^{M-1}\sum _i^N{c}_{ij}{\hat{\mathbf{u}}}_{j\mid i}g\left(\sum _i{c}_{ij}{\hat{\mathbf{u}}}_{j\mid i}\right)).\end{array}$$
即最大化通往目标类别胶囊的耦合系数, 最小化通往其它类别胶囊的耦合系数. 有许多针对改进路由算法的研究工作, 它们的不同点也仅仅在于用不同的方法来估计 ${\mathbf{C}}^{\ast }$.

最优化的 ${\mathbf{s}}_j^{\ast }$ 不需要耦合系数也能够通过学习得到:
$${\mathbf{s}}_j^{\ast }=\sum _i^N{c}_{ij}^{\ast }{\hat{\mathbf{u}}}_{j\mid i}=\sum _i^N{c}_{ij}^{\ast }{\mathbf{W}}_{ij}^{\ast }{\mathbf{u}}_i=\sum _i^N{\mathbf{W}}_{ij}^{\prime }{\mathbf{u}}_i$$
也就是说, 耦合系数可以通过变换矩阵 $\mathbf{W}$ 来隐式地学习到.

接下来作者分析了胶囊网络对于仿射变换具有较差的健壮性.
$${\hat{\mathbf{u}}}_{j\mid i}={\mathbf{W}}_{ij}{\mathbf{u}}_i=t_{j\mid i}\left({\mathbf{u}}_i\right)=t_{j\mid i}\left(p_i\left({\mathbf{x}}_i\right)\right).$$
${\mathbf{x}}_i$ 是输入样本, 函数 $p$ 代表卷积的过程, 输出为胶囊, $t_{j\mid i}$ 是对胶囊的线性变换. 作者认为, 线性变换矩阵 $\mathbf{W}$ 只能对存在于 $\mathbf{x}$ 上的实体做出有意义的变换, 设 $t_{j\mid i}(\cdot )$ 的输入域为 ${\mathbb{U}}_i$, 对于经过仿射变换的样本 ${\mathbf{x}}^{\prime }$ , $\mathbf{p}({\mathbf{x}}^{\prime })$ 不在 ${\mathbb{U}}_i$ 中, 因而 $\mathbf{W}$ 无法做出有意义的线性变换.

于是作者做出了一个非常简单但有效的解决方法, 那就是对所有的胶囊使同一个线性变换. 即一个共享的线性变换 $\mathbf{W}$. 这样一来就可以进行有意义的变换, 因为这一个矩阵学习了从所有低层次胶囊到高层次胶囊的变换.

补充

补充材料 (3) 式的推导

原文结果:
$$\begin{array}{rl}\frac{\partial c_{ij}^{(K)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}=c_{ij}^{(K)}\left(1-c_{ij}^{(K)}\right)& \frac{\partial B_{ij}^{(K-1)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}\\ & +\sum _{k=1\&k\ne j}^M{c}_{ij}^{(K)}{c}_{ik}^{(K)}\frac{\partial B_{ij}^{(K-1)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}\end{array}$$
原文上式中有两处错误, 应改为
$$\begin{array}{rl}\frac{\partial c_{ij}^{(K)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}=c_{ij}^{(K)}\left(1-c_{ij}^{(K)}\right)& \frac{\partial B_{ij}^{(K-1)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}\\ & -\sum _{k=1\&k\ne j}^M{c}_{ij}^{(K)}{c}_{ik}^{(K)}\frac{\partial B_{ik}^{(K-1)}}{\partial {\hat{\mathbf{u}}}_{m\mid i}}\end{array}$$