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

Tensor Methods for Generating Compact Uncertainty Quantification and Deep Learning Models

General-Purpose Bayesian Tensor Learning With Automatic Rank Determination and Uncertainty Quantification

High-Dimensional Uncertainty Quantification via Tensor Regression with Rank Determination and Adaptive Sampling

High-Dimensional Uncertainty Quantification via Active and Rank-Adaptive Tensor Regression

Statistical Inference in Tensor Completion: Optimal Uncertainty Quantification and Statistical-to-Computational Gaps

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

Uncertainty Quantification for Sparse Deep Learning

Tensor Networks for Big Data Analytics and Large-Scale Optimization Problems

Tensor Methods in High Dimensional Data Analysis: Opportunities and Challenges

Generalized Uncertainty of Deep Neural Networks: Taxonomy and Applications

A Tensor-Based Approach for Big Data Representation and Dimensionality Reduction

Uncertainty Quantification for Deep Learning

Tensor Networks for Dimensionality Reduction and Large-Scale Optimizations. Part 2 Applications and Future Perspectives

A Big-Data Approach to Handle Many Process Variations: Tensor Recovery and Applications

Uncertainty quantification in scientific machine learning: Methods, metrics, and comparisons

Scalable Uncertainty Quantification for Deep Operator Networks using Randomized Priors

Learning Multivariate CDFs and Copulas using Tensor Factorization

Bayesian tensor-on-tensor regression with efficient computation

Diffusion Tensor Estimation with Uncertainty Calibration

The Tensor-Train Stochastic Finite Volume Method for Uncertainty Quantification

High-dimensional Quantile Tensor Regression