Schrödinger's FP: Dynamic Adaptation of Floating-Point Containers for Deep Learning Training

TL;DR

This work addresses the high energy and latency cost of memory transfers during deep learning training by proposing Schrödinger's FP, a family of methods that dynamically adapt floating-point containers for activations and weights. It introduces per-tensor and time-varying mantissa/exponent sizing through Quantum Mantissa, Quantum Exponent, BitWave, and Gecko, achieving large memory-footprint reductions while preserving accuracy. The approach is validated with extensive training and a hardware-aware evaluation showing notable improvements in energy efficiency and throughput over traditional FP32, BFloat16, and FP8 baselines. This automatic, closed-loop datatype optimization lowers the trial-and-error burden and informs future hardware/software co-design for leaner training pipelines.

Abstract

The transfer of tensors from/to memory during neural network training dominates time and energy. To improve energy efficiency and performance, research has been exploring ways to use narrower data representations. So far, these attempts relied on user-directed trial-and-error to achieve convergence. We present methods that relieve users from this responsibility. Our methods dynamically adjust the size and format of the floating-point containers used for activations and weights during training, achieving adaptivity across three dimensions: i) which datatype to use, ii) on which tensor, and iii) how it changes over time. The different meanings and distributions of exponent and mantissas lead us to tailored approaches for each. We present two lossy pairs of methods to eliminate as many mantissa and exponent bits as possible without affecting accuracy. Quantum Mantissa and Quantum Exponent are machine learning compression methods that tap into the gradient descent algorithm to learn the minimal mantissa and exponent bitlengths on a per-layer granularity. They automatically learn that many tensors can use just 1 or 2 mantissa bits and 3 or 4 exponent bits. Overall, the two machine learning methods reduce the footprint by $4.74\times$. Alternatively, BitWave observes changes in the loss function during training to adjust mantissa and exponent bitlengths network-wide, yielding a $3.19\times$ reduction in footprint. Finally, we present an optional method, Gecko, to exploit the naturally emerging, lop-sided exponent distribution to losslessly compress resulting exponents from Quantum Exponent or BitWave and, on average, improve compression rates to $5.64\times$ and $4.56\times$.

Figures (12)

• Figure 1: QM and QE on ResNet18/ImageNet throughout training: (a) Validation accuracy, (b) Weighted mantissa bitlengths with their spread, (c) Weighted exponent bitlengths with their spread, and (d) Weighted total bitlengths with their spread.
• Figure 2: Quantum Mantissa and Quantum Exponent on ResNet18/ImageNet: mantissa ($+$), exponent ($\times$), and total ($\cdot$) bitlength datatypes of each tensor at the end of each epoch. Darker colors indicate multiple occurrences.
• Figure 3: Schrödinger's FP: Relative training footprint of ResNet18 with FP32, BFloat16, $\mathit{\mathrm SFP}_\mathit{\mathrm BW}$ and $\mathit{\mathrm SFP}_\mathit{\mathrm Q}$.
• Figure 4: BitWave on ResNet18/ImageNet: (a) Validation accuracy throughout training, (b) Average mantissa and exponent bitlengths per epoch throughout training, (c) Distribution of BitWave's mantissa bitlengths throughout the 5005 batches of epoch 5 of training.
• Figure 5: Gecko on ResNet18/ImageNet: (a) Cumulative distribution of exponent values. (b) Post-encoding cumulative distribution of exponent bitlength.