小内存有救了！Reversible ViT：显存减少15倍，大模型普及曙光初现！

摘要

本文介绍了可逆视觉Transformer，这是一种用于视觉识别的高效内存架构设计。通过将GPU内存需求与模型的深度解耦，可逆ViT能够以高效的内存使用来扩展架构。论文将两种流行的模型，即ViT和多尺度ViT，调整为可逆的变体，并在图像分类、目标检测和视频分类的模型大小和任务中进行广泛的基准测试。可ViT在大致相同的模型复杂性、参数和精度下实现了高达15.5倍的内存占用减少，证明了可逆ViT作为硬件资源有限训练机制的有效骨干的前景。最后论文发现，对于更深的模型来说，重新计算激活的额外计算负担已经被克服了，其中吞吐量可以比不可逆的模型增加2.3倍。

代码链接：https://github.com/facebookresearch/slowfast

简介

计算机视觉的深度学习革命建立在高性能硬件加速器的基础上。在专用人工智能加速器的推动下，对最先进模型的计算需求呈指数级增长。然而计算只是故事的一半。另一个，也是经常被忽视的一半，是内存带宽瓶颈，与峰值加速器FLOP[54]相比，它很难按比例扩展。特别是峰值加速器FLOP以每2年约3.1倍的速度增加。然而，峰值带宽仅以每2年约1.4倍的速率扩展。Transformer中的这种差异加剧了，在过去三年中，Transformer的所需计算量大约每三个月翻一番，导致了所谓的memory wall，在memory wall中，整体模型性能和训练速度都变得严格受限[34]。

因此，对于带宽受限模型，通过重新计算来换取内存实际上可能比使用工作优化算法更有效[70，71]。在训练神经网络模型的情况下，这可以通过重新计算激活来实现，而不是存储并从DRAM[31]加载它们。除了训练速度之外，scaling vision transformers自然会影响GPU的内存容量，尤其是在内存不足的情况下，如视频识别，由于中间激活的高内存占用率，最先进的模型通常被限制为批量大小1。

本文提出了可逆ViT，这是一个表达性视觉识别架构家族，与它们的不可逆变体相比，具有非常有利的激活记忆足迹（图1）。通过将GPU激活缓存与高效的动态激活重新计算进行权衡，可逆ViT有效地将激活内存的增长与模型的深度解耦。

总结来说，本文的主要贡献如下：

• 论文提出了可逆视觉Transformer（Rev ViT）和可逆多尺度视觉Transformer（Rev.MViT），这是对最先进的视觉识别主干的记忆高效可逆自适应；

• 论文观察到可逆Transformer比简单的网络具有更强的固有正则化。因此通过调整具有不同重复增强、增强幅度和下降路径率的原始配方来开发新的训练配方，以匹配其不可逆对应物的性能；

• 论文在以下几个任务中对模型进行基准测试：图像分类、目标检测和动作识别、准确性、内存、最大训练批量和模型复杂性。特别是，在匹配的复杂度（FLOP/参数）和最终精度下，Rev-ViT-B和Rev-ViT-L的每图像内存占用分别比ViT-B和ViT-L轻8.2倍和15.5倍。此外，论文还展示了深度可逆网络如何实现比普通网络高达2-4倍的吞吐量。

相关工作

Transformers：Transformers是一种流行的网络结构，最初被提出用于自然语言应用[68]，现在广泛应用于深度学习的所有领域，如强化学习、语音、音乐、多模态学习，最近还广泛应用于传统的视觉任务。自引入以来，ViT得到了广泛的采用，并已应用于多个视觉识别任务。

然而放大Transformer模型的一个关键问题是所需GPU内存随深度的增长而增长。这种内存的线性增长阻碍了非常深入的模型的开发，因为需要大幅减少批量大小才能适应在GPU上存储中间激活。这个问题在处理非常大的输入张量的视频模型中加剧，并且即使对于较浅的深度，也经常用批量大小1来训练。扩大传统Transformer架构的一个潜在系统级解决方案是模型并行[10]，它将模型的不同部分放在不同的GPU上。然而，在实践中，它相当慢，并且由于巨大的跨设备流量，需要特殊的高带宽网络基础设施。

可逆体系结构：可逆架构是一系列基于NICE可逆变换模型的神经网络架构，该模型是现代基于生成流的图像生成架构的前身。基于NICE可逆变换，Gomez等人提出了一种可逆ResNet架构，该架构采用可逆变换[12]在ResNets[27]中进行高效记忆的图像分类。可逆CNNs已应用于几种传统的图像任务，如压缩、重建[43]、检索[42]和去噪，以及压缩传感、紧凑分辨率、图像到图像翻译、遥感、医学图像分割和MRI重建。

