【Towards Efficient Tensor Decomposition-Based DNN Model Compression withOptimization Framework】论文

基于张量分解的DNN模型压缩

摘要

先进的张量分解技术，比如tensor train(TT)和tensor ring(TR)。然而，使用TT/TR压缩卷积神经网络(CNNs)总是遭受显著的精度损失。 本文提出了一个使用交替方向乘子法的(ADMM)基于张量分解的模型压缩的系统框架。将基于TT分解的模型压缩构造为具有张量秩约束的优化问题（an optimization problem with constraints on tensor ranks），我们利用ADMM技术以迭代的方式(in a iterative way)系统地解决这个优化问题。在此过程中，整个DNN模型以原始结构而不是TT格式进行训练，但逐渐获得所需的低张量秩特征。然后，我们将该未压缩的模型分解为TT格式，并对其进行微调，最终获得一个高精度的TT格式的DNN模型。我们的框架是非常通用的，它同时适用于cnn和rnn，并且可以很容易地修改以拟合其他张量分解方法。我们在不同的DNN模型上评估了我们提出的框架，包括图像分类和视频识别任务。实验结果表明，我们基于ADMM的TT格式模型（ADMM-based TT-format models）具有很高的压缩性能并保证高精度。特别是，在CIFAR-100上，2.3x和2.4x的压缩率，我们的模型相比原来的ResNet-20和ResNet-32分别有1.96%和2.21%的top-1准确率提升。压缩在ImageNet上的ResNet-18，我们的模型实现了在无精度损失下2.47x的FLOPs降低。

1、引言

 模型压缩中的张量分解（Tensor Decomposition）

 与其他模型压缩方法不同，张量分解可以独特地可以提供超高的压缩比，特别是对于循环神经网络(RNN)模型。先进的张量分解方法(如TT和TR)（48,38），可以给RNN模型的输入到隐藏层（input-to-hidden）带来1000x以上参数缩减，甚至可以显著提高视频识别任务中相应的分类精度。此外，为了充分利用这些模型所提供的优势，已经开发了几种TT格式的DNN硬件加速器，并适配了不同的芯片格式，比如digital CMOS ASIC，memristor ASIC和IoT board。

前沿技术的局限

尽管张量分解具有很好的潜力，但作为一种成熟的模型压缩方法，其性能还不够满足需求。目前，所有报道的张量分解成功仅局限于在视频识别任务中压缩RNN模型。所有最先进的张量分解方法，包括TT和TR，在面对CNN都遭受了非常显著的精度损失。对于CIFAR-10上的ResNet-32，即使最先进的使用TR的方法，在压缩2.7x的时候也损失了1%，5.8x时损失了1.9%。

为什么会限制性能

上述张量分解的局限性主要是由于训练张量分解的DNN模型存在独特的挑战。一般情况说，用张量分解得到压缩模型有两种方法：1）以分解后的格式从头开始进行训练；2)分解一个预先训练过的未压缩模型，然后再训练。在前一种情况下，如TT格式模型时，由于模型在训练前结构已经预先设置为低张量秩格式，与全秩结构相比，相应的模型容量通常有限，因此使得训练过程对初始化非常敏感，具有更高的精度挑战性。在后一种情况下，虽然预训练的未压缩模型提供了良好的初始化位置，但通过将全秩未压缩模型直接分解为低张量秩格式，会导致不可避免且不可忽略的近似误差，即使经过长时间的再训练，这些近似误差仍然很难恢复。此外，无论采用哪种训练策略，张量分解总是会导致网络深度的线性增加，这意味着训练张量分解格式的DNN通常更容易出现梯度消失问题，因此难以得到很好的训练。

Technical Preview and Contributions

