Accurate Neural Training with 4-bit Matrix Multiplications at Standard Formats

TL;DR

This paper tackles the challenge of reducing the computational footprint of neural network training by quantizing all GEMMs to 4 bits using standard formats. It introduces Logarithmic Unbiased Quantization (LUQ), a method that combines logarithmic quantization with unbiased stochastic rounding to quantize neural gradients in FP4 while keeping forward weights/activations in INT4, enabling full 4-bit training without hardware-special formats. The key contributions include a detailed LUQ design with a robust handling of underflow/overflow, a theoretical justification for using unbiased gradients in backprop, and practical optional methods (SMP, FNT) to further reduce degradation. Experimental results across multiple models demonstrate state-of-the-art 4-bit training performance (e.g., ~1.1% degradation on ResNet-50, dropping to ~0.32% with FNT and SMP), with overhead comparable to prior approaches, and notable hardware-relevant considerations such as data movement and scaling costs. Overall, the work provides a viable path to practical, standard-format 4-bit training that can significantly cut energy and bandwidth requirements in real-world training pipelines.

Abstract

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.

Figures (9)

• Figure 1: Comparison between stochastic rounding (SR) and round-to-nearest (RDN) quantization. In (a) we present the MSE of a uniform distributed tensor with the two different rounding schemes. Quantization to 4 bits of the activations and weights ($Q_{\mathrm{INT}})$(b) and neural gradients ($Q_{\mathrm{FP}}$) (c) of ResNet18 - Cifar100 dataset with SR and RDN. While MSE is important in $Q_{\mathrm{INT}}$ for the weights and activations, unbiasedness achieved with SR is crucial for the neural gradients in $Q_{\mathrm{FP}}$. The neural gradients in (b) and the weights and activations in (c), are in full precision to focus on the effect of the rounding scheme only in one part of the network in each experiment.
• Figure 2: (Left): ResNet50 top-1 validation accuracy in ImageNet dataset with different quantization schemes for the neural gradients. FP4 refers to standard logarithmic (1-3-0) floating point quantization. SR refers to stochastic rounding, which makes the quantization unbiased inside the FP range. SU refers to stochastic underflow which makes the quantization unbiased below minimum FP. "INT4 unbiased" refers to the combination of INT4 and SR which is fully unbiased. Notice that while biased logarithmic quantization ("FP4"), partially biased logarithmic ("FP4 + SR", "FP4 + SR + SU"), and uniform unbiased ("INT4 unbiased") lead to significant accuracy degradation, the proposed fully unbiased logarithmic quantization ("LUQ") has a minimal degradation. (Right): Illustration of the different logarithmic quantization schemes for FP2 ([1,1,0] format), where $2^0$ and $2^2$ are, respectively, the minimal and maximal FP representations. Two arrows for the same point mean SR - thicker lines represent higher probability. Only LUQ is able to achieve unbiasedness in all floating point ranges.
• Figure 3: Summary of the different formats used in each of the three GEMM in our 4 training regime. Boldface represent matrix and vectors. Different colors represent different elements of each layer (activations, weights, neural gradients, etc.)
• Figure 4: ResNet18 top-1 validation in Cifar10 dataset, with 4-bit quantization of the neural gradients using stochastic-rounding. The amortization is the numbers of iteration that we re-use the random samples.
• Figure 5: ResNet18 top-1 validation accuracy in CIFAR100 with quantization of the neural gradients to 2-bit (FP2 - [1,1,0] format) using different samples numbers to reduce the variance. Notice that 16 samples completely close the gap to the baseline.