然而，单词级输入划分比补丁级图像划分包含更丰富的语义内容，NLP Transformer往往在深度上更浅，但在通道维度上更宽。例如，Kiatev等人[38]专注于扩展输入序列维度，而不是模型深度，并且没有对最大批量大小、峰值GPU内存和训练吞吐量进行基准测试。

论文的的实验表明，可逆视觉Transformer的自适应对于更深（8个block）的模型表现不佳。这项工作首次提出了可逆视觉Transformer，将其适应两个最先进的Transformer网络，即ViT和MViT。此外，这项工作是首次将可逆主干用于目标检测和视频分类，这往往是视觉识别中最缺乏内存的领域之一。

方法

可逆block结构

可逆Transformer由一堆可逆block组成，这些可逆块遵循可逆变换的结构，以允许输出的分析可逆性。

可逆变换

考虑一个变换，它将被划分为二维张量的输入张量变换为具有任意可微函数的同样被划分为张量的输出张量O：

同样，考虑一个类似的转置变换：

现在考虑组合，其变换输入向量的分区并且被获得为：

朴素网络需要缓存激活

考虑反向传播机制。给定计算图节点M、其子节点以及子节点相对于最终损失的梯度，反向传播算法使用链式规则来计算的梯度：

应用上述反向传播算法来计算相对于父节点的导数并且使用输出Y作为唯一子节点，可以得到：

因此，由于函数雅可比，反向传播算法需要在前向传递期间的中间激活在后向传递中可用，以计算相对于权重的梯度。通常这是通过将中间激活缓存在GPU内存上以用于反向通道来实现的。这允许以额外内存为代价进行快速梯度计算。此外，网络的顺序性质要求在计算损失梯度和释放高速缓存的存储器之前，所有层的激活都要被高速缓存。这种依赖性显著影响峰值存储器使用，因此峰值存储器使用变得线性依赖于网络深度D。

无需缓存激活即可进行学习

如前文所示用可逆变换T变换的输入允许从变换的输出重新计算输入。因此，由这种可逆变换组成的网络不需要存储中间激活，因为它们可以在从输出向后的过程中容易地重新计算。然而，可逆变换T对学习函数的性质施加了重要的约束。

等尺寸约束：如前文所述，函数F和G在输入和输出空间中需要是等维的。因此，在T下，特征尺寸需要保持不变。虽然这种约束对其他视觉架构（如ResNets[27]）来说是一个障碍，这些架构需要改变特征尺寸，但在视觉Transformer架构[15]中很容易满足，因为它在整个层中保持恒定的特征尺寸。

可逆Vision Transformers

Adapting ViT to Two-Residual-Streams

图2a显示了适用于ViT架构的可逆变换T[15]。输入由两个分割张量和组成，这两个张量根据方程3.1.1进行变换，以保持可逆性。这导致了两个残差流架构，其中输入和中的每一个保持它们自己的残差流，同时使用函数F和G将信息彼此混合。在ViT[15]之后，分别使用多头注意力和MLP子块作为函数F和函数G。

边界条件

由于ViT体系结构仅使用单个残差流，因此需要对该体系结构进行修改，以支持双残差流设计。论文提出以下建议：

• 开始。我们保持stem的完整性，并将patchification输出激活发送到和。请注意，这种设计选择与[23]不同，后者建议沿通道尺寸一分为二。

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

重新配置残差连接

残差连接在深度网络中的信号传播中起着关键作用[27]。可逆变换T本身也关键地取决于两个流之间的剩余连接以保持可逆性。有趣的是，论文在可逆视觉Transformer中观察到残差连接和信号传播之间的关键关系。注意，虽然通常将神经网络块包裹在残差连接周围以获得更好的梯度流[27]，但或输入都没有这种连接。具体地，不存在用于和流的MLP和关注子块周围的内部残差连接。而是每个残差流的残余连接流过另一个流，通过可逆变换T中存在的固有跳跃连接进行操作。论文发现这些残差跳跃连接不利于深层模型的训练收敛，而不会为较浅模型带来额外增益，并选择在可逆视觉变换块中完全省略它们。

Reversible Multiscale Vision Transformers

最近提出的MViT体系结构通过对视觉分辨率进行下采样和对通道维度进行上采样，在模型内部开发了一个特征层次。它在图像和视频分类基准上都获得了最先进的结果。为了展示可逆设计的灵活性，论文将MViT模型应用于可逆多尺度视觉Transformer。我们建议以与MViT模型相同的结构组成可逆MViT架构，但使用两个不同的层——阶段转换和阶段保留块。

Stage-Transition Block

图2b描述了拟议的阶段敏感区块的架构。阶段转换块严格遵循MViT[18]中分辨率上采样块的设计，并进行了以下关键修改：

