Learning from Loss Landscape: Generalizable Mixed-Precision Quantization via Adaptive Sharpness-Aware Gradient Aligning

TL;DR

This work tackles the high search cost of mixed-precision quantization (MPQ) by learning transferable MPQ policies on small proxy datasets and generalizing them to large-scale targets. It introduces Adaptive Sharpness-Aware Gradient Aligning (ASGA), which combines sharpness-aware minimization, implicit gradient alignment, and an adaptive perturbation radius to produce flatter loss landscapes and better cross-domain generalization. The authors provide theoretical guarantees for generalization and convergence and demonstrate that proxy-based MPQ policy search achieves ImageNet-equivalent accuracy with up to 150% faster search and strong results on ImageNet and VOC across multiple architectures. By exploiting loss landscape information for transferability, the approach reduces MPQ search cost while maintaining or improving accuracy, enabling more practical deployment on diverse hardware and datasets.

Abstract

Mixed Precision Quantization (MPQ) has become an essential technique for optimizing neural network by determining the optimal bitwidth per layer. Existing MPQ methods, however, face a major hurdle: they require a computationally expensive search for quantization policies on large-scale datasets. To resolve this issue, we introduce a novel approach that first searches for quantization policies on small datasets and then generalizes them to large-scale datasets. This approach simplifies the process, eliminating the need for large-scale quantization fine-tuning and only necessitating model weight adjustment. Our method is characterized by three key techniques: sharpness-aware minimization for enhanced quantization generalization, implicit gradient direction alignment to handle gradient conflicts among different optimization objectives, and an adaptive perturbation radius to accelerate optimization. Both theoretical analysis and experimental results validate our approach. Using the CIFAR10 dataset (just 0.5\% the size of ImageNet training data) for MPQ policy search, we achieved equivalent accuracy on ImageNet with a significantly lower computational cost, while improving efficiency by up to 150% over the baselines.

Figures (8)

• Figure 1: Comparison of the generalization performance between the baseline MPQ methods and ASGA on ResNet18 with CIFAR10. $\sigma$ serves as a measure of the sharpness of the loss landscape (defined in Section \ref{['2.2.2']}) and Generalization means the change in terms of Top-1 accuracy of the model on CIFAR10 and ImageNet. Compared to the baselines (a and c), ASGA significantly reduces surrogate gap (i.e., the difference between the perturbed loss $\mathcal{L}_p(\theta)$ and the experience loss $\mathcal{L}(\theta)$), smoothing the sharpness of the loss landscape, thereby enhancing the model’s generalization on the target dataset (b and d).
• Figure 2: The illustration of our approach. We aim to search for an optimal quantization policy with a flat loss landscape on a proxy dataset, which can be applied to large-scale target datasets. In the policy searching stage, we seek a MPQ policy with a flat loss landscape by minimizing the empirical loss, complexity loss, and surrogate gap. In the deployment stage, the searched MPQ policy can be directly applied to model inference on large target datasets.
• Figure 3: (b) exhibits a flatter landscape compared to (a), as indicated by its smaller $h(\theta_2)$. Nevertheless, SAM favors (a) due to its lower local minimum $\mathcal{L}(\theta_1)$.
• Figure 4: Comparison of Top-1 accuracy (-- $\cdot$ --) and convergence epochs (-- $\cdot$ --) of ResNet18, ResNet50, and MobileNet-V2 on CIFAR10 with different $\rho$. The result shows that adaptive $\rho$ can significantly reduce the search cost without performance deterioration.
• Figure 5: Bitwidth assignment for each layer of ResNet18, ResNet50, and MobileNet-V2