Few-Bit Backward: Quantized Gradients of Activation Functions for Memory Footprint Reduction

TL;DR

This paper introduces a method to quantize activation function gradients to reduce memory usage during neural network training, maintaining performance while significantly lowering memory footprint.

Contribution

It presents a systematic approach for optimal low-bit quantization of activation gradients using dynamic programming, applicable to various nonlinearities and compatible with existing training pipelines.

Findings

Memory footprint is significantly reduced.

Training convergence is maintained.

Applicable to all popular nonlinearities.

Abstract

Memory footprint is one of the main limiting factors for large neural network training. In backpropagation, one needs to store the input to each operation in the computational graph. Every modern neural network model has quite a few pointwise nonlinearities in its architecture, and such operation induces additional memory costs which -- as we show -- can be significantly reduced by quantization of the gradients. We propose a systematic approach to compute optimal quantization of the retained gradients of the pointwise nonlinear functions with only a few bits per each element. We show that such approximation can be achieved by computing optimal piecewise-constant approximation of the derivative of the activation function, which can be done by dynamic programming. The drop-in replacements are implemented for all popular nonlinearities and can be used in any existing pipeline. We confirm the memory reduction and the same convergence on several open benchmarks.