点击下方卡片，关注 「AI视界引擎」 公众号
( 添加时备注：方向+学校/公司+昵称/姓名 )
视觉 Transformer （ViTs）在各种边缘应用中的普及，包括个性化学习，创造了对设备端微调的需求。然而，在边缘设备有限的内存和计算能力下进行训练仍然是一个重大挑战。特别是，训练所需的内存远高于推理所需的内存，这主要是因为需要存储所有层的激活值以计算进行权重更新的梯度。先前的研究通过冻结权重训练以及以压缩格式存储激活值来减少这一内存需求。然而，这些方法被认为效率低下，因为它们无法提供训练或推理速度的提升。
在本文中，作者首先调查了旨在减少内存和计算需求现有设备端训练方法的局限性。
然后作者提出了块选择性重编程（BSR），在其中作者仅对预训练模型的全部块中的一小部分进行微调，并根据冻结层的自注意力分数选择性地丢弃标记。
为了展示BSR的有效性，作者对ViT-B和DeiT-S模型在五个不同的数据集上进行了广泛的评估。与现有方法相比，作者的方法同时将训练内存减少高达，计算成本减少高达，同时保持相似的准确度。
作者还展示了Mixture-of-Expert（MoE）模型的结果，证明了在多任务学习场景中作者方法的有效性。
1 Introduction
在过去几年里，边缘设备上深度学习应用的部署呈现出前所未有的增长。多任务学习（MTL）因其能够以最小的开销动态适应不同任务而在边缘应用中越来越受到重视。目前，边缘设备主要处理推理任务，而训练或微调过程则在云端进行。这不仅涉及到模型和数据传输的通信开销巨大，也引发了数据隐私问题[27]。
尽管设备端训练更为理想，但边缘设备的内存限制带来了挑战。简单的解决方案包括直接使用预训练模型或仅微调最后一层。然而，如果新数据的分布与预训练数据差异显著，这些方法可能导致准确度大幅下降。
之前关于设备端高效学习的工作[3, 22]指出，训练内存主要被激活值占用，而不是参数。为了减少激活内存，已经在卷积神经网络（CNNs）[3, 38]中广泛探索了在冻结原始 Backbone 网络的同时训练残差模块的方法。
然而，由于自注意力、softmax和GELU等非线性层需要存储输入激活以进行梯度计算，即使权重被冻结，这种方法在视觉 Transformer （ViTs）中并不有效。另一项研究通过在激活张量中引入不规律稀疏性来降低ViTs的内存成本[22]。然而，由于稀疏表示的开销，这些方法可能不会产生显著的内存节省，并且没有提供任何潜在的加速。
相比之下，这项工作提出了一种内存和参数高效的微调方法，称为块选择重编程（BSR），它只微调预训练权重张量的一小部分，以最小化传播梯度的激活内存需求。
作者的贡献
作者首先调查并确定了现有SoTA方法在ViTs中的应用限制，以及现有设备端ViT微调方法的不足。特别是，作者观察到残差微调并不一定能减少激活内存，并且可能因为计算开销增加而受到影响。
基于作者的观察，作者提出了BSR，作者选择性地微调预训练网络中的一小部分块。此外，为了降低计算成本，作者采用了一种 Token 丢弃方法，主要根据冻结层的自注意力分数来丢弃 Token 。这节省了激活内存并产生了训练加速，与激活稀疏化[22]不同。
作者详细评估了块选择和 Token 丢弃位置对准确度的影响。作者证明了BSR在两种不同的 Transformer 模型和五个不同数据集上的有效性。具体来说，与现有替代方案相比，作者的方法最多可节省的训练内存，同时计算成本最多可节省。与在CIFAR-10上完全微调一个ImageNet预训练的ViT-B[13]模型（图1）相比，作者的方法在批量大小为32的情况下，提供了的训练内存减少以及的训练时间减少。
最后，作者考虑将多任务学习（MTL）作为作者微调方法的目标，因为在边缘应用中保持单一模型跨不同任务的好处。特别是，混合专家（MoE）模型[8, 19]由于能够解耦参数空间，从而提高MTL性能，在MTL中越来越受欢迎。作者使用BSR在接近 Baseline 准确度的情况下微调MoE，同时将训练内存和FLOPs分别减少2.3和1.5，证明了作者方法在MTL场景中的有效性。
2 Related Work
设备端训练 TinyTL [3] 是最早表明存储激活而不是参数是设备端训练的主要瓶颈的论文之一。因此，像权重剪枝 [15, 17, 20] 和量化 [11, 18] 这样的模型压缩技术并不能显著降低设备端训练的成本。参数高效训练方法，如 Bitfit [41] 或只训练批量归一化参数 [21]，减少了参数内存，这只是总训练内存的一小部分 [3, 22]。
在反向传播过程中重新计算丢弃的激活 [6, 16] 可以降低内存成本，但以计算开销大为代价，这使得它们不适合边缘设备。TinyTL 提出只训练冻结权重的偏置。为了补偿学习能力的降低，它们引入了宽度更小的残差块，并采用组卷积来减少内存占用。Rep-Net [38] 通过设计六个小型残差模块改进了这种残差学习方法，这些模块与冻结的预训练网络间歇性地交换特征，实现了更低的内存占用和更高的准确度。然而，引入残差块会导致参数数量的增加，以及训练和推理计算的增加。其他减少激活内存的方法包括激活量化 [14, 26, 33] 和剪枝 。Back-Razor [22] 提出在前向传播后剪枝用于反向传播的激活。
它们将激活稀疏化到95%，为 ViTs 实现了5.5倍的训练内存减少。它们的训练内存减少受到保存稀疏激活的二值 Mask 所需的内存限制，因此无法实现与激活稀疏性成比例的内存减少。值得注意的是，当前所有的设备端学习方法在推理时都没有提供加速或计算效率。
**高效的视觉 Transformer ** 设计高效的 ViTs 是一个活跃的研究领域，可以分为两类：一类是针对高效 ViTs 的架构设计，另一类是针对 ViTs 的操作优化。在第一类中，研究行人探索了不同的自注意力替代方案，包括 Linformer [37]、SAL-ViT [42] 和 Performer [10]，这些方案旨在减少其二次计算复杂度。其他工作则提出了资源和延迟友好的高效架构，包括 Mobile-former [9] 和 Efficientformer [29]。然而，这些工作并未针对设备端微调。
对于第二类，研究行人开发了各种压缩 [28, 39, 40]、 Token 丢弃 [31] 和 Token 合并 [1] 方案以产生计算效率，这些方案旨在提供推理加速。然而，尽管这些方法减少了计算量和延迟，但它们本身无法实现设备端微调所需的激活和梯度存储的理想减少。
在本文中，作者利用 Token 丢弃在受限的激活内存和计算预算下进行设备端微调。
3 Motivational Analysis
Preliminaries
ViT模型将输入图像划分为个块，这些块也称为token，并将每个token嵌入到一个长度为的嵌入向量中。在图像token集合中添加了一个额外的分类token ，从而生成了个token。负责聚合全局图像信息。一个可训练的位置嵌入被添加到每个token嵌入中，然后这些嵌入通过一系列 Transformer 编码器块，每个块由一个多头自注意力（MHSA）层和一个前馈网络（FFN）组成。token被映射到 Query （）、键（）和值（）矩阵中，每个矩阵的维度为，其中，是MHSA中的头数。每个头执行自注意力操作
FFN通常是一个具有GELU激活和隐藏维度的两层网络。
Analyzing Memory Cost
作者对ViT块在梯度计算中所需的激活内存进行了详细分析。线性层和非线性层的反向传播梯度分别由方程式2和方程式3给出，如下所示。
在这里，第层的输入和输出激活分别用和表示。对于线性层，，并且梯度与无关；而对于非线性层，梯度是的函数，由给出。权重的梯度由方程式4和方程式5给出，其中表示非线性层关于权重的梯度函数，如下所示。
无论层是线性的还是非线性的，权重的梯度都依赖于输入激活。因此，可训练的块必须存储线性层（包括LayerNorm或全连接层）和非线性层（包括Softmax、GELU和Attention）的激活。相比之下，位于梯度流路径中的冻结块只需要存储非线性层的激活。DeiT-S [36]中非线性层存储激活的记忆成本在表1中给出。
对每个块所需的梯度激活内存的详细分析显示，冻结块存储 MB，而完全可训练的块存储 MB，表明使用冻结权重进行训练可以将内存需求减少。从表1可以看出，激活内存主要取决于输入 Token 的数量。Softmax输入呈二次依赖，而所有保存的激活呈线性依赖。
基于这一动机，作者采用了一种有效的 Token 丢弃机制，作者选择性地丢弃非信息性 Token ，在不牺牲准确性的情况下实现激活内存的线性减少。
4 Proposed Approach
Block Reprogramming in ViTs
为了提升在微调ViTs方面的SoTA表现，作者提出了选择性重编程块（BSR），它有选择性地训练预训练模型中的一小部分块，并结合了 Token 丢弃。
Token 丢弃方法在ViTs中已被广泛研究，用于减少延迟和能耗[1, 31]。然而， Token 丢弃尚未在减少设备端训练激活内存的背景下使用。在这项工作中，作者将 Token 丢弃与冻结权重训练相结合，实现了高达的激活内存减少。
作者的 Token 丢弃方法受到EViT[30]的启发。作者使用自注意力分数计算 Token 重要性，并根据 Token 丢弃率融合低重要性 Token 。因此，更少的 Token 传递到可训练块，显著减少了激活内存。
Token 重要性由一个块的MHSA模块通过分类 Token 对所有其他 Token 的注意力计算得出，得分如下所示：
注意力图的第一行是分类 Token （）得到的 Query 向量与键矩阵的点积。作者使用这个点积来计算 Token 重要性，因为值矩阵中的 Token 根据这些分数（）线性组合来预测输出类别。根据重要性得分，保留排名前-K的 Token ，由 Token 丢弃率决定，而不重要的 Token 被融合并传递到后续的FFN层。Token 通过网络长度逐渐丢弃，这些块被称为 Token 丢弃位置。由于预训练模型和完全微调模型4在分类 Token 上显示出完全相同的分布，作者利用预训练的分类 Token 进行重要性计算。
理解梯度流
在本节中，作者导出了残差学习和BSR的梯度流，并解释了作者选择微调主网络的一部分而不是插入残差块的原因。作者进一步反驳了Rep-Net[38]的观点，即因为可训练块和冻结块之间的累加关系而获得激活内存的减少。
在图3中，左边，作者考虑了一个如Rep-net的残差学习网络，其中冻结的第主块，表示为，与可训练的侧块交换特征。第块的输入和输出激活表示为和。主块和侧块的个别输出表示为和，其中和。块的导数表示为。损失关于第侧块的可训练权重的导数如下。
作者将损失函数 对可训练权重 的梯度表示为：
其中，推导的最后一步利用了 与 独立的事实。
方程式7清楚地表明，可训练侧块的权重梯度取决于后续主块的梯度 。因为损失函数的导数依赖于非线性行为的输入激活，所以作者仍然需要沿着主路径存储这些输入激活。
在CNNs中冻结权重的优势源于这样一个事实：
线性层（如卷积）的损失导数只需要权重，而不需要其输入激活，如方程式2所示。ReLU虽然非线性，但只需要存储一个二进制 Mask 。因此，与Rep-Net的声称不同，冻结块与可训练块之间的累加关系在激活内存减少中并不起作用。
内存减少是通过减少可训练块的输入激活维度来实现的，这些可训练块将激活的下采样版本传递给残差块。然而，对于ViTs来说，这个优势就消失了，因为每个编码器块都包含自注意力、softmax和GELU，这些都需要存储激活输入以便于梯度传播（方程式3）。因此，引入残差块会进一步增加激活内存，如作者在第5.1节详细描述的实验所示。
此外，在残差学习中，正向传播会通过冻结的主块和可训练的残差块，导致在训练和推理期间参数和计算成本的增加。
在图3的右侧，作者展示了BSR中的梯度流，其中第主块被设置为可训练。损失函数 对主块 的可训练权重 的导数表示如下：[此处应继续原文中的公式，但提供的文本中未包含后续内容]。
缺乏残差路径简化了等式8，这意味着只需要存储主路径上的激活。
Design Choice Discussion
讨论1 在针对小规模数据集进行微调时，哪些模块更为关键？
作者分析了在CIFAR-10数据集上，从ImageNet预训练模型到完全微调模型之间中间特征的变化，以观察哪些层次显示出最大的差异。
图4展示了编码器块之前的三种可训练元素（ Patch 嵌入、位置嵌入和分类标记）的层次统计数据。
令人惊讶的是，作者发现这些元素在预训练模型和微调模型中的分布完全相同。这一点尤其对于负责聚合最终分类信息分类标记来说并不明显。通常，微调模型的最后几层更为重要[5, 12, 35]，这就是在迁移学习应用中通常只微调最后几层的原因。对于卷积神经网络（CNN）来说，前几层或特征提取层在各项任务中通常是相似的。
这引出了作者的下一个讨论话题：仅微调最后几层在准确性和激活内存方面是否足够？
讨论2**仅微调最后几块是否足够？*
表2展示了微调最后两块、三块和四块，以及是否进行标记剪枝的结果。作者观察到，仅训练靠近末端的块无法弥补准确性的差距，即使在没有标记剪枝的情况下也是如此。当标记剪枝时，因为末端只剩下少数标记，作者观察到准确度显著下降。这需要智能地选择可训练块和标记丢弃位置，以在准确性和激活内存之间实现最佳权衡。
讨论3标记丢弃位置与可训练块的相对位置 网络的整体标记丢弃率取决于标记丢弃位置和在这些位置丢弃的标记比例。例如，如果选择较高的标记丢弃率，将标记丢弃位置放置在网络末端与在较低丢弃率的较浅层放置标记丢弃位置具有相同的效果（表8）。EViT[30]显示，当在第三层之前丢弃标记时，标记丢弃方法会导致性能显著下降，因为 Transformer 模型无法那么早识别出重要的标记。他们进一步观察到，在这一层之后，无论丢弃位置和丢弃率如何，网络在相同水平的整体标记减少下表现出稳定的性能。
作者在第5.3节验证了这一观察。作者将标记丢弃位置放置在第4块、第7块和第10块，标记丢弃率为0.5，从而导致整体标记减少50%。在这种标记丢弃配置下，可训练块的位置被调整以找到在准确性和训练内存之间的合适权衡。
表7呈现了可训练块位置消融研究。作者做了几个有趣的观察。
- 在标记丢弃位置之前放置可训练块可以提高准确性，但会极大地增加训练内存消耗。
- 在标记丢弃位置之后放置所有可训练块可以带来显著的内存优势，但会导致准确度下降。
- 在均匀分布的深度上训练块的表现优于仅训练最后的块。
基于这些观察，作者将12个块中的3个（即第4块、第8块和第12块）保留为可训练。第4块也是一个标记丢弃位置，所有N个标记都通过MHSA进行标记重要性计算，然后将N/2个标记传递给FFN层。第8块和第12块分别只处理N/4和N/8个标记。
5 Experimental Results
模型与数据集：作者在DeiT-S [36]和ViT-B [13]上展示了结果。对DeiT-S在ImageNet1k上进行预训练的模型进行微调，对ViT-B在ImageNet22k上进行预训练的模型进行微调。MoE模型由一个在ImageNet上预训练的ViT-B主干网络构建，包含三个MoE层，每层由16个专家组成。在五个数据集上的迁移学习性能已展示：CIFAR-10，CIFAR-100 [23]，Flowers [32]，Pets [34]，和Food [2]。估计的内存结果是按照[3, 22]计算的，设备上的内存是在NVIDIA RTX A6000 GPU上测量的。
训练超参数：按照[22]中的训练设置，在Pytorch中训练了ViT-B模型。使用SGD优化器训练ViT-B模型，共20k步，采用余弦学习率衰减，每个数据集调整初始学习率。ViT MoE模型使用SGD训练，共8k步，初始学习率为0.01。DeiT-S模型使用AdamW优化器训练50个周期，采用余弦学习率衰减。默认图像大小设置为224，默认块大小设置为16。
Results and Analysis
作者的方法采用了根据第4.3节讨论3中描述的设置配置的ViT-B和DeiT-S模型，整体标记丢弃率为50%。表3展示了DeiT-S在CIFAR-10和CIFAR-100上的结果。作者的方法实现了6.03的训练内存减少，同时在CIFAR-10上的准确度下降了，在CIFAR-100上下降了。作者还实现了约的FLOPS减少。
在表4中，作者将全程微调（FT-Full）、仅微调最后一层或分类头（FT-Last）、Back Razor [22]以及作者的方法进行了比较。对CIFAR-100的计算训练内存、设备内存和FLOPS进行了计算。作者与80%和95%激活稀疏度的Back Razor进行比较。FT-Last提供了最高的内存节省，同时在像Flowers这样的数据集上表现合理，但在CIFAR-100和Food上则出现了实质性退化。作者的一些FT-Full准确度略低于Back Razor报告中 Baseline 准确度。例如，对于Flowers-102，作者的FT-Full准确度为99.2%，而他们报告的准确度为99.5%。Back Razor和作者的方法在Flowers-102上都出现了0.1%的退化。对于Pets，作者方法的准确度与95%的Back Razor相当，尽管作者的FT-Full准确度低了0.4%，这意味着作者的方法退化程度少了0.4%。作者的方法在CIFAR-100和Pets上的表现优于95%的Back Razor，分别高出1.02%和0.4%，在Flowers数据集上表现相当。作者在提供比95%的Back Razor更大的训练内存减少的同时，还实现了的FLOPS减少。
Extending BSR for Mixture-of-Experts
多任务学习（MTL）[7, 19]因其能够动态地适应任务，且开销最小，正在成为边缘智能的关键。MTL模型为不同的任务学习共享的表示，从而避免了训练以及存储单独模型的负担。"专家混合"（MoE）[8, 19]在MTL中构成了一种新的范式，它通过根据任务和输入标记激活模型的部分来分离参数空间，从而提供了改进的MTL性能。在本节中，作者证明作者的方法可以很容易地扩展用于微调MoE模型。据作者所知，先前未曾探索过针对MoE模型的设备端训练。
对于ViT MoE，ViT中的FFN被替换为一个MoE层。一个MoE层由几个表示为MLP的专家组成。一个任务依赖的路由器为每个输入标记稀疏地激活一组专家。作者使用了一个带有ViT-B Backbone 网络的MoE模型。作者只将三个块中的MLP层（对应于BSR中可训练块的位置）替换为MoE层。MoE层有16个专家候选者，路由器为每个标记选择前四个候选者，将结果相加生成MoE层的输出。表6展示了BSR对MoE微调的结果。单任务学习（STL）由每个任务的独立模型组成。作者观察到BSR对MoE微调表现出色，相对于CIFAR-10、CIFAR-100、Flowers和Food的MoE Baseline ，准确度下降仅为0.21%、0.78%、0.09%和0.78%，而在Pets上则比 Baseline 高出1.04%。作者在设备上获得了2.32倍的内存减少和1.5倍的FLOPs减少。值得注意的是，在大多数情况下，作者的MoE模型显著优于STL模型。
Ablation Studies
可训练块的数目与位置 表7展示了关于可训练块的数目和位置的消融研究。索引从0到11，对应于十二层的模型深度。使用四个训练块的训练并没有比三个训练块提供显著优势。实际上，三个训练块通常表现得比四个训练块更好，如训练块[4, 7, 11]与[4, 7, 10, 11]或[2, 5, 8, 11]与[2, 5, 8]的比较所见。使用两个训练块会导致退化程度更高。因此，作者选择可训练块的数量为三个。可训练块相对于 Token 丢弃位置的位置提供了准确性与内存的权衡，如第4.3节所讨论。
表5：在CIFAR-10上，CNN和ViT的残差学习和BSR。激活内存和FLOPS针对批量大小为1的情况报告。
表8展示了关于 Token 丢弃率和位置的消融研究，使得总体丢弃率约为50%。作者观察到，在相似的丢弃情况下，准确度变化很小（约为0.1%），同时内存和FLOPs的减少也相似。
6 Conclusions
作者提出了BSR方法，这是一种用于在设备上高效训练ViTs的方法，它将针对下游任务选择性微调预训练模型中的一小部分块与基于主要冻结权重的自注意力分数的标记丢弃相结合。与现有减少训练内存但增加计算开销的设备学习方法不同，BSR在将计算成本降低2倍的同时，将训练时间缩短了约2.5倍，内存减少了高达6.03倍。
作者在设备上训练的内存减少方面比现有最佳(SoTA)方法提高了1.4倍，同时保持了相似的性能。作者进一步证明了该方法在MTL应用中对MoE模型的有效性。作者的方法与激活量化正交，后者可以用于进一步减少激活内存。
参考
[1].Block Selective Reprogramming for On-device Training of Vision Transformers.
点击上方卡片，关注 「AI视界引擎」 公众号