Facebook团队新作 | Rev-ViT: 一种即插即用且对内存友好的Transformer架构

以下文章来源于CVHub ，作者派派星

Paper: https://arxiv.org/pdf/2302.04869.pdf

Code:   https: //github.com/facebookresearch/slowfast.

导读

SOTA模型在计算量和内存占用指标上的对比

近年来，计算机视觉领域中出现的许多基于深度学习的革命性工作均离不开高性能硬件的支持。在专用 AI 加速器的推动下，对最先进模型的计算量要求呈指数级增长。然而，计算量只是一部分，还有另一部分是被大家经常忽视的内存带宽瓶颈。 据统计，AI 加速器的FLOPs是以每2年约3.1倍的速度增加，而带宽仅是每2年约1.4倍。这种差异导致了所谓的内存墙，这尤其在transformer中尤其明显。

因此，对于带宽需求高的模型可以通过重新计算（re-computation），以计算换取内存可能比目前SOTA算法更有效。 例如，在模型训练阶段，可以通过重新计算激活值代替在内存中对激活值进行先存后取的操作。除了训练速度之外，缩放vision transformer网络深度会影响 GPU 内存容量。例如，视频识别领域的SOTA模型的batch size仅能为1，这是由于中间激活层的内存占用量过大造成的。

基于此，本文提出了可逆的vision transformer（Rev-ViT），与不可逆的版本相比，它的内存占用量显著降低，如上图所示。通过权衡 GPU 激活缓存与高效的即时激活重新计算，Rev-ViT有效地将激活内存增长与模型深度分离。 此外，通过重新配置ViT和Multiscale ViT中的残差路径，有效解决了模型训练收敛不稳定问题，使得Rev-ViT能够直接适应深层的常规ViT架构。

本文提出的Rev-ViT方法的主要贡献如下：

提出了Rev-ViT 和Rev-MViT，这是目前最先进的内存友好型的视觉识别backbone。

通过大量的消融实验发现，Rev-ViT比常规ViT具有更强的正则化能力。

在图像分类、检测和动作识别数据集上验证了本文方法的高效性，Rev-ViT比ViT在内存使用的指标上平均低10+倍。

欢迎大家关注我们的公众号CVHub，每日都给大家带来原创、多领域、有深度的前沿AI论文解读与工业成熟解决方案！

方法

Rev-ViT和Rev-MViT结构图

本节基于以下内容进行展开讲述：

Rev-ViT的两个残差流结构和相关约束

Rev-MViT架构及端到端的可逆训练原理

可逆的 Vision Transformers

适用于 ViT 的两个残差流

上图(a) 显示了适用于 Vision Transformer 架构的可逆转换。输入由两个分区张量  和  组成。这导致了一个双残差流架构，其中每个输入  和  都保持自己的残差流，同时使用函数  和  将信息相互混合。在 ViT 之后，我们使用多头注意力 MLP 子块分别作为函数  和 。

边界条件

由于 ViT 架构仅使用单个残差流，因此需要修改架构以支持双残差流设计。我们提出以下建议：

我们保持 stem 完整，并将 patchi-fication 输出激活发送到  和 。

这两个残差路径需要在最终分类器头之前融合以保留信息。我们建议首先对输入进行层归一化，然后进行级联，以减少融合计算开销。

重新配置残差连接

残差连接对深度网络的信号传播起着非常关键的作用，可逆变换  可行的的关键取决于两个流之间的残差连接。有趣的是，我们通过实验观察到 Rev-ViT 中残差连接和信号传播之间的关键关系。

虽然通常的做法是将神经网络块包裹在残差连接周围以获得更好的梯度流，但  或  输入都没有这种连接。具体来说，MLP 周围的内部残差连接和  和  流的注意力子块都不存在。相反，每个残差流的残差连接流经另一个流，通过可逆转换  中存在的固有跳过连接进行操作。我们发现这些内部跳跃连接不利于深层模型的训练收敛，同时不会为较浅的模型带来额外的增益，因此选择完全忽略它们用于可逆视觉转换器块。

可逆的多尺度Vision Transformers

最近提出的 MViT 架构通过对视觉分辨率进行下采样和对通道维度进行上采样，在模型内部建立了一个特征层次结构。它在图像和视频分类基准上都获得了SOTA结果。为了展示可逆设计的灵活性，我们设计了可逆多尺度 Vision Transformers。

阶段转化块（Stage-Transition Block）

上图(b) 展示了阶段转换块的架构。阶段转换模块紧跟 MViT 中分辨率上采样模块的设计，并进行了以下重要修改：

横向连接：残差流  和  在阶段转换块的开始与横向连接融合。这是一种有效计算分辨率下采样和特征上采样的方式，避免了在每个流中分别重复计算。

特征上采样：MViT 在分辨率上采样块之前的最后一个 MLP 块中执行特征上采样。我们将通道上采样阶段移动到阶段转换块的池化注意力子块内。具体来说，在池化通道卷积层之后的线性层中对Q、K和V向量进行上采样（上图 2b 和 2c）。

阶段保留块（Stage-Preserving Block）

上图(c) 显示了适用于多尺度Vison Transformer 架构的可逆转换 。该设计非常类似于Rev-ViT块（上图(a)）的设计，在此基础上增加了多头池化注意力。值得注意的是，即使注意力在K和V矩阵上使用池化，从而改变了序列长度，但其输出维度仍然保留不变。因此，阶段保留块仍然遵循等维约束，在没有缓存激活的情况下实现完全可逆和学习。

由于每个阶段转换块都会改变时空分辨率，因此它们在整个 MViT 网络中只出现有限的次数。换句话说，大部分计算和内存使用都是在阶段保留块内执行的，并且是完全可逆的。对于阶段转换和阶段保持块，我们遵循与可逆Rev-ViT块中相同的残差连接数。

实验

在ImageNet-1K分类数据集上与SOTA模型的性能对比

在同等精度下，Rev-ViT及其变体节省了2~15倍的推理所需内存，增加了1~13倍的最大输入数量。

在Kinetics-400视频分类数据集上与SOTA模型的性能对比

在同等精度下，Rev-MViT比MviT节省了1倍的推理所需内存，增加了一倍的最大输入数量。

在COCO检测数据集上与SOTA模型的性能对比

在同等精度下，Rev-MViT比MviT节省了将近1倍的推理所需内存。

消融实验

a(1)：没有激活缓存不会损害不同深度的 Rev ViT-B 的精度;

a(2): 内部残差连接对于浅层模型训练良好，但对于更深的模型训练发散;

(b)：Rev-MViT 对于更高的输入分辨率和更深的 MViT 模型具有更高的吞吐量；

(c)：不同规模大小Rev-ViT 的吞吐量均比对应规模大小ViT的高，充分证明了Rev-ViT有效性。

不同规模的Rev-ViT模型

总结

本文设计了两种内存友好的可逆ViT，分别为Rev-ViT和Rev-MViT。此外，本文提出的方法在多个指标上（模型复杂性、吞吐量、准确性和内存使用情况）对多个任务（图像分类、对象检测和视频分类）进行了基准测试并获得了SOTA性能。简单来说，Rev-ViT 和 Rev-MViT 模型分别比 ViT 和 MViT 模型实现了 15.5 倍和 4.5 倍的内存占用。

转自：“arXiv每日学术速递”微信公众号

如有侵权，请联系本站删除！