[论文解读] 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 applications被引用 4

一句话总结

本文提出了一种端到端的贝叶斯张量化训练框架，可在训练过程中自动确定最优张量秩，从而从零开始实现超紧凑的神经网络。通过结合变分推断与低秩张量分解（CP、Tucker、TT），该方法将模型参数量减少高达26,000倍（例如，从4.25B降至1.6×10⁵），且精度损失极小，实现了面向设备端的节能训练。

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.

研究动机与目标

• 解决训练大规模、过参数化的深度神经网络所带来的高计算与能耗成本问题。
• 克服在低秩分解中手动选择最优张量秩的挑战，该选择对模型精度与效率至关重要。
• 实现从零开始直接训练紧凑神经网络，从而消除昂贵的全精度预训练过程。
• 开发一种可扩展的端到端框架，支持多种张量格式（CP、Tucker、张量列车）并具备自动秩自适应能力。
• 在大规模模型上实现高稀疏化比率与低推理成本，同时保持或提升模型精度。

提出的方法

• 在低秩张量分解（CP、Tucker、张量列车、TTM）上构建贝叶斯概率模型，以实现在训练过程中自动确定张量秩。
• 开发一种定制化的随机变分推断算法，以高效优化张量因子与秩的后验分布。
• 引入具有自动相关性确定（ARD）的分层先验结构，以实现诱导稀疏性的秩自适应。
• 采用投影核化Stein差异（projected kernelized Stein discrepancy）以提升高维、结构化参数空间中的后验近似质量。
• 设计一种可微分的训练流水线，避免显式构建完整大小的嵌入表，从而降低大规模模型中的内存开销。
• 将该框架集成至深度学习推荐模型（DLRM）的端到端训练中，以在百亿参数模型上验证其可扩展性。

实验结果

研究问题

• RQ1 贝叶斯张量化训练框架是否能在无需超参数调优的情况下，在端到端训练过程中自动确定最优张量秩？
• RQ2 自动秩确定在多大程度上能减少大规模深度学习模型的参数量，同时保持或提升预测精度？
• RQ3 与训练后压缩相比，端到端张量化训练在内存占用、能耗和推理效率方面表现如何？
• RQ4 所提出的方法是否能实现在FPGA等资源受限平台上的紧凑模型实用化设备端训练？
• RQ5 在真实世界基准测试中，不同张量格式（CP、Tucker、TT）对压缩率、精度和训练效率有何影响？

主要发现

• 在大规模DLRM模型中，该方法将参数量从42.5亿减少至1.6×10⁵，实现了26,000倍的压缩率，且精度损失极小。
• 在DLRM基准上，ARD-LU变体以仅227万训练参数实现了78.67%的准确率，优于参数量更高的固定秩基线模型。
• 与嵌入式CPU上的非张量化训练相比，该方法在FPGA上实现了最高123倍的能效提升和59倍的加速比。
• 通过ARD实现的自动秩确定在泛化性能上优于固定秩方法，ARD-HC在TT格式上仅用163,976个最终参数即达到78.73%的准确率。
• 该框架成功实现了无需显式构建完整嵌入表的超紧凑模型训练，避免了训练后压缩的计算开销。
• 贝叶斯求解器实现了稳定且可扩展的大规模百亿参数模型训练，结合低秩张量化与自动秩自适应。