[论文解读] TRP: Trained Rank Pruning for Efficient Deep Neural Networks

Yuhui Xu, Yuxi Li+7 more|arXiv (Cornell University)|Apr 30, 2020

Sparse and Compressive Sensing Techniques被引用 4

一句话总结

本文提出训练秩剪枝（TRP），一种将低秩近似与核范数正则化整合到深度神经网络训练过程中的方法，实现了高效且高精度的模型压缩。通过在截断奇异值分解（TSVD）与核范数的随机次梯度下降之间交替进行，TRP 在训练过程中保持了完整的网络容量，在分解后实现近乎零的精度损失，并在 CIFAR-10 和 ImageNet 上实现了高达 2.23 倍的推理加速，优于以往的低秩压缩方法。

ABSTRACT

To enable DNNs on edge devices like mobile phones, low-rank approximation has been widely adopted because of its solid theoretical rationale and efficient implementations. Several previous works attempted to directly approximate a pretrained model by low-rank decomposition; however, small approximation errors in parameters can ripple over a large prediction loss. As a result, performance usually drops significantly and a sophisticated effort on fine-tuning is required to recover accuracy. Apparently, it is not optimal to separate low-rank approximation from training. Unlike previous works, this paper integrates low rank approximation and regularization into the training process. We propose Trained Rank Pruning (TRP), which alternates between low rank approximation and training. TRP maintains the capacity of the original network while imposing low-rank constraints during training. A nuclear regularization optimized by stochastic sub-gradient descent is utilized to further promote low rank in TRP. The TRP trained network inherently has a low-rank structure, and is approximated with negligible performance loss, thus eliminating the fine-tuning process after low rank decomposition. The proposed method is comprehensively evaluated on CIFAR-10 and ImageNet, outperforming previous compression methods using low rank approximation.

研究动机与目标

• 解决压缩 DNN 中因后训练低秩近似导致的性能下降问题。
• 通过将该过程嵌入训练阶段，消除低秩分解后对大量微调的依赖。
• 通过在优化过程中保持完整网络容量并强制实施低秩结构，提升模型效率与精度。
• 开发一种训练方案，以实现深层低秩网络中有效的秩选择与稳定优化。

提出的方法

• TRP 每隔 m 个迭代周期交替执行标准反向传播与通过截断 SVD（TSVD）进行的低秩近似，同时在优化过程中保留原始权重。
• 在训练过程中应用经随机次梯度下降优化的核范数正则化，以促进权重矩阵的低秩结构。
• 该方法在逐步通过迭代 TSVD 与梯度更新将权重推向低秩形式的同时，保持了原始网络的完整容量。
• 该方法兼容通道级与空间级的低秩分解，支持灵活压缩。
• 通过随机次梯度下降优化核范数正则化，以在不增加计算开销的前提下增强低秩促进效果。
• 训练过程确保了奇异值衰减的理论边界得以维持，从而实现稳定收敛。

实验结果

研究问题

• RQ1 低秩近似能否有效集成到训练过程中，以避免后压缩微调中出现的性能下降？
• RQ2 通过随机次梯度下降优化的核范数正则化如何提升低秩网络训练的稳定性和精度？
• RQ3 在标准基准上，TRP 在实现高倍推理加速的同时，能在多大程度上保持模型精度？
• RQ4 在通道级与空间级分解设置下，TRP 是否优于现有的低秩压缩方法？

主要发现

• 在 ImageNet 上，TRP 实现了 74.06% 的 Top-1 准确率与 92.07% 的 Top-5 准确率，推理速度提升 1.80 倍，优于 Luo 等人（2017）与 He 等人（2017）在相似压缩率下的表现。
• 在 2.23 倍加速下，TRP 仅造成 1.4% 的 Top-5 准确率下降，优于 He 等人（2017）与 Zhou 等人（2019）在相同加速水平下的表现。
• 在 ResNet-18 上，采用空间级分解的 TRP 将 CPU 上的推理时间降低至 0.31ms，每张图像处理时间仅为 49.88ms，优于基线模型（118.02ms）。
• 核范数正则化在基础方法上显著提升了性能，但在 TRP 上收益递减，表明 TRP 已能有效诱导出低秩结构。
• 秩分布分析显示，奇异值在训练过程中持续衰减，能量集中于前几个奇异值，当最小能量比达到阈值 e=0.05 时趋于收敛。
• 在整个训练过程中，奇异值衰减的理论边界均被维持，验证了 TRP 优化过程的稳定性。