FLIQS: One-Shot Mixed-Precision Floating-Point and Integer Quantization Search

TL;DR

FLIQS tackles the challenge of automatically selecting per-layer mixed-precision formats for both integer and low-precision floating-point networks under real hardware costs. It introduces a one-shot reinforcement-learning framework with a cosine entropy regularization schedule to search quantization formats and optionally neural architectures without retraining. The approach yields state-of-the-art results on ImageNet across ResNet and MobileNetV2 for both integer and FP8 formats and extends to joint quantization and NAS (FLIQNAS), achieving Pareto-optimal accuracy-cost tradeoffs. Overall, FLIQS enables efficient hardware-aware co-design and scalable exploration of large mixed-precision search spaces.

Abstract

Quantization has become a mainstream compression technique for reducing model size, computational requirements, and energy consumption for modern deep neural networks (DNNs). With improved numerical support in recent hardware, including multiple variants of integer and floating point, mixed-precision quantization has become necessary to achieve high-quality results with low model cost. Prior mixed-precision methods have performed either a post-training quantization search, which compromises on accuracy, or a differentiable quantization search, which leads to high memory usage from branching. Therefore, we propose the first one-shot mixed-precision quantization search that eliminates the need for retraining in both integer and low-precision floating point models. We evaluate our search (FLIQS) on multiple convolutional and vision transformer networks to discover Pareto-optimal models. Our approach improves upon uniform precision, manual mixed-precision, and recent integer quantization search methods. With integer models, we increase the accuracy of ResNet-18 on ImageNet by 1.31% and ResNet-50 by 0.90% with equivalent model cost over previous methods. Additionally, for the first time, we explore a novel mixed-precision floating-point search and improve MobileNetV2 by up to 0.98% compared to prior state-of-the-art FP8 models. Finally, we extend FLIQS to simultaneously search a joint quantization and neural architecture space and improve the ImageNet accuracy by 2.69% with similar model cost on a MobileNetV2 search space.

Figures (26)

• Figure 1: FLIQS -- The explosion of model architectures, numerical support, and deployment platforms requires automated methods for searching model configurations to utilize platform-specific numerical formats. We establish FLIQS as the first one-shot quantization and neural architecture search framework for searching for integer and floating point formats.
• Figure 2: FLIQS Overview -- (a) FLIQS begins with pure training to allow the reward signal to stabilize before updating its policy. The activation quantization is delayed to allow the activation statistics to stabilize. (b) The RL then controller proposes per-layer formats and architectural decisions during training
• Figure 3: FLIQS Examples -- In these quantization search examples, FLIQS allocates more precision to the first and last layers and the small pointwise convolutions of ResNet-18, and to the attention block within DeiT-B16. More configurations are listed in Appendix \ref{['sec:config']}.
• Figure 4: FLIQS Analysis -- (a) The switching error grows relatively large when either bitwidth is small and affects model convergence. In addition, the optimal clipping threshold depends on the current bitwidth, which motivates swapping thresholds. (b) Accuracy improves for higher entropy regularization, and the entropy regularization affects the policy convergence.
• Figure 5: ImageNet FLIQS Quantization Search -- FLIQS reaches higher accuracies at lower costs, and in general FLIQS-L achieves higher accuracies. Models are evaluated at multiple widths ranging .25$\times$ to 2$\times$ of their original channel width to generate each data point.