Mesa：一种面向 Transformer 的省内存训练框架
计算机视觉与模式识别
2022-08-30 v3 机器学习
摘要
设计高性能 Transformer 的兴趣呈现爆发式增长。尽管 Transformer 带来了显著的性能提升，但训练此类网络极为耗费内存，原因在于需存储所有中间激活值以用于反向传播中的梯度计算，尤其对于长序列而言。为此，我们提出 Mesa，一种面向 Transformer 的省内存训练框架。具体而言，Mesa 在前向传播中使用精确激活值，同时存储低精度版本的激活值以降低训练期间的内存消耗。随后在反向传播中对低精度激活值进行反量化以计算梯度。此外，为解决多头自注意力层中激活值分布异构的问题，我们提出了逐头激活量化策略，其基于每个头的统计量对激活值量化以最小化近似误差。为进一步提升训练效率，我们通过运行估计来学习量化参数。更重要的是，通过将节省的内存重新用于采用更大批量或扩大模型规模，我们可在受限计算资源下进一步提升性能。在 ImageNet、CIFAR-100 与 ADE20K 上的大量实验表明，Mesa 在训练期间可实现灵活的内存节省（高达 50%），同时取得可比甚至更优的性能。代码见 https://github.com/ziplab/Mesa。
引用
@article{arxiv.2111.11124,
title = {Mesa: A Memory-saving Training Framework for Transformers},
author = {Zizheng Pan and Peng Chen and Haoyu He and Jing Liu and Jianfei Cai and Bohan Zhuang},
journal= {arXiv preprint arXiv:2111.11124},
year = {2022}
}
备注
Tech report