[论文解读] Network Pruning via Transformable Architecture Search

Xuanyi Dong, Yi Yang|arXiv (Cornell University)|May 23, 2019

Advanced Malware Detection Techniques参考文献 49被引用 140

一句话总结

论文提出 Transformable Architecture Search (TAS) 通过可微分的 NAS 学习最优宽度（通道数）和深度（层数）来裁剪网络，使用按通道的插值对多尺寸特征图进行聚合，并将知识从未裁剪的网络迁移到裁剪后的网络（KD）。在 CIFAR-10/100 和 ImageNet 上的实验显示相较于传统裁剪方法的改进。

ABSTRACT

Network pruning reduces the computation costs of an over-parameterized network without performance damage. Prevailing pruning algorithms pre-define the width and depth of the pruned networks, and then transfer parameters from the unpruned network to pruned networks. To break the structure limitation of the pruned networks, we propose to apply neural architecture search to search directly for a network with flexible channel and layer sizes. The number of the channels/layers is learned by minimizing the loss of the pruned networks. The feature map of the pruned network is an aggregation of K feature map fragments (generated by K networks of different sizes), which are sampled based on the probability distribution.The loss can be back-propagated not only to the network weights, but also to the parameterized distribution to explicitly tune the size of the channels/layers. Specifically, we apply channel-wise interpolation to keep the feature map with different channel sizes aligned in the aggregation procedure. The maximum probability for the size in each distribution serves as the width and depth of the pruned network, whose parameters are learned by knowledge transfer, e.g., knowledge distillation, from the original networks. Experiments on CIFAR-10, CIFAR-100 and ImageNet demonstrate the effectiveness of our new perspective of network pruning compared to traditional network pruning algorithms. Various searching and knowledge transfer approaches are conducted to show the effectiveness of the two components. Code is at: https://github.com/D-X-Y/NAS-Projects.

研究动机与目标

• Motivate network pruning to reduce computation without sacrificing accuracy in over-parameterized CNNs.
• Shift pruning from hand-crafted fixed structures to learning network size using NAS.
• Honor compute budgets by optimizing width and depth under a cost constraint.
• Leverage knowledge transfer from the unpruned network to the pruned architecture.
• Demonstrate generality across datasets (CIFAR-10/100, ImageNet) and architectures.

提出的方法

• Introduce TAS to search for per-layer channel counts and per-stage depths via differentiable architecture parameters.
• Attach learnable distributions to candidate channel numbers and layer counts; optimize via Gumbel-Softmax to enable backpropagation.
• Aggregate multiple size fragments of feature maps using channel-wise interpolation (CWI) and a weighted sum guided by sampled sizes.
• Compute final outputs as an aggregation over depths; back-propagate to both width (alpha) and depth (beta) parameters.
• Use a computation-cost term in the validation loss to encourage meeting a target FLOPs, with a piecewise cost function.
• Transfer knowledge from the unpruned network to the pruned architecture using Knowledge Distillation (KD) to improve pruning performance.

实验结果

研究问题

• RQ1 Can NAS be used to directly optimize network size (width and depth) for pruning, rather than topology alone?
• RQ2 Does differentiable sampling of width/depth with a cost-aware objective yield superior pruned architectures?
• RQ3 Is knowledge transfer from the unpruned model beneficial for the pruned network’s performance?
• RQ4 How do sampling strategies and feature map alignment (CWI) affect NAS effectiveness?
• RQ5 How do TAS-derived architectures compare to traditional pruning and other NAS baselines on CIFAR and ImageNet?

主要发现

• TAS with KD consistently improves pruning performance over traditional methods on CIFAR-10/100 and ImageNet.
• Searching both width and depth yields better accuracy at similar FLOPs than searching either alone.
• Knowledge transfer (KD) from the unpruned network enhances pruned network accuracy across experiments.
• Using channel-wise interpolation and differentiable architecture parameters enables effective alignment and aggregation of multi-size feature maps.
• Compared to state-of-the-art pruning methods, TAS achieves higher accuracy at comparable or lower FLOPs on several ResNet variants and datasets.