• 横向连接。残差流或在阶段过渡块的开始处与横向连接融合。这允许分辨率下采样和特征上采样的有效计算，而无需在每个流中单独重复计算；

• 特征上采样。MViT在分辨率上采样块之前的最后一个MLP块中执行特征上采样。我们建议将通道上采样阶段移动到阶段转换块的池化注意力子块内。具体而言建议在池化通道卷积层之后的线性层中对q、k和v进行上采样（图2b和2c）。这是双重好处：（A）允许所有特征维度的变化在同一块内同步发生，并允许其他块保持特征维度不变，这是可逆架构的优点；（B）避免了在先前的MLP和池化层中使用额外的计算。论文在阶段转换块处遵循与可逆视觉Transformer架构中相同的边界条件。

Stage-Preserving Block

图2c显示了适用于多尺度视觉Transformer架构[18]的可逆变换T。该设计与可逆视觉转换块（图2a）非常相似，增加了多头池化注意力[18]。请注意，即使注意力在k和v上使用了池化，从而改变了序列长度，但输出维度仍然保留下来。因此，阶段保留块仍然遵循等维约束，因此可以在不缓存激活的情况下使其完全可逆并学习。

由于每个阶段转换块都会改变分辨率，因此它们在整个MViT网络中只发生有限的次数。换句话说，大部分计算以及内存使用都是在阶段保留块内执行的，并且是完全可逆的。论文遵循与可逆视觉Transformer块中相同的残差连接电路，用于stage-transition和stage-preserving块。

结果

图像分类

数据集：我们在图像分类（ImageNet[11]）、视频分类（Kinetics 400[36]和Kinetics 600[5]）和对象检测（MS-COCOCO[45]）方面对可逆视觉转换器和可逆多尺度视觉转换器架构进行了广泛的基准测试。在所有的基准测试中，我们观察到通过使用可逆体系结构可以显著节省内存，而精度变化可以忽略甚至没有。所有呈现的结果和消融都是从随机初始化中训练的，除了COCO，我们从ImageNet权重中初始化。

结果：表1显示了不同模型和FLOP匹配下可逆ViT和可逆多尺度ViT的结果。论文在224 x 224图像大小和其他相同条件下，在单个16 GB V100 GPU上对所有模型进行基准测试。最大批处理大小是作为一个批处理中可以在不耗尽GPU内存的情况下训练的图像的最高数量而获得的。每个图像的内存被测量为每个图像在训练期间占用的峰值GPU内存。

论文注意到，可逆ViT与其不可逆对应物的FLOP和参数规格完全匹配，这是由于可逆视觉变换块的简约设计。可逆多尺度ViT由于阶段转换阶段而具有略高的FLOPs，同时由于阶段保留阶段而仍然具有非常高的GPU内存效率。

随着深度的增加而增加内存节省。在表1中观察到，Rev ViT在所有模型变体（Small、Base和Large）和FLOP方案中都与朴素ViT的性能非常接近。由于每个图像使用的内存线性地依赖于香草网络的模型深度，因此可逆模型的内存增益随着网络深度的缩放而增加。值得注意的是，尽管可逆ViT-S已经相对于普通ViT-S模型节省了约86.8%的内存（相当于减少了7.6倍），但可逆ViT-L模型的增益进一步增加到15.5倍，或者节省了约93%的内存。

视频分类

表2和表3分别给出了Kinetics-400[36]和Kinetics-600[5]上动作识别任务的结果。对于动作识别，论文对经过调整的可逆MViT模型进行了基准测试，并报告了两个数据集的前1和前5性能。与图像分类基准类似，可以观察到可逆MViT模型以很小的内存成本与不可逆模型的精度非常匹配。

目标检测

论文还在MS-COCO[45]上对所提出的Rev-ViT-B和Rev-MViTB模型的对象检测进行了基准测试。所有模型都在118K的训练图像上进行训练，并在5K的验证图像上进行评估。论文采用在IN上预先训练的ViT-B和MViT-B骨干，并使用标准Mask R-CNN[26]作为检测框架。所有模型都使用标准的3x时间表（36个时期）进行训练。对于MViT，论文将多尺度主干与特征金字塔网络相集成[44]。参考表4观察到，Rev-MViT-B模型与MViT-B上的AP性能非常匹配，仅为内存成本的54.8%。

消融

更强的内在正则化。在不同的模型和数据集中，论文发现在相同的FLOP和参数规范下，可逆模型往往比不可逆模型具有更强的内在正则化。因此，可逆视觉Transformer的训练配方具有更轻的重复增强，更小的增强幅度，因此具有更高的重量衰减。表5显示了这些配方变化对Rev-ViT-B的影响。在其他Rev-ViT和Rev-MViT模型上也观察到了类似的效果，其中具有较轻增强和较高重量衰减的修改训练配方在匹配性能中发挥着至关重要的作用。

