ActNN: Reducing Training Memory Footprint via 2-Bit Activation Compressed Training

2021/04/29 by Jianfei Chen, Lianmin Zheng, Chen, Jianfei +11 · 10 citations
Computer Science · Engineering · Mathematics · #Advanced Memory and Neural Computing #Advanced Neural Network Applications #Computer Vision and Pattern Recognition (cs.CV) #FOS: Computer and information sciences #Ferroelectric and Negative Capacitance Devices #Machine Learning (cs.LG) #Machine Learning (stat.ML) #cs.CV #cs.LG #stat.ML

paper · pdf · doi:10.48550/arxiv.2104.14129

to be published in ICML 2021

openalex publication_date 2021/04/29 · arxiv created 2021/07/06 · arxiv updated 2021/07/07 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

The increasing size of neural network models has been critical for improvements in their accuracy, but device memory is not growing at the same rate. This creates fundamental challenges for training neural networks within limited memory environments. In this work, we propose ActNN, a memory-efficient training framework that stores randomly quantized activations for back propagation. We prove the convergence of ActNN for general network architectures, and we characterize the impact of quantization on the convergence via an exact expression for the gradient variance. Using our theory, we propose novel mixed-precision quantization strategies that exploit the activation's heterogeneity across feature dimensions, samples, and layers. These techniques can be readily applied to existing dynamic graph frameworks, such as PyTorch, simply by substituting the layers. We evaluate ActNN on mainstream computer vision models for classification, detection, and segmentation tasks. On all these tasks, ActNN compresses the activation to 2 bits on average, with negligible accuracy loss. ActNN reduces the memory footprint of the activation by 12x, and it enables training with a 6.6x to 14x larger batch size.

Citations

• Training Deep Nets with Sublinear Memory Cost
• DoReFa-Net: Training Low Bitwidth Convolutional Neural Networks with Low Bitwidth Gradients
• Training Deep Neural Networks with 8-bit Floating Point Numbers
• MMDetection: Open MMLab Detection Toolbox and Benchmark
• Deep High-Resolution Representation Learning for Visual Recognition
• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Improved Baselines with Momentum Contrastive Learning
• Learning Multiple Layers of Features from Tiny Images

Cited by

• FedHybrid: Breaking the Memory Wall of Federated Learning via Hybrid Tensor Management
• LANCE: Low Rank Activation Compression for Efficient On-Device Continual Learning
• QuantV2X: A Fully Quantized Multi-Agent System for Cooperative Perception
• Rethinking the Potential of Layer Freezing for Efficient DNN Training
• Adacc: An Adaptive Framework Unifying Compression and Activation Recomputation for LLM Training
• DAF: An Efficient End-to-End Dynamic Activation Framework for on-Device DNN Training
• DipSVD: Dual-importance Protected SVD for Efficient LLM Compression
• Mesa: A Memory-saving Training Framework for Transformers
• TAH-QUANT: Effective Activation Quantization in Pipeline Parallelism over Slow Network
• Low-Precision Training of Large Language Models: Methods, Challenges, and Opportunities

Related

• Less Memory Means smaller GPUs: Backpropagation with Compressed Activations
• COMET: A Novel Memory-Efficient Deep Learning Training Framework by Using Error-Bounded Lossy Compression
• GACT: Activation Compressed Training for Generic Network Architectures
• Training Compact Neural Networks with Binary Weights and Low Precision Activations
• Few-Bit Backward: Quantized Gradients of Activation Functions for Memory Footprint Reduction