[Paper Review] Mesa: A Memory-saving Training Framework for Transformers

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

Advanced Neural Network Applications25 references4 citations

TL;DR

Mesa is a memory-efficient training framework for Vision Transformers that reduces GPU memory usage by up to 50% through low-precision activation quantization during backpropagation, while maintaining exact forward passes. It introduces head-wise quantization to handle heterogeneous attention distributions and learns quantization parameters via running estimates, enabling larger batch sizes or model scales without performance loss.

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.

Motivation & Objective

• To address the high memory consumption of training large Vision Transformers, especially due to storing full-precision activations during backpropagation.
• To enable memory-efficient training for researchers with limited GPU resources by reducing activation memory footprint without sacrificing model performance.
• To develop a general-purpose activation quantization framework compatible with all key operations in Transformers, including MatMul, Softmax, GELU, and LayerNorm.
• To address the challenge of heterogeneous activation distributions across multi-head self-attention heads, which undermines standard quantization strategies.
• To design a training framework that is orthogonal to existing techniques like AMP, checkpointing, and gradient quantization, allowing for complementary memory savings.

Proposed method

• Mesa uses exact forward pass activations but stores low-precision (8-bit) quantized versions to reduce memory usage during training.
• During backpropagation, the low-precision activations are dequantized to compute gradients, preserving training accuracy.
• A head-wise activation quantization strategy is proposed, where each self-attention head has its own quantization parameters (clipping range α and offset β) to minimize approximation error.
• Quantization parameters (α and β) are learned via running estimates during training, avoiding per-sample statistics or gradient-based optimization to reduce overhead.
• The framework supports all major operations in Transformers: matrix multiplication, Softmax, GELU, and LayerNorm, enabling end-to-end memory savings.
• The saved memory is re-invested in increasing batch size or model size to further improve performance under constrained hardware.

Experimental results

Research questions

• RQ1 Can a memory-saving training framework be designed for Vision Transformers that reduces activation memory without degrading performance?
• RQ2 How does head-wise quantization improve accuracy compared to uniform or layer-wise quantization in multi-head self-attention layers?
• RQ3 Can running estimates for quantization parameters achieve comparable or better results than per-sample statistics or gradient-based learning with lower computational cost?
• RQ4 To what extent can memory savings from Mesa be leveraged to scale up model size or batch size in practice?
• RQ5 How does Mesa compare to existing memory-saving techniques like checkpointing, gradient accumulation, and AMP in terms of memory reduction and training efficiency?

Key findings

• Mesa reduces memory consumption by up to 50% during training on Vision Transformers like Swin-T, Swin-B, and DeiT, with no performance degradation on ImageNet, CIFAR-100, and ADE20K.
• Head-wise quantization significantly improves training stability and performance compared to uniform quantization, especially in multi-head attention layers with heterogeneous activation distributions.
• Running estimates for quantization parameters achieve comparable or better performance than per-sample statistics while reducing training overhead and memory cost.
• The framework enables re-investment of saved memory into larger batch sizes or model scales, leading to improved generalization and faster convergence.
• On CIFAR-100, Mesa achieves slightly higher accuracy than baseline training when compressing selected operations, demonstrating a favorable speed-memory-accuracy trade-off.
• Visualization shows that quantization parameters (α and β) evolve differently across heads and layers, with β values skewed toward negative values in MSA layers, indicating a prevalence of negative activations.