Information theoretic-limits and algorithms for model quantization in high-dimensional learning
Rajarshi Saha
Abstract:
High-dimensional models often have a large memory footprint and must be quantized after training before being deployed on resource-constrained edge devices for inference tasks. In this work, we develop an information-theoretic framework for the problem of quantizing the parameters of a model theta, learned from training data (X,y) generated according to some statistical relationship y = f(X; theta). The learned model, which is an estimate of the d-dimensional latent parameter theta, is constrained to be representable using only Bd bits, where B in (0, infty) is a pre-specified budget. We consider the learning tasks of: (i) linear regression, and (ii) binary classification with linear classifiers, and derive information-theoretic limits for the minimax risk under these settings. We also propose matching upper bounds using randomized embedding-based algorithms which are tight up to constant factors. These lower and upper bounds together characterize the minimum threshold bit-budget required to achieve a performance risk comparable to the unquantized setting. We also propose randomized Hadamard embeddings that yield computationally efficient algorithms which are optimal up to a mild logarithmic factor of the lower bound. Our model quantization strategy can be generalized and we show its efficacy by extending the method and upper-bounds to two-layer ReLU neural networks for non-linear regression. Numerical simulations show the improved performance of our proposed scheme as well as its closeness to the lower bound.