[Paper Review] Finding the Task-Optimal Low-Bit Sub-Distribution in Deep Neural Networks

Runpei Dong, Zhanhong Tan+3 more|arXiv (Cornell University)|Dec 30, 2021

Advanced Neural Network Applications4 citations

TL;DR

This paper proposes DGMS, a differentiable Gaussian mixture weight sharing method that automatically learns a task-optimal low-bit sub-distribution within deep neural networks through end-to-end co-tuning of weights and a concrete Gaussian mixture approximation. By directly optimizing the task objective instead of relying on ELBO or manual configurations, DGMS achieves state-of-the-art quantization performance with up to 7.46× inference speedup on mobile CPUs and strong zero-shot transferability across domains.

ABSTRACT

Quantized neural networks typically require smaller memory footprints and lower computation complexity, which is crucial for efficient deployment. However, quantization inevitably leads to a distribution divergence from the original network, which generally degrades the performance. To tackle this issue, massive efforts have been made, but most existing approaches lack statistical considerations and depend on several manual configurations. In this paper, we present an adaptive-mapping quantization method to learn an optimal latent sub-distribution that is inherent within models and smoothly approximated with a concrete Gaussian Mixture (GM). In particular, the network weights are projected in compliance with the GM-approximated sub-distribution. This sub-distribution evolves along with the weight update in a co-tuning schema guided by the direct task-objective optimization. Sufficient experiments on image classification and object detection over various modern architectures demonstrate the effectiveness, generalization property, and transferability of the proposed method. Besides, an efficient deployment flow for the mobile CPU is developed, achieving up to 7.46$ imes$ inference acceleration on an octa-core ARM CPU. Our codes have been publicly released at \url{https://github.com/RunpeiDong/DGMS}.

Motivation & Objective

• To address the performance degradation caused by distribution divergence in low-bit weight quantization.
• To eliminate manual hyperparameter tuning in quantization by learning an optimal sub-distribution automatically.
• To improve generalization and transferability of quantized models across different vision domains.
• To enable efficient deployment on mobile CPUs through a streamlined inference pipeline.
• To reframe quantization as a task-optimization problem by learning a model-inherent sub-distribution instead of relying on fixed priors or ELBO objectives.

Proposed method

• Introduces a differentiable Gaussian Mixture (GM) approximation to model the latent low-bit sub-distribution, replacing the fixed Dirac delta posterior.
• Co-tunes network weights and GM parameters end-to-end by directly minimizing the task loss, avoiding the need for ELBO optimization.
• Employs temperature-based softmax reparameterization to bridge the gap between training and inference in the quantization process.
• Projects weights onto the GM-approximated sub-distribution during training, enabling smooth, differentiable quantization.
• Uses a self-adaptive, automated schema where the GM evolves with weight updates, eliminating manual configuration.
• Deploys the quantized model efficiently on mobile CPUs by preserving parallelism and optimizing cache access.

Experimental results

Research questions

• RQ1 Can a model-inherent, task-optimal low-bit sub-distribution be automatically discovered without manual hyperparameter tuning?
• RQ2 Does direct optimization of the task objective lead to better quantization performance than ELBO-based variational inference?
• RQ3 Can the learned sub-distribution generalize across different vision domains, including zero-shot transfer?
• RQ4 How effective is DGMS in achieving high compression with minimal accuracy drop on both image classification and object detection tasks?
• RQ5 To what extent can DGMS accelerate inference on mobile CPUs without sacrificing model accuracy?

Key findings

• DGMS achieves up to 7.46× inference speedup on an octa-core ARM CPU compared to full-precision models, with minimal accuracy drop.
• On CIFAR-10, 4-bit DGMS achieves 70.25% top-1 accuracy with ResNet-20, outperforming full-precision baselines in some cases.
• With zero-shot transfer, DGMS achieves 77.46% top-1 accuracy on CUB-200-2011 and 81.46% on Stanford Cars, rivaling full-precision models.
• After only one fine-tuning epoch on ImageNet, DGMS reaches 82.28% top-1 accuracy, matching full-precision performance.
• The learned sub-distribution exhibits strong domain-invariant properties, enabling effective zero-shot quantization across diverse datasets.
• The quantization error (measured by MSE) spontaneously decreases during training, correlating with accuracy improvement despite not being used as a training objective.