面向设备训练的视觉 Transformer 块级选择性重编程
计算机视觉与模式识别
2024-05-21 v1 机器学习
摘要
视觉 Transformer（ViT）的普及已广泛应用于包括个性化学习在内的各种边缘应用，催生了对设备端微调的需求。然而，由于边缘设备的有限内存和计算资源，进行训练仍是一个重大挑战。特别是，训练所需的内存远高于推理需求，这主要是因为需要存储所有层的激活值，以便计算所需权重更新的梯度。以往的工作通过冻结权重训练或存储压缩格式的激活来降低内存需求，但这些方法效率低下，无法提供训练或推理加速。本文首先探讨了旨在降低内存和计算需求的现有设备训练方法的局限性。随后，我们提出块级选择性重编程（BSR），该方法仅对预训练模型的部分块进行微调，并基于冻结层的自注意力得分选择性地丢弃标记。为展示 BSR 的有效性，我们在 ViT-B 和 DeiT-S 上进行了广泛评估，使用五个不同数据集。与现有替代方法相比，我们的方法在保持类似准确率的同时，训练内存降低最高可达1.4倍，计算成本降低最高可达2倍。我们还展示了针对 Mixture-of-Expert（MoE）模型的结果，证明了该方法在多任务学习场景中的有效性。
引用
@article{arxiv.2405.10951,
title = {Block Selective Reprogramming for On-device Training of Vision Transformers},
author = {Sreetama Sarkar and Souvik Kundu and Kai Zheng and Peter A. Beerel},
journal= {arXiv preprint arXiv:2405.10951},
year = {2024}
}