Source-linked AI summary
N2N Learning: Network to Network Compression via Policy Gradient Reinforcement Learning
Anubhav Ashok, Nicholas Rhinehart, Fares Beainy, Kris M. Kitani
TL;DR
Large neural networks can be difficult to deploy, while manually selecting compressed student architectures is a challenging search problem. The paper uses reinforcement learning with two-stage layer removal and shrinkage policies, rewarding accuracy and compression, and reports effectiveness across tasks and architectures with transfer to larger teachers. It concludes that the procedure can also incorporate practical constraints such as power or inference time.
Problem
Large networks can exceed deployment resources, while manually designing compressed student architectures makes it difficult to identify an optimal reduced network.
Method
The method uses reinforcement learning to select a compressed student architecture through sequential layer removal and layer shrinkage, with rewards based on accuracy and compression.
Results
The method performs well across several datasets and architectures and shows transfer learning that makes compression more efficient on larger networks.
Takeaways & Limitations
The procedure can support practical constraints such as power or inference time and can be made more efficient through transfer learning.
Takeaways & Limitations
Straightforward reinforcement-learning formulations can become intractable because their search space grows exponentially with action-space size and sequence length.
Abstract
from arXiv · showhide
While bigger and deeper neural network architectures continue to advance the state-of-the-art for many computer vision tasks, real-world adoption of these networks is impeded by hardware and speed constraints. Conventional model compression methods attempt to address this problem by modifying the architecture manually or using pre-defined heuristics. Since the space of all reduced architectures is very large, modifying the architecture of a deep neural network in this way is a difficult task. In this paper, we tackle this issue by introducing a principled method for learning reduced network architectures in a data-driven way using reinforcement learning. Our approach takes a larger `teacher' network as input and outputs a compressed `student' network derived from the `teacher' network. In the first stage of our method, a recurrent policy network aggressively removes layers from the large `teacher' model. In the second stage, another recurrent policy network carefully reduces the size of each remaining layer. The resulting network is then evaluated to obtain a reward -- a score based on the accuracy and compression of the network. Our approach uses this reward signal with policy gradients to train the policies to find a locally optimal student network. Our experiments show that we can achieve compression rates of more than 10x for models such as ResNet-34 while maintaining similar performance to the input `teacher' network. We also present a valuable transfer learning result which shows that policies which are pre-trained on smaller `teacher' networks can be used to rapidly speed up training on larger `teacher' networks.
1 INTRODUCTION
The paper frames student-architecture design as a large, difficult search problem and proposes reinforcement learning to compress teacher networks systematically. Its two-stage policies improve scalability, while the learned policies are evaluated across architectures and tasks, including transfer to larger networks.
- Motivation: Hand-designed student architectures make it difficult to identify an optimal compressed network.Deep networks can exceed smaller devices’ power, memory, and computational limits, while architecture design often relies on tedious trial and error.
- Approach: Reinforcement learning models teacher-to-student compression as a Markov Decision Process over network architectures and structural actions.States represent architectures, while actions can remove filters or reduce layer sizes.
- Approach: A two-stage action mechanism first removes layers and then shrinks remaining layers to make architecture exploration more efficient.The policies select macro-scale layer-removal actions followed by micro-scale layer-shrinkage actions.
- Approach: The reward combines student accuracy and compression rate, with a proposed formulation incorporating hardware-based computational budgets as linear constraints.This reward provides supervision for searching reduced architectures.
- Evaluation: The method is evaluated across multiple architectures and visual-learning tasks, with policies also transferring from ResNet-18 to ResNet-34 to accelerate reinforcement-learning training.The experiments cover MNIST, SVHN, CIFAR-10, CIFAR-100, and Caltech-256.
2 RELATED WORK
Prior work compresses manually designed networks through pruning and knowledge distillation, while architecture-search methods automate network construction. N2N differs by learning the compressed student architecture itself through reinforcement learning.
- Pruning: Pruning methods remove redundant weights while preserving weights judged important, whereas N2N searches over both model weights and architecture.This places N2N in a larger search space than weight-focused pruning methods.
- Knowledge Distillation: Knowledge distillation trains a smaller student network to mimic a teacher, but prior approaches generally use hand-designed student architectures.N2N trains a policy to learn the student architecture instead of designing it manually.
- Architecture Search: Architecture-search research explores neural-network design spaces, providing an automation direction distinct from the compression task addressed here.The related work distinguishes constructing high-performance networks from compressing an existing teacher network.
3 APPROACH
The approach models teacher-to-student architecture compression as a Markov decision process and uses policy-gradient reinforcement learning to search reduced networks. It combines layer removal, layer shrinkage, knowledge-distillation evaluation, and rewards balancing compression, accuracy, and resource constraints.
- Markov Decision Process: The method treats each reduced network architecture as an MDP state and uses actions that deterministically transform one architecture into another.Actions include layer removal and layer-parameter reduction; the policy seeks to maximize expected total reward.
- Architecture Evaluation: Each generated architecture is trained with knowledge distillation and receives a final-state reward, while intermediate incomplete architectures receive zero reward.The optimization repeatedly samples architectures and updates both policies from their achieved rewards.
- Two-Stage Compression: A two-stage procedure first selects which layers to remove and then shrinks parameters in the remaining layers, reducing the search burden of incremental actions.The first policy operates over layer-level keep/remove decisions, while the second produces shrinkage actions for layer configuration variables.
- Reward Function: The reward multiplies compression and accuracy terms, with compression computed relative to teacher parameters and accuracy normalized by teacher validation accuracy.The compression term uses Rc = C(2 −C), while the accuracy term is Ra = A/Ateacher; degenerate architectures receive −1.
- Constraints as Rewards: Hardware and resource budgets are incorporated through constraints on architecture variables, with penalties discouraging policy outputs that violate the specified limits.The constraints can represent resources such as memory, time, power, or accuracy, but a non-smooth penalty can reduce exploration and worsen local convergence.
- Optimization: Policy gradients use sampled rollouts, where each action probability is weighted by the rollout reward to update the policy parameters.The rollout count is m, the trajectory length is T, and Pθ(at|ht) denotes the action probability given the policy hidden state.
4 EXPERIMENTS
Experiments across datasets and teacher architectures show that the method finds highly compressed, high-performing student networks, often matching or exceeding teachers. It also outperforms pruning and hand-designed distillation baselines and benefits from pretrained policies.
- Compression experiments: Reward, accuracy, and compression improved over iterations in the MNIST compression experiments, with policies moving toward a better size-accuracy trade-off.The figures plot these quantities against iteration for both policy stages.
- Compression experiments: 10x compression on CIFAR-10 still yielded student networks performing almost as well as or better than their teacher networks.The comparison used VGG-19, ResNet-18, and ResNet-34 teachers.
- Compression experiments: On CIFAR-100, compressed networks retained reasonably high accuracy despite fewer parameters, and a compressed ResNet-34 student outperformed ResNet-18 with fewer parameters.The experiments also found that ReLU layers within residual blocks were removed from many final compressed models.
- Baseline comparisons: The method outperformed pruning on both evaluated datasets despite pruning achieving good compression rates.The comparison used compression rate defined by parameter count rather than number of bits.
- Baseline comparisons: The method performed better than hand-designed models on both datasets despite containing fewer parameters.The comparison included VGG-small, FitNet-4, and SqueezeNet; the authors also identify model selection as important in knowledge distillation.
- Transfer learning: Pretrained policies achieved higher average reward, while a slight accuracy drop reflected the tradeoff between compression and accuracy.Reusing policies from smaller teacher models provides evidence of faster policy search on larger models.
5 CONCLUSION
The method compresses large neural networks through two-stage layer removal and shrinkage, using accuracy and compression signals to search architectures across datasets and networks. It also supports transfer learning and practical constraints such as power or inference time.
- Two-stage layer removal and layer shrinkage compress large neural networks through learned architecture search.The procedure uses accuracy and compression as supervision.
- The method performs well across a variety of datasets and architectures.
- Transfer learning generalizes the method and makes the compression procedure more efficient.
- Power and inference-time constraints can be incorporated into the compression process.
6 ACTOR-CRITIC
The section compares Actor-Critic and vanilla REINFORCE for policy search. Actor-Critic improves stability in some settings but can explore less effectively or perform worse than REINFORCE.
- Actor-Critic replaces the reward baseline with a learned value-function estimate at each time step.
- Actor-Critic uses an additional fully connected layer that maps the LSTM hidden state to a scalar value.
- On MNIST, Actor-Critic provides a slight stability improvement, while both methods converge at a similar rate.
- On CIFAR-10, Actor-Critic is more stable but performs worse than vanilla REINFORCE.
7 LEARNING RATE AND BATCH SIZE
Learning rate and batch size were selected through grid searches on MNIST. The chosen settings balanced convergence speed, stability, and additional computation.
- Grid searches evaluated learning rates and batch sizes by measuring convergence rates on MNIST.
- lr=0.003 was selected because lr=0.03 did not converge and lr=0.0003 converged too slowly.
- Average reward was compared across learning rates and batch sizes over 3 runs on MNIST.
- A batch size of 5 worked best because batch size 1 was unstable and batch size 10 added computation without substantial improvement.
8 TRANSFER LEARNING EXPERIMENTS
Pretrained policies begin transfer-learning experiments with higher reward than policies trained from scratch, supporting faster policy search on larger networks.
- Pretrained policies start with a high reward, unlike policies trained from scratch.
9 ADDITIONAL EXPERIMENTS
This section introduces additional compression experiments, including transfer learning and ResNet-18 evaluations on SVHN and Caltech.
- Additional experiments evaluate transfer learning for the compression policies.
- ResNet-18 experiments on SVHN: ResNet-18 experiments on SVHN report separate results for Stage 1 and Stage 2.
- ResNet-18 experiments on Caltech: ResNet-18 experiments on Caltech report separate results for Stage 1 and Stage 2.
10 IMPLEMENTATION DETAILS
The implementation details describe the experimental setup, reward design, degenerate architectures, and computational limitations of evaluating candidate students.
- Implementation details: The removal and shrinkage policies use two hidden layers, Adam optimization, and learning rates of 0.003 and 0.1, respectively.Each policy uses 30 or 50 hidden units and is trained for at least 100 epochs with batches of five rollouts.
- Reward design: The reward penalizes high compression with low accuracy more strongly than low compression with high accuracy while increasing with both quantities.These criteria are introduced to prevent compression from dominating model performance.
- Reward design: The naive reward R = A ∗C assigns 0.25 to both contrasting accuracy–compression cases, whereas the paper’s reward assigns 0.25 and 0.4375.The paper reports that the non-linear reward outperformed the naive reward empirically.
- Reward design: The reward-manifold comparison shows that the proposed reward gives lower scores to low-accuracy, high-compression models while remaining monotonically increasing.
- Degenerate architectures: Degenerate outputs include empty architectures, impractically large fully connected layers, and specialized-architecture dependencies.Residual networks require matching feature-map dimensionality at the start and end of each residual block.
- Limitations and future directions: Evaluating each student requires training it for a few epochs to obtain a reward, which can be computationally expensive depending on the dataset.The authors suggest random-weight initialization, hypernetworks, or informative data subsets as possible alternatives.