Source-linked AI summary
Network Pruning via Transformable Architecture Search
Xuanyi Dong, Yi Yang
TL;DR
Existing pruning methods restrict pruned networks with predefined widths and depths, limiting structural choices for reducing computation. The paper proposes Transformable Architecture Search, which differentiably learns flexible network sizes and transfers knowledge to the resulting networks. Across CIFAR-10, CIFAR-100, and ImageNet, searched architectures with knowledge distillation outperform previous state-of-the-art pruning methods.
Problem
Existing pruning methods predefine pruned-network widths and depths, constraining the structures available for reducing over-parameterized networks’ computation.
Method
Transformable Architecture Search differentiably learns channel and layer sizes from pruned-network loss by aggregating sampled feature-map fragments, then transfers knowledge from unpruned networks.
Results
Searched architectures with parameters transferred by knowledge distillation outperform previous state-of-the-art pruning methods on CIFAR-10, CIFAR-100, and ImageNet.
Takeaways & Limitations
The experiments support a pruning pipeline that combines learned network size with knowledge transfer, with improvements also reported across different architectures and transfer approaches.
Takeaways & Limitations
TAS is presented as preliminary, and the pruning pipeline could be improved through more effective searching and knowledge-transfer methods.
Abstract
from arXiv · showhide
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.
1 Introduction
The paper reframes pruning as learning network size rather than applying hand-crafted structural rules. TAS searches flexible widths and depths differentiably, then transfers knowledge to the searched network.
- Over-parameterized CNNs are difficult to deploy on resource-constrained devices, motivating computation reduction through pruning.
- Traditional pruning removes filters and fine-tunes the resulting network, but its accuracy is bounded by hand-crafted widths, depths, or sparsity rules.
- TAS searches network size rather than topology, optimizing channel and layer counts with computation-cost regularization such as FLOPs.
- TAS learns candidate channel and layer probabilities by back-propagating pruned-network loss through aggregated feature-map fragments aligned by channel-wise interpolation.
- The maximum probability in each learned distribution determines the pruned network’s width and depth, whose parameters are learned through knowledge transfer.
- Knowledge-distilled searched architectures outperform previous state-of-the-art pruning methods on CIFAR-10, CIFAR-100, and ImageNet.
2 Related Studies
Related work spans compression, structured pruning, architecture search, network transformation, and knowledge transfer. TAS differs by explicitly and differentiably optimizing flexible network width and depth.
- Network pruning compresses and accelerates CNNs for hardware with limited storage and computation, using techniques including low-rank decomposition, weight pruning, channel pruning, dynamic computation, and quantization.
- Structured pruning removes convolutional filters or whole layers, producing networks that are easier to deploy than unstructured sparse networks.
- Unlike prior pipelines that implicitly or manually set pruning ratios, the proposed approach explicitly optimizes the number of channels and layers for high performance and low cost.
- Most NAS methods search topology, whereas TAS searches network depth and width with a differentiable method that supports flexible channel sizes.
- Earlier network-transformation methods manually or heuristically grow networks, while TAS can differentiably enlarge or shrink CNNs.
- Knowledge transfer from unpruned networks is established in pruning, and the paper applies knowledge distillation to obtain robust performance for searched architectures.
3 Methodology
The method searches network width and depth differentiably, then derives a compact architecture and trains it through knowledge transfer. TAS aggregates sampled feature-map fragments with channel-wise interpolation and optimizes classification and computation-cost objectives.
- The pruning pipeline trains an unpruned network, searches a smaller network’s width and depth with TAS, then transfers knowledge using KD.The procedure alternates architecture search with subsequent knowledge distillation on the searched architecture.
- TAS assigns learnable probabilities to candidate channel counts and layer depths, allowing architecture parameters to receive gradients from pruned-network losses.Width uses α and depth uses β; the final architecture selects the maximum-probability candidates.
- Gumbel-Softmax makes candidate sampling differentiable, while channel-wise interpolation aligns feature maps with different channel sizes before weighted aggregation.A small sampled subset can be aggregated to reduce memory costs and training time.
- The searched architecture is selected through a bilevel objective that minimizes validation loss after training weights on the training loss.TAS searches candidates sharing the same topology structure but with smaller widths and depths.
- The validation objective combines classification loss with a computation-cost penalty, using expected and actual costs to guide the architecture toward a target budget.The cost metric is FLOP in the described implementation and can be replaced by latency.
- After search, knowledge distillation trains the compact network to match soft targets from the pretrained unpruned network while also predicting the true labels.The distillation objective uses temperature-scaled softmax cross-entropy together with ordinary target supervision.
4 Experimental Analysis
Experiments evaluate TAS across datasets, search strategies, knowledge-transfer methods, architectures, and computational settings. TAS generally improves pruning results, while joint width-depth search and knowledge transfer are especially effective, with search cost and large search spaces remaining practical boundaries.
- Experimental setup: Experiments cover CIFAR-10, CIFAR-100, and ImageNet, with TAS evaluated through hyperparameters, sampling, transfer methods, and comparisons against pruning algorithms.The datasets include 50K/10K images for CIFAR-10, 50K/10K for CIFAR-100, and 1.28M/50K for ImageNet.
- Search strategies: TAS with CWI and Gumbel-softmax becomes more confident in selecting suitable widths and can constrain FLOPs to the target range.Without FLOP constraints, TAS finds maximum width and depth; strategies without CWI fail to optimize architecture parameters.
- Architecture comparison: TAS outperforms the NAS random-search baseline and finds better structures using different knowledge-transfer methods.The comparison is reported in Table 1 for ResNet-32 on CIFAR-100 with about 40% FLOPs pruned.
- Knowledge transfer: Knowledge transfer consistently improves pruned-network accuracy, while KD improves CIFAR-100 accuracy by more than 2%.The comparison includes initialization and KD; KD is reported as robust across the evaluated architectures.
- Search dimensions: Jointly searching width and depth achieves better accuracy at similar FLOPs than searching either depth or width alone.Searching depth alone performs worse than searching width alone in the reported comparison.
- Efficiency and limitations: Searching ResNet-32 takes about 3.8 hours on one V100 GPU, whereas larger search spaces such as ResNet-164 are challenging to explore.The authors identify 8163 × 183 candidate structures for ResNet-164 and note over-fitting on CIFAR-10.
- State-of-the-art comparisons: TAS achieves 72.25% accuracy when pruning ResNet-56 on CIFAR-100 and 76.20% accuracy when pruning 43.5% FLOPs of ResNet-50 on ImageNet.The CIFAR-100 result exceeds FPGM’s 69.66%, while the ImageNet result exceeds FPGM by 0.7.
- Efficiency and limitations: TAS is presented as a preliminary pruning pipeline that could be improved through more effective search and knowledge-transfer methods.The authors identify both components as directions for future work.
5 Conclusion
The paper presents network pruning as a two-component paradigm: search for suitable depth and width with TAS, then optimize the searched network by transferring knowledge from the unpruned network.
- TAS applies NAS to search for a network’s depth and width rather than its topology.
- Differentiable TAS is designed to find suitable network depth and width efficiently and effectively.
- The searched network is optimized by transferring knowledge from the unpruned network.
- A simple KD algorithm performs knowledge transfer, while other transfer approaches test the effectiveness of this component.
- The results suggest that jointly improving architecture search and knowledge transfer may produce further advances in network pruning.