可逆神经网络（Invertible Neural Networks）详细解析：让神经网络更加轻量化-技术圈

来源：PaperWeekly
本文约3600字，建议阅读7分钟
本文以可逆残差网络（The Reversible Residual Network: Backpropagation Without Storing Activations）作为基础进行分析。

为什么要用可逆网络呢？

因为编码和解码使用相同的参数，所以 model 是轻量级的。可逆的降噪网络 InvDN 只有 DANet 网络参数量的 4.2%，但是 InvDN 的降噪性能更好。
由于可逆网络是信息无损的，所以它能保留输入数据的细节信息。
无论网络的深度如何，可逆网络都使用恒定的内存来计算梯度。

其中最主要目的就是为了减少内存的消耗，当前所有的神经网络都采用反向传播的方式来训练，反向传播算法需要存储网络的中间结果来计算梯度，而且其对内存的消耗与网络单元数成正比。这也就意味着，网络越深越广，对内存的消耗越大，这将成为很多应用的瓶颈。

下面是 Pytorch summary 的结果，Forward/backward pass size（MB）: 218.59 就是需要保存的中间变量大小，可以看出这部分占据了很大部分显存（随着网络深度的增加，中间变量占据显存量会一直增加，resnet152（size=224）的中间变量更是占据总共内存的 606.6÷836.79≈0.725 ）。如果不存储中间层结果，那么就可以大幅减少 GPU 的显存占用，有助于训练更深更广的网络。

import torch
from torchvision import models
from torchsummary import summary

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
vgg = models.vgg16().to(device)

summary(vgg, (3, 224, 224))

结果：

----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 64, 224, 224]           1,792
              ReLU-2         [-1, 64, 224, 224]               0
            Conv2d-3         [-1, 64, 224, 224]          36,928
              ReLU-4         [-1, 64, 224, 224]               0
         MaxPool2d-5         [-1, 64, 112, 112]               0
            Conv2d-6        [-1, 128, 112, 112]          73,856
              ReLU-7        [-1, 128, 112, 112]               0
            Conv2d-8        [-1, 128, 112, 112]         147,584
              ReLU-9        [-1, 128, 112, 112]               0
        MaxPool2d-10          [-1, 128, 56, 56]               0
           Conv2d-11          [-1, 256, 56, 56]         295,168
             ReLU-12          [-1, 256, 56, 56]               0
           Conv2d-13          [-1, 256, 56, 56]         590,080
             ReLU-14          [-1, 256, 56, 56]               0
           Conv2d-15          [-1, 256, 56, 56]         590,080
             ReLU-16          [-1, 256, 56, 56]               0
        MaxPool2d-17          [-1, 256, 28, 28]               0
           Conv2d-18          [-1, 512, 28, 28]       1,180,160
             ReLU-19          [-1, 512, 28, 28]               0
           Conv2d-20          [-1, 512, 28, 28]       2,359,808
             ReLU-21          [-1, 512, 28, 28]               0
           Conv2d-22          [-1, 512, 28, 28]       2,359,808
             ReLU-23          [-1, 512, 28, 28]               0
        MaxPool2d-24          [-1, 512, 14, 14]               0
           Conv2d-25          [-1, 512, 14, 14]       2,359,808
             ReLU-26          [-1, 512, 14, 14]               0
           Conv2d-27          [-1, 512, 14, 14]       2,359,808
             ReLU-28          [-1, 512, 14, 14]               0
           Conv2d-29          [-1, 512, 14, 14]       2,359,808
             ReLU-30          [-1, 512, 14, 14]               0
        MaxPool2d-31            [-1, 512, 7, 7]               0
           Linear-32                 [-1, 4096]     102,764,544
             ReLU-33                 [-1, 4096]               0
          Dropout-34                 [-1, 4096]               0
           Linear-35                 [-1, 4096]      16,781,312
             ReLU-36                 [-1, 4096]               0
          Dropout-37                 [-1, 4096]               0
           Linear-38                 [-1, 1000]       4,097,000
================================================================
Total params: 138,357,544
Trainable params: 138,357,544
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 218.59
Params size (MB): 527.79
Estimated Total Size (MB): 746.96
----------------------------------------------------------------

接下来我将先从可逆神经网络讲起，然后是神经网络的反向传播，最后是标准残差网络。对反向传播算法和标准残差网络比较熟悉的小伙伴，可以只看第一节：可逆神经网络。如果各位小伙伴不熟悉反向传播算法和标准残差网络，建议先看第二节：反向传播（BP）算法和第三节：残差网络（Residual Network）。本文1.2和1.3.4摘录自 @阿亮。

可逆神经网络

可逆网络具有的性质：

网络的输入、输出的大小必须一致。
网络的雅可比行列式不为 0。

1.1 什么是雅可比行列式？