为了克服目前张量分解的局限性，充分释放其模型压缩的潜力，本文提出了一种利用交替方向乘数法(ADMM)提出的基于张量分解的模型压缩系统框架。通过将基于TT分解的模型压缩构造为具有张量秩约束的优化问题，我们利用ADMM技术系统地迭代解决该优化问题。在此过程中，整个DNN模型以原始结构而不是TT格式进行训练，但逐渐获得所需的低张量秩特征。然后，我们将该未压缩的模型分解为TT格式，并对其进行微调，最终获得一个高精度的TT格式的DNN模型。综上所述，本文的贡献总结如下：

• 我们提出了一个系统的框架来制定和解决基于张量分解的模型压缩问题。通过将这个问题表述为一个受约束的非凸优化问题，我们的框架逐渐将DNN模型限制在目标张量秩（target tensor ranks）上，而没有直接在TT格式上训练，从而保持了模型的容量，避免了巨大的近似误差和增加的网络深度。
• 我们建议使用ADMM通过单独解决两个子问题来有效地解决这个重新表述的优化问题：一是通过随机梯度下降正则化DNN直接优化损失函数，二是利用引入的投影（introduced projection）对张量秩进行解析约束（constraint the tensor ranks analytically）。
• 我们在不同的DNN模型上评估了我们提出的框架，用于图像分类和视频识别任务。实验结果表明，我们基于ADMM的TT格式模型具有很高的压缩性能和高精度。值得注意的是，在CIFAR-100上，在压缩比为2.3x和2.4x的情况下，我们的模型分别比原来的ResNet-20和ResNet-32的top-1精度高1.96%和2.21%。对于在ImageNet上的ResNet-18，我们的模型实现了2.47x压缩，没有精度损失。

2、相关工作

稀疏（Sparsigication） ：

稀疏化是最流行的DNN压缩方法。不同层次的网络结构可以是稀疏的，比如权重，过滤器和通道。为了获得稀疏性，可以对DNN模型进行修剪或者用稀疏性感知的正则化进行训练。此外，引入的稀疏性既可以是结构化的，也可以是非结构化的。非结构化稀疏模型虽然精度高，压缩比高，但会给底层硬件平台带来内存访问不规则和工作负载不平衡的问题。结构化稀疏模型对硬件更友好；然而，它们的压缩比和精度通常不如非结构化模型。

量化（Quantization）:

张量分解（Tensor Decomposition）:

基于张量理论,张量分解方法将权值张量分解成更小的张量，以减小模型大小。在[24]中，利用矩阵导向奇异值分解(SVD)作为张量分解的低维实例进行模型压缩。然而，使用该方法或其他经典的高维张量分解方法，如Tucker[42]和CP分解[18]，在有限的压缩比情况下，会导致显著的精度损失（>0.5%）。从[12]开始，先进的张量分解方法，如张量序列(TT)和张量环(TR)分解，已经成为更流行的选择。这些方法具有非常诱人的优点--压缩比可以非常高（例如>1000x），因为它们具有独特的数学性质。这些好处已经在视频识别任务中的RNN压缩中得到了证明。如[48,38]中所说，可通过对RNN模型的输入-隐层（input-to-hidden layer）进行TT或TR分解实现17.560x~34203x压缩比，在视频识别中。然而，TT和TR方法在CNN模型上的表现并不好。例如，即使是最近的进展，[43,32]在2.7x压缩比下仍然遭受1.0%的精度损失，在5.8x压缩比下甚至损失1.9%，对于CIFAR-10数据集上的ResNet-32模型。从实际部署的角度来看，这种不可忽略的精度下降严重阻碍了张量分解在许多CNN模型压缩场景中的广泛采用。

3、背景和初步研究（Background and Preliminaries）

3.1. Notation

3.2. Tensor Train (TT) Decomposition (TT分解)

给出一个tensor，它可以被分解成如下一系列三阶张量通过Tensor Train Decomposition(TTD)：

 3.3. TT格式的DNN

TT 全连接层

TT卷积层

在TT格式的DNN上训练

