Accurate Neural Training with 4-bit Matrix Multiplications at Standard Formats
This work addresses the computational cost of training for machine learning practitioners, offering an incremental improvement by combining existing ideas for 4-bit quantization.
The paper tackles the problem of reducing the computational footprint of deep neural network training by quantizing neural gradients to 4-bit, proposing a logarithmic unbiased quantization (LUQ) method that achieves state-of-the-art results, such as 1.1% degradation in ResNet50 on ImageNet, improvable to 0.32% with fine-tuning.
Quantization of the weights and activations is one of the main methods to reduce the computational footprint of Deep Neural Networks (DNNs) training. Current methods enable 4-bit quantization of the forward phase. However, this constitutes only a third of the training process. Reducing the computational footprint of the entire training process requires the quantization of the neural gradients, i.e., the loss gradients with respect to the outputs of intermediate neural layers. Previous works separately showed that accurate 4-bit quantization of the neural gradients needs to (1) be unbiased and (2) have a log scale. However, no previous work aimed to combine both ideas, as we do in this work. Specifically, we examine the importance of having unbiased quantization in quantized neural network training, where to maintain it, and how to combine it with logarithmic quantization. Based on this, we suggest a $\textit{logarithmic unbiased quantization}$ (LUQ) method to quantize both the forward and backward phases to 4-bit, achieving state-of-the-art results in 4-bit training without the overhead. For example, in ResNet50 on ImageNet, we achieved a degradation of 1.1%. We further improve this to a degradation of only 0.32% after three epochs of high precision fine-tuning, combined with a variance reduction method -- where both these methods add overhead comparable to previously suggested methods.