[Paper Review] 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 Techniques4 citations

TL;DR

This paper proposes Trained Rank Pruning (TRP), a method that integrates low-rank approximation and nuclear norm regularization into the training process of deep neural networks, enabling efficient, accurate model compression. By alternately applying truncated SVD and stochastic sub-gradient descent on the nuclear norm, TRP maintains full network capacity during training, achieves near-zero accuracy loss after decomposition, and outperforms prior low-rank compression methods on CIFAR-10 and ImageNet with up to 2.23× inference speedup.

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.

Motivation & Objective

• To address the performance degradation caused by post-training low-rank approximation in compressed DNNs.
• To eliminate the need for extensive fine-tuning after low-rank decomposition by embedding the process into training.
• To improve model efficiency and accuracy by maintaining full network capacity while enforcing low-rank structure during optimization.
• To develop a training scheme that enables effective rank selection and stable optimization in deep low-rank networks.

Proposed method

• TRP alternates between standard backpropagation and low-rank approximation via truncated SVD (TSVD) every m iterations, preserving original weights during optimization.
• A stochastic sub-gradient descent-optimized nuclear regularization is applied to promote low-rank structure in the weight matrices during training.
• The method maintains the full capacity of the original network while gradually pushing weights toward a low-rank form through iterative TSVD and gradient updates.
• The approach is compatible with both channel-wise and spatial-wise low-rank decomposition, enabling flexible compression.
• The nuclear norm regularization is optimized via stochastic sub-gradient descent to enhance low-rank promotion without increasing computational overhead.
• The training process ensures that the theoretical bound for singular value decay is maintained, enabling stable convergence.

Experimental results

Research questions

• RQ1 Can low-rank approximation be effectively integrated into the training process to avoid performance degradation seen in post-compression fine-tuning?
• RQ2 How does nuclear norm regularization optimized via stochastic sub-gradient descent improve low-rank network training stability and accuracy?
• RQ3 To what extent can TRP maintain model accuracy while achieving high-speedup in inference on standard benchmarks?
• RQ4 Does TRP outperform existing low-rank compression methods in both channel-wise and spatial-wise decomposition settings?

Key findings

• TRP achieves 74.06% Top-1 and 92.07% Top-5 accuracy on ImageNet with 1.80× speedup, outperforming Luo et al. (2017) and He et al. (2017) at similar compression rates.
• With 2.23× acceleration, TRP maintains 1.4% Top-5 accuracy degradation, surpassing He et al. (2017) and Zhou et al. (2019) at the same speedup level.
• On ResNet-18, TRP with spatial-wise decomposition reduces inference time to 0.31ms on CPU, achieving 49.88ms per image, outperforming baseline (118.02ms).
• The nuclear norm regularization improves performance significantly on basic methods, but provides diminishing returns on TRP, indicating that TRP already induces low-rank structure effectively.
• Rank distribution analysis shows that singular values decay steadily over training, with energy concentrating on the first few singular values, converging when the smallest energy ratio reaches the threshold e=0.05.
• The theoretical bound for singular value decay is maintained throughout training, validating the stability of the TRP optimization process.