Towards Compact Neural Networks via End-to-End Training: A Bayesian Tensor Approach with Automatic Rank Determination

While post-training model compression can greatly reduce the inference cost of a deep neural network, 5 uncompressed training still consumes a huge amount of hardware resources, run-time and energy. It 6 is highly desirable to directly train a compact neural network from scratch with low memory and low 7 computational cost. Low-rank tensor decomposition is one of the most effective approaches to reduce 8 the memory and computing requirements of large-size neural networks. However, directly training 9 a low-rank tensorized neural network is a very challenging task because it is hard to determine 10 a proper tensor rank a priori, which controls the model complexity and compression ratio in the 11 training process. This paper presents a novel end-to-end framework for low-rank tensorized training 12 of neural networks. We first develop a flexible Bayesian model that can handle various low-rank 13 tensor formats (e.g., CP, Tucker, tensor train and tensor-train matrix) that compress neural network 14 parameters in training. This model can automatically determine the tensor ranks inside a nonlinear 15 forward model, which is beyond the capability of existing Bayesian tensor methods. We further 16 develop a scalable stochastic variational inference solver to estimate the posterior density of large17 scale problems in training. Our work provides the first general-purpose rank-adaptive framework 18 for end-to-end tensorized training. Our numerical results on various neural network architectures 19 show orders-of-magnitude parameter reduction and little accuracy loss (or even better accuracy) in 20 the training process. Specifically, on a very large deep learning recommendation system with over 21 4.2× 10 model parameters, our method can reduce the variables to only 1.6× 10 automatically in 22 the training process (i.e., by 2.6 × 10 times) while achieving almost the same accuracy. 23