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

©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。受到激活值存储只用于后向传播的启发，本文提出了非对称剪枝策略，即旨在后向传播过程中使用稀疏的激活层权重。