TT-FC，TT-CONV和相应的前向传播方案，标准的随机梯度下降（SGD）能被用于更新将秩设为r的TT-cores，r决定了目标的压缩比。TT-cores的初始化可以随机设置或者经过直接TT分解预训练的未压缩模型得到。

4. 系统压缩框架（Systematic Compression Framework）

对吸纳有TT格式DNN训练的分析 

 如上一段所述，目前有一个TT格式的DNN是：1）用随机初始化的张量核进行训练而来；或者2)对预训练过的模型的直接分解进行训练得到。对于第一种策略，它不利用有关高精度未压缩模型的任何信息；而其他的模型压缩方法，如修剪和知识蒸馏，已经表明，正确利用预先训练好的模型对于DNN压缩是非常重要的。对于第二种策略，虽然确实利用了预训练模型的知识，但由于预训练的模型通常缺乏TT低 秩性质，在直接低秩张量分解后，近似误差太大，即使使用长时间的再训练也无法正确恢复。因此，现有训练策略的固有局限性导致压缩TT格式DNN模型的精度损失。

我们的关键Idea.

 我们认为，克服这些限制的关键是最大限度地保留未压缩模型中所包含的知识，或者换句话说，在给定的目标张量秩下，使张量分解后的近似误差最小化。为此，我们提出了一个使具有低张量秩约束的未压缩模型的损失函数最小化的优化问题。通过适当的先进优化技术(如ADMM)-正则化训练过程，未压缩的DNN模型可以逐渐表现出低张量秩的性质。经过ADMM正则化训练阶段后，显式低秩张量分解所带来的近似误差可以忽略不计，并且通过基于SGD的微调可以轻松恢复。Figure 2显示了我们提出的总体框架的主要步骤。

4.1.   问题定式化（Problem Formulation）

如上所述，我们的框架的第一阶段是逐步将低张量秩特征强加到一个高精度的未压缩的DNN模型上。 在数学上，该目标可以表述为一个优化问题，最小化对象模型（object model）的损失函数，模型的每层都具有对TT秩的约束（全连接或者卷积）：

4.2. 用ADMM进行优化

 显然，由于rank（.）是不可微的，使用常规优化算法解决问题通常是困难的。为了克服这个挑战，我们首先将它重写为：

 （7）

因此，目标的形式（7）是一个经典的带约束的非凸优化问题，可以用ADMM正确地求解。 具体来说，我们可以首先引入一个辅助变量Z和一个s的指标函数g(.)

 然后问题（7）等价于以下形式:

为了保证不需要ℓ严格凸性或有限性等假设的收敛，代替拉格朗日量，给出了上述问题的尺度对偶形式的增广拉格朗日量：

此处省略一万个数学

4.3. 微调

 经过ADMM正则化训练后，我们首先将训练后的未压缩DNN模型分解为TT格式。这里对张量核进行分解，目标是TT-ranks为r*。由于ADMM优化过程已经对未压缩的模型施加了所需的低TT-rank结构，因此与其他现有的方法不同，不会带来显著的近似误差。（更多的细节将在第5.1节中进行分析）。然后，利用标准SGD对分解后的TT格式模型进行微调。请注意，在微调阶段，损失函数是，ADMM没有引入其他正则化术语。通常，这个微调阶段非常快，只需要几次迭代。这是因为在这一阶段的起点，分解的TT模型已经对原始的未压缩模型具有非常接近的精度。

5. 实验

对于图像分类任务，我们在MNIST、CIFAR-10、CIFAR-100和ImageNet数据集上评估了多个CNN模型。 对于视频分类任务，我们在UCF11和HMDB51数据集上评估了不同的LSTM模型。我们采用相同的秩选择方案，以满足目标压缩比的需要。为了简化选择过程，同一层（layer）中的大部分秩被设置为相等。

