Probabilistic weight fixing: large-scale training of neural network weight uncertainties for quantisation

Subia-Waud, Christopher and Dasmahapatra, Srinandan (2023) Probabilistic weight fixing: large-scale training of neural network weight uncertainties for quantisation. Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M. and Levine, S. (eds.) In NIPS '23: Proceedings of the 37th International Conference on Neural Information Processing Systems. Curran Associates, Inc. pp. 59410-59424 . (doi:10.5555/3666122.3668718).

Abstract

Weight-sharing quantization has emerged as a technique to reduce energy expenditure during inference in large neural networks by constraining their weights to a limited set of values. However, existing methods often assume weights are treated solely based on value, neglecting the unique role of weight position. This paper proposes a probabilistic framework based on Bayesian neural networks (BNNs) and a variational relaxation to identify which weights can be moved to which cluster center and to what degree based on their individual position-specific learned uncertainty distributions. We introduce a new initialization setting and a regularization term, enabling the training of BNNs with complex dataset-model combinations. Leveraging the flexibility of weight values from probability distributions, we enhance noise resilience and compressibility. Our iterative clustering procedure demonstrates superior compressibility and higher accuracy compared to state-of-the-art methods on both ResNet models and the more complex transformer-based architectures. In particular, our method outperforms the state-of-the-art quantization method top-1 accuracy by 1.6\% on ImageNet using DeiT-Tiny, with its 5 million+ weights now represented by only 296 unique values.

Code available at https://github.com/subiawaud/PWFN.

This record has no associated files available for download.

Contributors

Author: Srinandan Dasmahapatra

Download statistics