雅可比行列式通常称为雅可比式（Jacobian），它是以 n 个 n 元函数的偏导数为元素的行列式。事实上，在函数都连续可微（即偏导数都连续）的前提之下，它就是函数组的微分形式下的系数矩阵（即雅可比矩阵）的行列式。若因变量对自变量连续可微，而自变量对新变量连续可微，则因变量也对新变量连续可微。这可用行列式的乘法法则和偏导数的连锁法则直接验证。也类似于导数的连锁法则。偏导数的连锁法则也有类似的公式；这常用于重积分的计算中。

1.2 雅可比行列式与神经网络的关系

为什么神经网络会与雅可比行列式有关系？这里我借用李宏毅老师的 ppt（12-14页）。想看视频的可以到 b 站上看。

简单的来讲就是，他们的分布之间的关系就变为，又因为有，所以这个网络的雅可比行列式不为 0 才行。

顺便提一下，flow-based Model 优化的损失函数如下：

其实这里跟矩阵运算很像，矩阵可逆的条件也是矩阵的雅可比行列式不为 0，雅可比矩阵可以理解为矩阵的一阶导数。

假设可逆网络的表达式为：

它的雅可比矩阵为：

其行列式为 1。

1.3 可逆残差网络（Reversible Residual Network）

论文标题：

The Reversible Residual Network: Backpropagation Without Storing Activations

论文链接：

https://arxiv.org/abs/1707.04585

多伦多大学的 Aidan N.Gomez 和 Mengye Ren 提出了可逆残差神经网络，当前层的激活结果可由下一层的结果计算得出，也就是如果我们知道网络层最后的结果，就可以反推前面每一层的中间结果。这样我们只需要存储网络的参数和最后一层的结果即可，激活结果的存储与网络的深度无关了，将大幅减少显存占用。令人惊讶的是，实验结果显示，可逆残差网络的表现并没有显著下降，与之前的标准残差网络实验结果基本旗鼓相当。

1.3.1 可逆块结构

可逆神经网络将每一层分割成两部分，分别为和，每一个可逆块的输入是，输出是。其结构如下：

正向计算图示：

公式表示：

逆向计算图示：

公式表示：

其中 F 和 G 都是相似的残差函数，参考上图残差网络。可逆块的跨距只能为 1，也就是说可逆块必须一个接一个连接，中间不能采用其它网络形式衔接，否则的话就会丢失信息，并且无法可逆计算了，这点与残差块不一样。如果一定要采取跟残差块相似的结构，也就是中间一部分采用普通网络形式衔接，那中间这部分的激活结果就必须显式的存起来。

1.3.2 不用存储激活结果的反向传播

为了更好地计算反向传播的步骤，我们修改一下上述正向计算和逆向计算的公式：

尽管和的值是相同的，但是两个变量在图中却代表不同的节点，所以在反向传播中它们的总体导数是不一样的。的导数包含通过产生的间接影响，而的导数却不受的任何影响。

在反向传播计算流程中，先给出最后一层的激活值和误差传播的总体导数，然后要计算出其输入值和对应的导数，以及残差函数 F 和 G 中权重参数的总体导数，求解步骤如下：

1.3.3 计算开销

一个 N 个连接的神经网络，正向计算的理论加乘开销为 N，反向传播求导的理论加乘开销为 2N（反向求导包含复合函数求导连乘），而可逆网络多一步需要反向计算输入值的操作，所以理论计算开销为 4N，比普通网络开销约多出 33% 左右。但是在实际操作中，正向和反向的计算开销在 GPU 上差不多，可以都理解为 N。那么这样的话，普通网络的整体计算开销为 2N，可逆网络的整体开销为 3N，也就是多出了约 50%。

1.3.4 雅可比行列式的计算

其编码公式如下：

其解码公式如下：

为了计算雅可比矩阵，我们更直观的写成下面的编码公式：

它的雅可比矩阵为：

其实上面这个雅可比行列式也是1，因为这里，它们的系数是一样的。

有另外一种解释方式就是把这种对偶的形式切成两半：

其行列式为 1.

因为是对偶的形式，所以这里的行列式也为 1.

因为，所以其行列式也为 1。

反向传播（BP）算法

上图中符号的含义：

x1，x2，x3：表示 3 个输入层节点。
：表示从 t-1 层到 t 层的权重参数，j 表示 t 层的第 j 个节点，i 表示 t-1 层的第 i 个节点。
：表示 t 层的第 i 个激活后输出结果。
g(x)：表示激活函数。

正向传播计算过程：

隐藏层（网络的第二层）

输出层（网络的最后一层）

反向传播计算过程：

以单个样本为例，假设输入向量是 [x1,x2,x3]，目标输出值是 [y1,y2]，代价函数用 L 表示。反向传播的总体原理就是根据总体输出误差，反向传播回网络，通过计算每一层节点的梯度，利用梯度下降法原理，更新每一层的网络权重 w 和偏置 b，这也是网络学习的过程。误差反向传播的优点就是可以把繁杂的导数计算以数列递推的形式来表示，简化了计算过程。

以平方误差来计算反向传播的过程，代价函数表示如下：