5.1. 收敛性和敏感性分析（Convergence and Sensitivity Analysis）

 如（10）所示，ρ是在ADMM正则化训练阶段引入的附加超参数。为了研究ρ对性能的影响以及促进超参数选择，我们研究了不同ρ设置下的ResNet-32模型的ADMM正则化训练在CIFAR10数据集上的收敛性和敏感性。

收敛性（Convergence）

Figure 3a 显示了ADMM正则化训练阶段的损失曲线。

 可以看出，不同曲线的ρ值非常不同，表现出非常相似的收敛速度。因此，这一现象表明，ρ对ADMM正则化训练的收敛性影响不大。

敏感性（Sensitivity）

 考虑到相似的收敛行为并不一定意味着不同的ρ会带来相似的精度，然后我们分析了ADMM正则化训练相对于ρ的性能敏感性。需要注意的是，在ADMM正则化训练后，W虽然采用非压缩格式，但应表现出较强的低TT秩特征，同时具有较高的精度。一旦W同时满足这两个标准，这意味着初始值从W分解的甚至在微调之前就已经具有很高的精度。

 

为了检验W所需的低TT秩行为（low TT-rank behavior），我们观察，它在ADMM正则化训练中衡量W和Z的相似度（Figure 3b）。由于（17）Z总是在低TT-rank约束下进行更新，图3b所示的曲线显示，除了ρ=0.001外，W在训练过程中确实很快表现出低TT-rank特征。这一现象意味着，为了确保权值张量被ADMM很好地正则化到目标TT秩，ρ不应该太小(小于0.001)。另一方面，Figure 3c显示了W在训练过程中的测试精度。在这里可以看到，较小的ρ往往会带来更好的性能。根据这些观察结果，ρ=0.005可以是让训练后的W满足上述两个标准的合适选择。

5.2. 图像分类

MINIST 

表1显示了LeNet-5模型在MNIST数据集上的实验结果。 我们将ADMM-TT-format模型于未压缩模型和最先进的TT/TR格式的模型进行了比较。可以看出，我们基于ADMM的压缩可以达到最高的压缩比和最好的精度。

 

CIFAR-10

表2将我们基于ADMM的TT格式ResNet-20和ResNet-32模型与CIFAR-10数据集上最先进的TT/TR格式的工作进行了比较。 对于ResNet-20，可以看到TT/TR格式模型的标准训练会导致严重的精度损失。即使是使用一些最先进的设计，如启发式秩选择(PSTRN-M/S)和强化学习(TR-RL)，性能的下降仍然是巨大的，特别是在高压缩比6.8x的情况下。另一方面，在同样的高压缩比下，我们基于ADMM的TT格式模型的精度仅下降了0.22%，即比最先进的PSTRN-M高出2.53%。此外，在中等压缩比4.5x下，我们的方法甚至可以优于未压缩模型，精度提高了0.22%。

同样，对于ResNet-32，使用TT或TR分解对压缩模型进行标准训练会导致巨大的性能下降。先进的PSTRN-S/M确实带来了性能改进，但测试精度仍然不令人满意。相反，我们的高压缩（5.8x）TT格式模型只有0.53%的精度损失，这意味着它的精度比相同压缩比的PSTRN-M高1.36%。更重要的是，当压缩比放宽到4.8x时，我们基于ADMM的TT格式模型达到了92.87%，甚至比未压缩模型高出0.38%。

CIFAR-100

表3显示了在CIFAR-100数据集上的实验结果。同样，我们基于ADMM的TT格式模型的性能优于最先进的工作。对于ResNet-20，在更高的压缩比下(我们的5.6x vs 4.7x)，我们的模型实现了1.3%的精度提高。在2.3x压缩比下，我们的模型达到了67.36%，甚至比未压缩模型高1.96%。对于ResNet-32，具有相同的5.2倍压缩比，我们的方法比最先进的PSTRN-M带来了0.4%的精度提高。在相同的2.4x压缩比下，我们的方法的精度比PSTRN-s高2.26%。我们的模型甚至优于未压缩模型，精度提高了2.21%。