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

Cole Hawkins, Xing Liu+1 more|arXiv (Cornell University)|Oct 17, 2020

Tensor decomposition and applications4 citations

TL;DR

This paper proposes an end-to-end Bayesian tensorized training framework that automatically determines optimal tensor ranks during training, enabling ultra-compact neural networks from scratch. By combining variational inference with low-rank tensor decompositions (CP, Tucker, TT), it reduces model parameters by up to 26,000× (e.g., from 4.25B to 1.6×10⁵) with minimal accuracy loss, enabling energy-efficient on-device training.

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 imes 10^9$ model parameters, our method can reduce the variables to only $1.6 imes 10^5$ automatically in the training process (i.e., by $2.6 imes 10^4$ times) while achieving almost the same accuracy.

Motivation & Objective

• Address the high computational and energy cost of training large, over-parameterized deep neural networks.
• Overcome the challenge of manually selecting optimal tensor ranks in low-rank factorization, which is critical for model accuracy and efficiency.
• Enable direct training of compact neural networks from scratch to eliminate expensive full-precision pre-training.
• Develop a scalable, end-to-end framework that supports multiple tensor formats (CP, Tucker, tensor train) with automatic rank adaptation.
• Achieve high compression ratios and low inference costs while maintaining or improving model accuracy on large-scale models.

Proposed method

• Formulate a Bayesian probabilistic model over low-rank tensor decompositions (CP, Tucker, tensor-train, TTM) to enable automatic rank determination during training.
• Develop a customized stochastic variational inference algorithm to efficiently optimize the posterior distribution over tensor factors and ranks.
• Introduce a hierarchical prior structure with automatic relevance determination (ARD) to enable sparsity-inducing rank adaptation.
• Use a projected kernelized Stein discrepancy to improve posterior approximation in high-dimensional, structured parameter spaces.
• Design a differentiable training pipeline that avoids explicit construction of full-sized embedding tables, reducing memory overhead in large-scale models.
• Integrate the framework with end-to-end training of deep learning recommendation models (DLRM) to demonstrate scalability on billion-parameter models.

Experimental results

Research questions

• RQ1 Can a Bayesian tensorized training framework automatically determine optimal tensor ranks during end-to-end training without prior hyperparameter tuning?
• RQ2 To what extent can automatic rank determination reduce model parameters while preserving or improving predictive accuracy in large-scale deep learning models?
• RQ3 How does end-to-end tensorized training compare to post-training compression in terms of memory, energy, and inference efficiency?
• RQ4 Can the proposed method enable practical on-device training of compact models on resource-constrained platforms like FPGAs?
• RQ5 What is the impact of different tensor formats (CP, Tucker, TT) on compression ratio, accuracy, and training efficiency in real-world benchmarks?

Key findings

• The proposed method reduced model parameters from 4.25 billion to 1.6×10⁵ in a large-scale DLRM, achieving a 26,000× compression ratio with minimal accuracy loss.
• On the DLRM benchmark, the ARD-LU variant achieved 78.67% accuracy with only 2.27 million training parameters, outperforming fixed-rank baselines despite lower parameter count.
• The method achieved up to 123× higher energy efficiency and 59× speedup on FPGA compared to non-tensorized training on embedded CPUs.
• Automatic rank determination via ARD led to better generalization than fixed-rank methods, with ARD-HC achieving 78.73% accuracy on the TT format with only 163,976 final parameters.
• The framework successfully trained ultra-compact models without explicitly forming full-sized embedding tables, avoiding the computational cost of post-training compression.
• The Bayesian solver enabled stable, scalable training of billion-parameter models using low-rank tensorization with automatic rank adaptation.