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

Journal Article · Thu Jan 13 00:00:00 EST 2022 · SIAM Journal on Mathematics of Data Science

DOI:https://doi.org/10.1137/21M1391444· OSTI ID:2889955

Hawkins, Cole ^[1]; Liu, Xing ^[2]; Zhang, Zheng ^[1]

---

1. University of California, Santa Barbara, CA (United States)
2. Facebook AI Systems Hardware/Software Co-design, Menlo Park, CA (United States)

Post-training model compression can reduce the inference costs of deep neural networks, but uncompressed training still consumes enormous hardware resources and energy. To enable low-energy training on edge devices, it is highly desirable to directly train a compact neural network from scratch with a low memory cost. Low-rank tensor decomposition is an effective approach to reduce the memory and computing costs of large neural networks. However, directly training low-rank tensorized neural networks is a very challenging task because it is hard to determine a proper tensor rank a priori, and the tensor rank controls both model complexity and accuracy. Here, this paper presents a novel end-to-end framework for low-rank tensorized training. We first develop a Bayesian model that supports various low-rank tensor formats (e.g., CANDECOMP/PARAFAC, Tucker, tensor-train, and tensor-train matrix) and reduces neural network parameters with automatic rank determination during training. Then we develop a customized Bayesian solver to train large-scale tensorized neural networks. Our training methods shows orders-of-magnitude parameter reduction and little accuracy loss (or even better accuracy) in the experiments. On a very large deep learning recommendation system with over 4.2 ×10⁹ model parameters, our method can reduce the parameter number to 1.6 ×10⁵ automatically in the training process (i.e., by 2.6 ×10⁴ times) while achieving almost the same accuracy. Code is available at https://github.com/colehawkins/bayesian-tensor-rank-determination.

• Research Organization: University of California, Santa Barbara, CA (United States)
• Sponsoring Organization: National Science Foundation (NSF); USDOE
• Grant/Contract Number: SC0021323
• OSTI ID: 2889955
• Journal Information: SIAM Journal on Mathematics of Data Science, Journal Name: SIAM Journal on Mathematics of Data Science Journal Issue: 1 Vol. 4; ISSN 2577-0187
• Publisher: Society for Industrial & Applied Mathematics (SIAM)
• Country of Publication: United States
• Language: English

Related Subjects

compact neural networks
compressed training
low-rank tensor
tensor rank determination