Compact Model Training by Low-Rank Projection with Energy Transfer

TL;DR

LRPET presents a novel training-from-scratch framework for low-rank neural networks by alternating SGD updates with low-rank projections and augmenting the process with energy transfer to preserve gradient energy and BN rectification to neutralize BN’s impact on optimal low-rank solutions. It introduces a practical rank-setting strategy (unique ratio or Bayesian optimization) and a ready-to-use inference scheme that factorizes weights into two smaller matrices for efficient computation. Empirical results on CIFAR-10 and ImageNet across a range of architectures show competitive or superior accuracy at substantial reductions in FLOPs and parameters, with strong generalization to object detection and semantic segmentation. The approach remains compatible with other compression techniques, and ablations highlight the critical roles of energy transfer and BN rectification in achieving robust performance.

Abstract

Low-rankness plays an important role in traditional machine learning, but is not so popular in deep learning. Most previous low-rank network compression methods compress networks by approximating pre-trained models and re-training. However, the optimal solution in the Euclidean space may be quite different from the one with low-rank constraint. A well-pre-trained model is not a good initialization for the model with low-rank constraints. Thus, the performance of a low-rank compressed network degrades significantly. Compared with other network compression methods such as pruning, low-rank methods attract less attention in recent years. In this paper, we devise a new training method, low-rank projection with energy transfer (LRPET), that trains low-rank compressed networks from scratch and achieves competitive performance. We propose to alternately perform stochastic gradient descent training and projection of each weight matrix onto the corresponding low-rank manifold. Compared to re-training on the compact model, this enables full utilization of model capacity since solution space is relaxed back to Euclidean space after projection. The matrix energy (the sum of squares of singular values) reduction caused by projection is compensated by energy transfer. We uniformly transfer the energy of the pruned singular values to the remaining ones. We theoretically show that energy transfer eases the trend of gradient vanishing caused by projection. In modern networks, a batch normalization (BN) layer can be merged into the previous convolution layer for inference, thereby influencing the optimal low-rank approximation of the previous layer. We propose BN rectification to cut off its effect on the optimal low-rank approximation, which further improves the performance.

Figures (4)

• Figure 1: An illustration of the procedure of energy transfer. For each convolution layer, we reshape its weight data as matrix Mi and decompose it by SVD method, then we clip the singular value matrix diag(s) with energy transfer and finally replace the convolution layer weight data. The dotted arrow denotes extracting/recovering the singular values from/to the diagonal matrix. The solid arrow denotes transferring the energy.
• Figure 2: Rank distribution comparison between manual fixed layer rank ratio setting and searched layer rank setting for ResNet-56 on CIFAR-10.
• Figure 3: Learning Curves for LRP and its variants. The learning curves of training loss, test loss, and test accuracy for ResNet-56 and VGG-16 on CIFAR-10 are illustrated.
• Figure 4: Test accuracy curves of different fixed $\alpha$ and our adaptive amplification coefficient $\alpha$ setting for ResNet-56 and VGG-16 on CIFAR-10 .

Theorems & Definitions (1)

• Theorem 1