NeurIPS 2022 | 基于激活值稀疏化的内存高效迁移学习

科技

2023-02-13 05:02

©PaperWeekly 原创 · 作者 | GlobalTrack

论文标题：

Back Razor: Memory-Efficient Transfer Learning by Self-Sparsified Backpropagation

论文链接：

https://openreview.net/pdf?id=mTXQIpXPDbh

源码链接：

https://github.com/VITA-Group/BackRazor_Neurips22

简介

从大的数据集到下游任务的迁移学习广泛使用。训练得到的模型作为预训练模型可以提升模型泛化能力。

有时需要在边缘端微调云端训练得到的大模型。一个主要的挑战是边缘端的内存约束。ResNet50 模型以 batchsize16 的设定很容易超过机载设备的内存限制。新出现的 VIT 模型会达赖更大的内存消耗。近期的研究指出在后向传播过程中存储的激活值占据了训练过程中内存消耗的很大一部分。

一些已有的解决方法：参数高效训练算法 Bitfit 不能通过压缩激活值满足内存限制。梯度检查点（Gradient Checkpointing）只需要存储激活值的子集，但是由于需要重新计算丢弃的激活值，后向传播过程中需要更多的 FLOPs。

本文给出了一种一般的激活层压缩算法 Back Razor。受到激活值存储只用于后向传播的启发，本文提出了非对称剪枝策略，即旨在后向传播过程中使用稀疏的激活层权重。

本文方法

考虑线性层，为一个例子，需要计算的梯度为:

如上述公式所示，激活值直接用于梯度的计算过程。这产生了在前向过程中保留激活值的需要。否则需要额外的 FLOPs 重新计算激活值。

2.1 算法细节

在前向过程中，计算密集激活值但只保留稀疏版本。在后向传播过程中，只使用稀疏激活值计算梯度。

常规的后向传播梯度计算过程为：

对于本文提出的 Back Razor 算法，后向传播过程中梯度计算过程为：

这里是剪枝后的激活值。

这里对于后向激活值剪枝使用一个简单的剪枝策略，剪枝最小幅度值的激活值。最小幅度值与较大幅度值的权重相比一般有较不重要的作用。实际使用时，本文方法以固定的稀疏度剪枝所有的激活值。另外这里考虑对于每个样本分别考虑激活值的排名，并使用样本感知阈值。另外发现不同样本的阈值非常接近，本文的样本感知策略与全局排名接近。

这里的稀疏矩阵由两部分存储：一个位图指示非 0 元素的位置和一个较小的密集矩阵存储非0元素值。位置与原始矩阵有相同形状但只存储二进制类型。

之前的激活值剪枝方法，源于剪枝的错误会在前向过程中逐渐累加。由于本文方法激活值剪枝只用于滋生的梯度计算，剪枝误差不会累积。另外本文方法在相同剪枝率条件下与其他激活值剪枝算法有相同的内存消耗。

2.2 其他方法比较

TinyTL（Reduce activations, not trainable parameters for efficient on-device learning）指出可以通过冻结权重避免储存激活值因为其他项不需要激活值参与。这对于大部分的卷积网络是符合的。但 VIT 模型出现打破了这个假定。很多操作（self-attention，Softmax，GeLU）都设计激活值梯度计算。

DiffPrune（Parameter-efficient transfer learning with diff pruning.）指出只需要调整 0.5% 比例权重微调的权重就可以与完全微调的性能相当。这里使用 Gumbel-Softmax 使遮挡的参数可微。但该方法不能获得内存效率

实验

表1给出了本文的 Back Razor 算法与先前提出的迁移学习方法在卷积神经网络和不同数据集下的性能比较。FT-Last 指微调网络最后一个全连接层的方法。FT-Norm+Last 指微调网络最后一个全连接和 BatchNorm 层的算法。FT-Full 指微调整个网络的算法。

表中结果所示：本文算法在稀疏度 90% 条件下与 FT-Full 算法相比在相近的精度条件下可以获得 8.7 倍小内存。与 FT-Norm+Last 相比获得了 4.6 倍小内存同时精度低 [0.7%, 19.5%, 2.1%, 3.1%, 2.6%, 12.2%, 4.7%, 6.2%]。FT-Last 与本文相比有接近的内存，但精度明显较低：[1.6%, 42.7%, 11.0%, 14.5%, 6.8%, 39.2%, 7.7%, 14.5%] 。

与之前 SOTA 算法 TinyTL 算法相比，内存接近，但本文精度取得了 [1.1%, 12.2%, 1.0%, 1.9%, 1.4%, 5.1%, 3.9%, 3.5%] 的提升。在较大的分辨率（320）情况下，本文方法比 TinyTL 相比减少了 1.5 倍小内存。精度取得了 [0.0%, 5.3%, 0.8%, 1.8%, 0.1%, 1.3%, 0.0%, 0.3%] 提升。

前向传播剪枝的影响：这里研究了在前向和后向传播过程都剪枝的策略（Prune Both Forward & backward）与只在后向传播剪枝的策略（Only Prune Backward）的比较。图 3 给出了不同稀疏度的实验结果。可以看出模型在大稀疏度情况下容易崩溃。只有后向传播剪枝的策略可以在较大的稀疏度（70%）条件下保持较高精度，比较的方法精度已经崩溃（低于 5%）。

表 3 给出了本文方法与先前在 VIT 模型设计的迁移学习方法的性能比较结果。与 FT-Full 相比，在 80% 稀疏度条件下 Back Razor 可以获得相当的性能，获得了 3.9 倍小内存。尽管 FT-Last 获得了更小的内存，但精度获得了 [2.2%, 32.9%, 2.1%, 6.4%, 0.1%, 26.6%, 0.9%, 4.2%] 提升。