面向紧凑神经网络端到端训练的贝叶斯张量方法与自动秩确定

虽然训练后模型压缩可大幅降低深度神经网络的推理成本，但未压缩的训练仍消耗大量硬件资源、运行时间与能量。非常期望能以低内存和低计算成本直接从零开始训练紧凑神经网络。低秩张量分解是减少大型神经网络内存与计算需求的最有效方法之一。然而，直接训练低秩张量化神经网络是一项极具挑战的任务，因为难以先验地确定合适的张量秩，它控制着训练过程中的模型复杂度与压缩比。本文提出一种用于神经网络低秩张量化训练的端到端新框架。我们首先开发了一个灵活的贝叶斯模型，可处理多种低秩张量格式（如 CP、Tucker、张量列和张量列矩阵）以在训练中压缩神经网络参数。该模型能在非线性前向模型中自动确定张量秩，这超出了现有贝叶斯张量方法的能力。我们进一步开发了一种可扩展的随机变分推断求解器，用于估计训练中大规模问题的后验密度。我们的工作提供了首个用于端到端张量化训练的通用秩自适应框架。我们在多种神经网络架构上的数值结果表明，训练过程中参数减少了数个数量级且精度损失极小（甚至更高）。具体而言，在一个拥有超过 $4.2\times 10^9$ 个模型参数的超大型深度学习推荐系统上，我们的方法能在训练过程中将变量自动减少至仅 $1.6\times 10^5$（即减少 $2.6\times 10^4$ 倍），同时达到几乎相同的精度。