横向融合策略的消融实验如表6所示。

重新配置残差连接：论文在图3a中省略了ViT架构的这种设计选择。在不改变任何其他模型尺寸的情况下改变模型深度，并观察两种可逆变体的性能，即有和没有内部跳跃连接。注意到，虽然具有内部跳跃连接的naive版本在浅层模型中稳定训练，但在深层模型中，精度显著下降。另一方面，Rev MViT可以很好地随深度缩放，在较浅的深度下与其他变体一样准确，但在较深的模型下明显更好。

不缓存激活的学习效果：图3a还比较了使用和不使用缓存激活训练的Rev-ViT-B架构的图像分类性能。这使我们能够将所提出的可逆视觉Transformer所需的残差配置的效果与在没有缓存激活的情况下学习可能产生的任何伪影区分开来。然而对于所有深度，注意到在没有缓存激活的情况下训练的Rev-ViT-B性能，以密切跟踪使用缓存训练的相同架构的性能。深度12结果的微小差异可能源于训练配方适用于在没有激活的情况下训练的实际Rev-ViT-B架构。

模型大小和输入分辨率与吞吐量：图3b显示了不同模型大小在224和320输入分辨率下的训练吞吐量比较。我们注意到，虽然对于较小的模型，如深度为12层的MViT-B，Rev-MViT-B的训练吞吐量略小（98.5对86.0），但在更高分辨率的训练和更深的模型中，中间激活的额外重新计算负担都很容易克服。特别地，在224分辨率下，24层和48层Rev-MViT模型具有与MViT模型相似的吞吐量，对于80层深度模型，在384分辨率下的吞吐量增加到2.1倍，在384解析度下的吞吐量提高到2.3倍。此外，与普通网络相比，可逆变体的更深模型的内存增长率要低得多，允许扩展到更深的模型，而不需要任何额外的训练基础设施负担或内存需求，如梯度检查点或模型并行。

最大batch：在图3c中对Rev-ViT-Base（B）、Large（L）和Huge（H）及其不可逆对应物的最大可能批量大小进行了基准测试。我们通过深度缩放ViT-L和ViT-H（保持其他模型尺寸不变）将趋势（表示为）外推到更大的模型，并对其可逆对应物进行最大批量大小基准测试。

模型大小与GPU内存占用：图1绘制了Rev-ViT和Rev-MViT系列模型以及其他几种先前网络（如MViT[18]、ViT[15]、ResNets和Reg-NetY[58]）的GPU内存占用图。我们注意到，在固定的GFLOP下，可逆变体的记忆效率极高，MViT高达4.5倍，ViT高达15.5倍，超过了先前的卷积变体几个数量级。

结论

本文提出了可逆视觉Transformer，这是对ViT和MViT模型进行的高效记忆可逆架构调整。论文对图像分类、目标检测和视频分类等多项任务以及模型复杂性、吞吐量、准确性和内存使用等多项指标进行了基准测试。在给定任何规格的情况下，本文的Rev ViT和Rev MViT以很小的内存成本与不可逆变体的精度相匹配，同时对较小的模型保持类似的训练吞吐量，对较大的模型保持高达2.3的吞吐量。具体而言，可以观察到Rev ViT和Rev MViT型号的内存占用分别比ViT和MViT轻15.5倍和4.5倍。

后续工作

Chen等人[77]将所提出的可逆视觉Transformer架构设计应用于可逆Swin Transformer，并将该模型应用于记忆有效的时间动作定位。时间动作定位涉及检测动作的开始和结束边界的精确时间帧位置，因此需要在密集采样的视频上执行。密集的帧采样导致训练期间的GPU内存开销，从而禁止在时间动作定位任务（TAL）上对骨干进行端到端微调。可逆主干减轻了内存开销，并允许高效的端到端TAL训练，从而显著提高了定位性能。

同时，[80]提出了一种可逆Transformer的训练程序，该程序允许加速训练，同时确保原始计算的精确复制。特别地，[80]提出在使用先前块的重新计算的激活的梯度计算之前将激活重新计算交错一个变换器块。这使得一旦上一个块完成，激活就可用于下一个块的梯度计算。因此，梯度计算步骤不需要等待激活重新计算，有效地将重新计算负担的延迟隐藏在graident计算的延迟之后，从而有效地加快了训练。这需要维护单独的cuda工作流，以异步方式处理上述两个步骤。根据硬件和计算规模的不同，这种运算符并行化可能会显著加速。

参考

[1] Reversible Vision Transformers