Precision Neural Network Quantization via Learnable Adaptive Modules

2025-04-24

AI Summary

To address the limitations of fixed scaling factors in quantization-aware training (QAT)—including reduced inference flexibility, poor adaptability to diverse activation distributions, and accuracy degradation—this paper proposes Adaptive Scalable Quantization (ASQ). ASQ introduces three key innovations: (1) a dynamic, trainable scaling module that adapts in real time to layer-wise activation distributions; (2) a POST non-uniform quantization scheme with base $2^{1/2}$, balancing accurate modeling of bell-shaped weight distributions and hardware efficiency; and (3) lookup-table (LUT)-based inference acceleration. Evaluated on 4-bit ResNet-34/ImageNet, ASQ achieves only a 1.2% top-1 accuracy drop relative to the full-precision baseline—outperforming state-of-the-art quantization methods by a significant margin, and even surpassing the baseline in certain configurations. Moreover, ASQ enhances deployment flexibility and operational efficiency without compromising accuracy or hardware compatibility.

Technology Category

• Machine Learning: Hardware-aware ML
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Search and Retrieval-Augmented AI: Efficiency and scalability of Web search engines
• Security and Privacy: Large-scale security measurements
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Quantization Aware Training (QAT) is a neural network quantization technique that compresses model size and improves operational efficiency while effectively maintaining model performance. The paradigm of QAT is to introduce fake quantization operators during the training process, allowing the model to autonomously compensate for information loss caused by quantization. Making quantization parameters trainable can significantly improve the performance of QAT, but at the cost of compromising the flexibility during inference, especially when dealing with activation values with substantially different distributions. In this paper, we propose an effective learnable adaptive neural network quantization method, called Adaptive Step Size Quantization (ASQ), to resolve this conflict. Specifically, the proposed ASQ method first dynamically adjusts quantization scaling factors through a trained module capable of accommodating different activations. Then, to address the rigid resolution issue inherent in Power of Two (POT) quantization, we propose an efficient non-uniform quantization scheme. We utilize the Power Of Square root of Two (POST) as the basis for exponential quantization, effectively handling the bell-shaped distribution of neural network weights across various bit-widths while maintaining computational efficiency through a Look-Up Table method (LUT). Extensive experimental results demonstrate that the proposed ASQ method is superior to the state-of-the-art QAT approaches. Notably that the ASQ is even competitive compared to full precision baselines, with its 4-bit quantized ResNet34 model improving accuracy by 1.2% on ImageNet.

Problem

• Resolves conflict between trainable quantization parameters and inference flexibility
• Dynamically adjusts quantization scaling factors for varying activations
• Proposes non-uniform quantization to handle rigid POT resolution issues

Innovation

• Dynamic scaling factors via trained adaptive modules
• Non-uniform POST quantization for bell-shaped weights
• LUT method maintains computational efficiency