Training with Mixed-Precision Floating-Point Assignments

TL;DR

This paper introduces a technique for assigning mixed-precision levels to tensors during neural network training, reducing memory usage while maintaining or improving accuracy compared to prior low-precision methods.

Contribution

It presents a novel method for generating precision assignments that optimize memory-accuracy tradeoffs in CNN training, outperforming existing approaches.

Findings

Typically achieves > 2x memory reduction

Maintains training accuracy with reduced precision

Avoids divergence issues seen in other baselines

Abstract

When training deep neural networks, keeping all tensors in high precision (e.g., 32-bit or even 16-bit floats) is often wasteful. However, keeping all tensors in low precision (e.g., 8-bit floats) can lead to unacceptable accuracy loss. Hence, it is important to use a precision assignment -- a mapping from all tensors (arising in training) to precision levels (high or low) -- that keeps most of the tensors in low precision and leads to sufficiently accurate models. We provide a technique that explores this memory-accuracy tradeoff by generating precision assignments for convolutional neural networks that (i) use less memory and (ii) lead to more accurate convolutional networks at the same time, compared to the precision assignments considered by prior work in low-precision floating-point training. We evaluate our technique on image classification tasks by training convolutional networks on CIFAR-10, CIFAR-100, and ImageNet. Our method typically provides > 2x memory reduction over a baseline precision assignment while preserving training accuracy, and gives further reductions by trading off accuracy. Compared to other baselines which sometimes cause training to diverge, our method provides similar or better memory reduction while avoiding divergence.