[论文解读] Mesa: A Memory-saving Training Framework for Transformers

Zizheng Pan, Peng Chen+4 more|arXiv (Cornell University)|Nov 22, 2021

Advanced Neural Network Applications参考文献 25被引用 4

一句话总结

Mesa 是一种针对视觉 Transformer 的内存高效训练框架，通过在反向传播过程中采用低精度激活量化，将 GPU 内存使用量最多减少 50%，同时保持前向传播的精确性。它引入了按头的量化策略以应对多头注意力中激活分布的异质性，并通过运行估计学习量化参数，从而在不损失性能的前提下支持更大的批量大小或模型规模。

ABSTRACT

There has been an explosion of interest in designing high-performance Transformers. While Transformers have delivered significant performance improvements, training such networks is extremely memory intensive owing to storing all intermediate activations that are needed for gradient computation during backpropagation, especially for long sequences. To this end, we present Mesa, a memory-saving training framework for Transformers. Specifically, Mesa uses exact activations during forward pass while storing a low-precision version of activations to reduce memory consumption during training. The low-precision activations are then dequantized during back-propagation to compute gradients. Besides, to address the heterogeneous activation distributions in the multi-head self-attention layers, we propose a head-wise activation quantization strategy, which quantizes activations based on the statistics of each head to minimize the approximation error. To further boost training efficiency, we learn quantization parameters by running estimates. More importantly, by re-investing the saved memory in employing a larger batch size or scaling up model size, we may further improve the performance under constrained computational resources. Extensive experiments on ImageNet, CIFAR-100 and ADE20K demonstrate that Mesa can achieve flexible memory-savings (up to 50%) during training while achieving comparable or even better performance. Code is available at https://github.com/ziplab/Mesa.

研究动机与目标

• 为解决训练大型视觉 Transformer 时高内存消耗的问题，特别是由于在反向传播期间存储全精度激活所导致的内存开销。
• 通过减少激活内存占用而不牺牲模型性能，使资源有限的 GPU 研究人员能够实现内存高效的训练。
• 开发一种通用的激活量化框架，兼容 Transformer 中所有关键操作，包括矩阵乘法、Softmax、GELU 和 LayerNorm。
• 解决多头自注意力头之间激活分布异质性的问题，该问题会削弱标准量化策略的效果。
• 设计一种与现有技术（如 AMP、检查点技术和梯度量化）正交的训练框架，从而实现互补的内存节省。

提出的方法

• Mesa 使用精确的前向传播激活，但存储低精度（8 位）量化版本以减少训练期间的内存使用。
• 在反向传播期间，将低精度激活反量化以计算梯度，从而保持训练精度。
• 提出一种按头的激活量化策略，其中每个自注意力头拥有独立的量化参数（截断范围 α 和偏移 β），以最小化近似误差。
• 量化参数（α 和 β）通过训练过程中的运行估计学习，避免使用逐样本统计或基于梯度的优化，从而降低开销。
• 该框架支持 Transformer 中所有主要操作：矩阵乘法、Softmax、GELU 和 LayerNorm，实现端到端的内存节省。
• 节省的内存可重新投入用于增大批量大小或模型规模，从而在硬件受限条件下进一步提升性能。

实验结果

研究问题

• RQ1 能否为视觉 Transformer 设计一种内存节省型训练框架，在不降低性能的前提下减少激活内存占用？
• RQ2 与统一或逐层量化相比，按头量化在多头自注意力层中如何提升准确率？
• RQ3 运行估计是否能在计算成本更低的前提下，实现与逐样本统计或基于梯度学习相当或更优的性能？
• RQ4 在实际应用中，Mesa 所节省的内存能在多大程度上用于扩大模型规模或批量大小？
• RQ5 与检查点技术、梯度累积和 AMP 等现有内存节省技术相比，Mesa 在内存减少和训练效率方面表现如何？

主要发现

• Mesa 在 Swin-T、Swin-B 和 DeiT 等视觉 Transformer 上训练时，内存消耗最高可减少 50%，且在 ImageNet、CIFAR-100 和 ADE20K 上无性能下降。
• 与统一量化相比，按头量化显著提升了训练稳定性和性能，尤其在激活分布异质性明显的多头注意力层中表现更优。
• 运行估计在保持性能相当或更优的同时，显著降低了训练开销和内存成本，优于逐样本统计方法。
• 该框架使节省的内存可重新用于增大批量大小或模型规模，从而提升泛化能力并加快收敛速度。
• 在 CIFAR-100 上，Mesa 在压缩部分操作时，准确率略高于基线训练，展现出有利的速度-内存-准确率权衡。
• 可视化结果显示，不同头和层的量化参数（α 和 β）演化方式不同，MSA 层中的 β 值多呈负向分布，表明负激活更为普遍。