Non-Uniform Adversarially Robust Pruning
Abstract Neural networks often are highly redundant and can thus be effectively compressed to a fraction of their initial size using model pruning techniques without harming the overall prediction accuracy.
Additionally, pruned networks need to maintain robustness against attacks such as adversarial examples.
Recent research on combining all these objectives has shown significant advances using uniform compression strategies, that is, parameters are compressed equally according to a preset compression ratio.
In this paper, we show that employing non-uniform compression strategies allows to improve clean data accuracy as well as adversarial robustness under high overall compression—in particular using channel pruning.
We leverage reinforcement learning for finding an optimal trade-off and demonstrate that the resulting compression strategy can be used as a plug-in replacement for uniform compression ratios of existing state-of-the-art approaches. Our code is available at https://intellisec.de/research/heracles
1 Introduction
Deploying deep neural networks on resource-constrained hardware is often hindered by the sheer size of the network.
Neural network pruning effectively removes redundancy at different structural granularity to reduce a model’s size.
In safety-critical environments, these networks additionally need to be robust against attacks, such as adversarial examples (Szegedy et al., 2014)
With adversarial training (Madry et al., 2018; Shafahi et al., 2019; Wong et al., 2020) it is possible to significantly improve robustness by introducing adversarial examples into the training process.
However, recent research (Zhang et al., 2019) suggests that large networks have higher adversarial robustness.
Consequently, it is inherently difficult to strike a balance between the compactness and robustness against attacks when pruning neural networks.
The typical network pruning procedure consists of three stages (Liu et al., 2019):
First, an over-parameterized model is trained.
Second, this pre-trained model is pruned based on a specific criterion and strategy.
Finally, the pruned network is fine-tuned to recover the potentially lost performance.
The most critical step in the procedure is the second one that defines the pruning objective and any additional objectives next to network compression itself.
Han et al.
(2015) propose to prune network connections following the order of weight magnitude (OWM), which later on has also been shown effective for robustness-aware pruning by Sehwag et al.
Ye et al.
and
Gui et al.
(2019) inherit this criterion and define network pruning as an optimization problem that can be solved by the alternating direction method of multipliers (ADMM), initially proposed by
Boyd et al.
Similarly, Sehwag et al.
(2020) formulate the pruning criterion as an importance score-based optimization problem that, however, anchors adversarial robustness deeply in the pruning process itself.
While both OWM (Ye et al., 2019) and optimization-based criteria (Sehwag et al., 2020) yield good results for robust-aware pruning, they require the specification of the compression ratio as an hyper-parameter that is then used uniformly across all layers.
Madaan et al.
propose ANP-VS to combine adversarial training with pruning and thus they merge the previously mentioned steps one and two.
As such, the method pursues a different goal for which compression does not need to be adjustable.
However, ANP-VS learns an implicit non-uniform compression that yields promising results.
In this paper, we follow this intuition and investigate the possibility of improving both compression and adversarial robustness of existing state-of-the-art approaches using non-uniform compression strategies .
The necessity of non-uniform compression is most evident for channels for which Table 1 provides a first glimpse of the improvements made by our method
Heracles.
We prune a network’s layers based on the order of weight magnitude (OWM), but determine the compression rate per layer.
Inspired by He et al.
(2018), we leverage deep reinforcement learning (Deep-RL) to automatically find this global pruning strategy to yield an optimal trade-off between accuracy and adversarial robustness of the pruned network.
The determined compression strategy is then used with approaches for pruning a pre-trained model which allows for increasing accuracy on benign as well as adversarial inputs.
[Table]
Contributions and impact.
We show that a non-uniform , global compression strategy is beneficial for effective network pruning when considering adversarial robustness.
The compression strategy learned by Heracles can be applied to state-of-the-art pruning techniques as a plug-in replacement for manually specified compression rates to improve original (benign) and adversarial accuracy whilst yielding the same overall compression.
In extensive experiments with the CIFAR-10, SVHN, and ImageNet datasets, we show to surpass the performance of Robust-ADMM (Ye et al., 2019) and
Hydra (Sehwag et al., 2020) in pruning channels that originally use uniform compression strategies.
As shown in Table 1, for channel pruning on VGG16, we yield up to 10 .
80 percentage points higher begin accuracy and 8 .
73 percentage points higher accuracy under adversarial inputs using RobustADMM.
For Hydra, we even successfully escape from a completely damaged model and achieve a remarkable performance improvement.
In summary, we are able to significantly improve channel pruning over related work and maintain at least on-par performance in weight pruning.
In practice, the additionally gained model robustness helps increase the security and safety of applications on hardware-constraint platforms, for instance, in autonomous driving or edge AI.
Moreover, channel pruning is particularly suitable for hardware deployment as it straightforwardly reduces the dimensionality of the necessary computations and thus also speeds up inference.
2 Background
We begin by briefly recapping concepts that are central to our approach, such as basic background on network pruning, adversarial training, and reinforcement learning.
2.1 Network Pruning
Network pruning enables to compress over-parameterized neural networks by removing structural redundancy (Han et al., 2015, 2016).
For this, usually a binary mask 𝑀 with elements in { 0 , 1 } is introduced to cancel out redundant network connections at weight level or channel level.
We represent this masking operation by the Hadamard product (cid:12) that transforms the model (its parameters) at the 𝑙 th layer of the network, 𝜃 ( 𝑙 ) , to a sparse (pruned) representation (cid:101) 𝜃 ( 𝑙 ) : (cid:101) [Equation] =
[Equation] (cid:12) 𝜃 ( 𝑙 ) .
Note that determining the importance of connections, and thus populating the binary mask 𝑀 , depends on the criterion used in the pruning stage.
The order of weight magnitude (OWM) has been shown to outperform other criteria such as Variational Dropout (Molchanov et al., 2017), Soft
Weight-Sharing (Karen Ullrich, 2017), or Filter Standard Deviation (Sun et al., 2019).
Thus, it is seen as the gold standard in network pruning (Liu et al., 2019).
Consequently, for Heracles, we pick up the OWM criterion for pruning as well but learn a global strategy.
Similar to Sehwag et al.
(2020), we use scored masks to binarize the pruning mask and initially assign scores to each element of the pruning masks 𝑀 based on scaled-absolute-initialization:
[Equation]
where | 𝜃 ( 𝑙 ) | takes the absolute values of model’s parameters of layer 𝑙 .
Note that for channel pruning, score masks are commonly initialized by the sum of absolute weights along each channel to comply with the OWM criterion.
2.2 Adversarial Training
To date, adversarial training (Madry et al., 2018) in its different manifestations (Zhang et al., 2019;
Wong et al., 2020) is the most efficient defense against adversarial examples (Szegedy et al., 2014).
It generates attacks and incorporates them in the training process, solving a min-max optimization problem, which is formally expressed as:
[Equation]
Input pairs of a data sample 𝑥 ∈ R 𝑑 and its label 𝑦 ∈
[ 𝑘 ] are drawn from the training data distribution D 𝑡 , where 𝑘 represents the number of classes.
As the normal training procedure, the outer minimization reduces the loss function L adv , for instance, the cross-entropy loss.
The inner maximization is formulated to increase the maximally allowed (adversarial) perturbation 𝛿 for each input data sample 𝑥 , and is solved by projected gradient descent (PGD) (Madry et al., 2018).
Building on top of this concept, several approaches have been proposed that improve upon the performance of PGD-based adversarial training (Zhang et al., 2019; Shafahi et al., 2019; Wong et al., 2020).
However, Guo et al.
(2018) and Ye et al.
(2019) show that increasing adversarial robustness is accompanied with stronger parameter distribution, which commonly hinders network pruning.
By striving for a globally optimal compression strategy with varying compression ratios per layer, we show that adversarial robustness and large compression rates are not mutually exclusive.
2.3 Reinforcement Learning
For reinforcement learning (RL), an agent strives for an action strategy to maximize the reward R over multiple episodes 𝑖 that provides feedback about the effectivity of certain actions in a specific environment (Sutton and Barto, 2018):
[Equation]
Here, 𝑠 refers to the agent’s state and 𝛾 represents the discount rate in each episode 𝑖 .
Policy 𝜋 aims to maximize the cumulative reward by optimizing the mapping from states to actions taken by the
RL agent.
To tailor this process to a particular application, such as network pruning, we have to define a state space representing the environment as well as an action space that specifies allowed actions.
The agent then outputs a so-called action space vector to influence its “location” in the environment.
In our case, this environment is the model 𝜃 we operate on.
For instance, Huang et al.
(2018) deploy a RL agent for a filter pruning, where the state space is composed of the number of input feature maps and the shape in each filter.
The agent returns a discrete action vector that scores the importance of each filter. [Figure]
3 Non-Uniform Adversarially Robust Pruning
Heracles searches for the globally optimal pruning strategy that increases compression of an adversarially trained network with minimal degradation on both benign accuracy and adversarial robustness.
A schematic depiction is provided in Fig. In contrast to focusing on direct connections in a network as implemented in related work (Sehwag et al., 2020), we consider the relations of all layers to each other (pre and post relations) and observe that these far-reaching contexts effect the robustness after pruning and fine-tuning.
Finding an optimal compression strategy [ 𝑎 ( 1 ) , . .
𝑎 ( 𝐿 ) ] under these constraints for all 𝐿 layers is challenging and is best solved automatically.
Model compression.
We consider Θ ( 𝑙 ) as the 𝑙 th layer’s total number of parameters and define the compression rate 𝑎 ( 𝑙 ) as the ratio of preserved parameters Θ ( 𝑙 ) 𝑠𝑎𝑣𝑒𝑑 to all parameters of layer 𝑙 ,
[Equation]
The compression rate for the entire network, 𝑎 , is computed analogously.
In line with He et al. (2018), we make sure that the network is not compressed below a specified global compression rate, 𝑎 𝑚𝑖𝑛 .
For this, the layer-specific compression rate 𝑎 ( 𝑙 ) is constrained to ensure that the overall compression is lower than the sum of i) already pruned parameters of all layer up to 𝑙 − 1, (cid:101) 𝜈 , ii) parameters that are about to be pruned in layer 𝑙 , (cid:101) 𝜃 ( 𝑙 ) , and iii) potentially removed parameters by the most aggressive compression rate 𝑎 𝑚𝑖𝑛 in layers from 𝑙 + 1 onward.
This mechanism allows to precisely control the network’s size.
For Heracles, we additionally adapt the action range to fit different network compression rates and allow for weight as well as channel pruning.
In Appendix A.1, we provide further details on the process.
Learning Globally Optimal Compression
Based on the above definition of network compression, we resume to define the details of learning a globally optimal compression strategy using reinforcement learning (Algorithm 1).
At each iteration of the searching process, we prune the model as outlined in Section 2.1 to determine the accuracy and robustness of the current state.
In the following, we detail the definition of the state and action space, specify the reward function used, and elaborate on the exploration phase.
State space.
For reinforcement learning, we define the RL state 𝑠 ( 𝑙 ) for layer 𝑙
based on the following eleven features:
[Equation]
All features but the compression rate of the previous layer, 𝑎 𝑝𝑟𝑒𝑣 , are dependent on layer 𝑙
For instance, the 𝑙 th layer and its output have shape 𝑘 × 𝑘 × 𝑐 𝑖𝑛 × 𝑐 𝑜𝑢𝑡 and ℎ × 𝑤 × 𝑐 𝑜𝑢𝑡 , respectively. stride refers to the striding offset used for convolutional layers, which may vary depending on input size of subsequent layers.
Additionally, we use Θ ( 𝑙 ) to denote the number of parameters of a specific layer, and specify the number of compressed parameters, (cid:101) 𝜈 , that are produced by pruning so far in preceding layers, as well as parameters remaining in latter layers, 𝜈 .
Moreover, we normalize all states to avoid overfitting.
Algorithm
Input: Pretrained Model 𝜃 , The number layers 𝐿 , RL-Agent RLA , Target rate 𝑎 𝑡𝑎𝑟𝑔𝑒𝑡 , Rate range [ 𝑎 𝑚𝑖𝑛 , 𝑎 𝑚𝑎𝑥 ] , Warm-up episodes 𝑁 𝑤𝑢𝑝 , Search episodes 𝑁 𝑠𝑟𝑐ℎ , Valid-set D 𝑣𝑎𝑙
Output: Global optimal non-uniform strategy [ 𝑎 ( 1 ) , ..., 𝑎 ( 𝐿 ) ] 1: Mask scores initialization: 𝜓 = | 𝜃 | max (| 𝜃 |) 2: for Episode = 1 . . .
𝑁 𝑠𝑟𝑐ℎ do 3: for 𝑙 = 1 . . .
𝐿 do 4: if Episode ≤ 𝑁 𝑤𝑢𝑝 then 5: 𝑎 ( 𝑙 ) =
𝑟𝑎𝑛𝑑𝑜𝑚 _ 𝑢𝑛𝑖𝑓 𝑜𝑟𝑚 ( 0 , 1 ) # Use random compression rate 6: else 7: Train RL-Agent with sampled data 8: 𝑎 ( 𝑙 ) =
N 𝑡𝑟𝑢𝑛𝑐 (cid:0) RLA ( 𝑠 ( 𝑙 ) ) , 𝜎 2 , 0 , 1 (cid:1) # Predict compression rate 9: end if 10: 𝑎 ( 𝑙 ) =
𝑎 𝑚𝑖𝑛 + 𝑎 ( 𝑙 ) · ( 𝑎 𝑚𝑎𝑥 − 𝑎 𝑚𝑖𝑛 ) # Re-scale rate 11: 𝑎 ( 𝑙 ) 𝑎𝑙𝑙𝑜𝑤 =
Max-Allow-Action ( 𝑎 ( 𝑙 ) , 𝑎 𝑡𝑎𝑟𝑔𝑒𝑡 ) # Compute maximal allowed rate 12: 𝑎 ( 𝑙 ) =
min ( 𝑎 ( 𝑙 ) , 𝑎 ( 𝑙 ) 𝑎𝑙𝑙𝑜𝑤 )
# Action control by 𝑎 𝑎𝑙𝑙𝑜𝑤 13: 𝑀 ( 𝑙 ) = 1 (cid:16) 𝜓 ( 𝑙 ) ≥ 𝜓 ( 𝑙 ) 𝐾 (cid:17) # Binary mask transformation with 𝑎 ( 𝑙 ) 14: (cid:101) 𝜃 ( 𝑙 ) =
𝑀 ( 𝑙 ) (cid:12) 𝜃 ( 𝑙 ) # Layer pruning 15: end for 16:
Robustness evaluation on (cid:101) 𝜃 with D 𝑣𝑎𝑙 17: end for Action space.
The action space of the RL agent here is (roughly speaking) the range of valid compression ratios.
In contrast to prior work (Huang et al., 2018), we do not directly produce a discrete binary mask for all layers, but use the Deep Deterministic Policy Gradient (DDPG) algorithm (Lillicrap et al., 2016) to predict a continuous compression rate along each layer.
This allows us to approach finer granularity and prune layers that have different shapes.
Consequently, the action space used for Heracles is in the range of ( 0 , 1 ] .
To facilitate more stable reinforcement learning, we use a replay buffer that is initialized in the RL agent’s warm-up stage using a random uniform distribution to generate 𝑎 ( 𝑙 ) (line 5).
In the exploration-exploitation stage of the RL process, we then use a truncated normal distribution to add noise to the action predicted by the RL agent ( RLA ) with 𝜎 = 0 . 5 which exponentially decays with each episode (line 8):
[Equation]
Further details on the action range and the action control algorithm are specified in Algorithmof Appendix A.1, where we introduce the used thresholds and elaborate on the function to selected the maximally allowed action (line 10– 12).
Exploration.
The RL agent operates on layer-based states 𝑠 ( 𝑙 ) and predicts a compression rate 𝑎 ( 𝑙 ) .
We then order the values by magnitude (OWM) and introduce a threshold 𝜓 ( 𝑙 ) 𝐾 that implements the determined compression rate 𝑎 ( 𝑙 ) .
Values lower than 𝜓 ( 𝑙 ) 𝐾 are zeroed out to construct the binary pruning mask 𝑀 (line 13).
We evaluate the robustness as well as benign accuracy of the pruned network on the validation dataset to determine the agent’s reward R (line 16) and distribute it to all state vectors.
Additionally, these are stored in the replay buffer to facilitate more stable reinforcement learning.
Reward function.
Next to the accuracy on clean, benign data Acc 𝑏𝑒𝑛 , we additionally incorporate the adversarial robustness as “adversarial accuracy” Acc 𝑎𝑑𝑣 (adversarial examples that are still classified correctly) in the reward function to yield an optimal trade-off between both:
[Equation]
For effective and fast exploration, the reward is obtained on the validation dataset only, which is sampled homogeneously from each class of the training data.
For CIFAR-10, as an example, we choose 500 images from every class, such that we yield an overall number of 5 , 000 samples for our validation dataset and thus 10 % of the training dataset.
4 Evaluation
We evaluate the performance of Heracles’s non-uniform compression strategies by enhancing state-of-the-art robust-aware pruning methods (Section 4.1), before we analyze the found strategies (Section 4.2) and discuss our method’s convergence (Section 4.3).
For this, we experiment with multiple architectures that are adversarially pre-trained on different datasets (CIFAR-10, SVHN, and ImageNet).
The pruning methods then attempt to maintain accuracy and robustness whilst achieving high compression rates of either channels or weights.
In the following, we use CIFAR-10 as the representative for small-scale datasets and report corresponding results for SVHN in the appendix, for which the class-wise imbalance (see Table 5) makes pruning even more challenging.
For both small-scale datasets, we consider ResNet18 (He et al., 2016), VGG16 (Simonyan and Zisserman, 2015) and WRN-28-4 (Zagoruyko and Komodakis, 2016), and thereby align with the experiments in related work.
Since the approaches we compare to all use slightly different variants of VGG16, we settle for the definition of Sehwag et al.
(2020) for all our experiments.
For strategy search, we bootstrap the pruning stage with 𝑁 𝑤𝑢𝑝 = 100 episodes (warm-up), and 𝑁 𝑠𝑟𝑐ℎ = 300 episodes for Deep-RL exploration-exploitation.
Further details on the experimental setup are provided in Appendix A.2.
Considered Adversaries.
We use PGD adversarial training for pre-training and fine-tuning, and also Heracles’s RL agent uses PGD adversarial examples (Madry et al., 2018) to validate the pruned network during strategy search.
To generate these, we initialize with random noise and make 10 perturbation steps per sample.
For CIFAR-10 and SVHN, models are trained with the maximal 𝑙 ∞ perturbation budget and step sizes of 8 ⁄ 255 and 2 ⁄ 255 , respectively.
For ImageNet, we use “free adversarial training” (Shafahi et al., 2019) with 4 replays, where the perturbation parameters are set to 4 ⁄ 255 and 1 ⁄ 255 .
The robustness (accuracy on adversarial examples) of the pruned models is then evaluated with multiple attack strategies each applied to the entire testing dataset with the same perturbation strength considered during training:
FGSM (Goodfellow et al., 2015), PGD-10 and PGD-20 (Madry et al., 2018), and C&W ∞ (Carlini and Wagner, 2017) optimized by PGD (20 steps). CO 2 Emission.
We have conducted all our experiments on Nvidia RTX-3090 GPU cards and have consumed about 960 GPU hours in total.
This amounts to an estimated total CO 2 emissions of 204 .
96 kgCO2eq when using Google Cloud Platform in region europe-west3 1 .
4.1 Improving related work using Heracles
We consider two approaches, Hydra and Robust-ADMM, that use uniform compression strategies for pruning neural networks, whilst maintaining both benign accuracy ( Acc 𝑏𝑒𝑛 ) and adversarial robustness, that is, the accuracy on adversarially modified inputs ( Acc 𝑎𝑑𝑣 ).
In the following, we show that it is possible to learn a non-uniform compression strategy that improves adversarial robustness when applied to Hydra or Robust-ADMM.
Moreover, Heracles is applicable to channel and weight pruning likewise—channel pruning yields a larger potential for improvement, while weight pruning is on par with related work.
We simply replace the uniformly used compression ratio of Hydra and Robust-ADMM with the strategy found by our method and present the results in Tables 2a and 2b for channel and weight pruning, respectively [Table] [Figure]
CIFAR-10.
Heracles’s compression strategy can improve the performance of VGG16, ResNet18, and WRN-28-4 pruned by Hydra as well as Robust-ADMM.
As an example, for pruning channels at a compression rate of 𝑎 𝑐 = 0 . 5, the benign and adversarial accuracy in VGG16 pruned by Hydra increase 6.90 and up to 7.03 (C&W ∞ ) percentage points, respectively.
A similar trend is observed for Robust-ADMM and other architectures as well, with WRN-28-4 being the most challenging setting.
For aggressive channel pruning ( 𝑎 𝑐 = 0 . 1), Heracles enables Hydra to even avoid an completely damaged model that yields 10 % accuracy and thus random outputs.
For weight pruning, the results are less obvious.
While Heracles’s non-uniform strategies do not yield similar high levels of improvement they are still slightly better or on par with the uniform compression rates.
Again, the results are consistent across architectures for both, Hydra and Robust-ADMM. SVHN.
As a second small-sized dataset, we have conducted experiments on the SVHN datasets.
While similar in size to the CIFAR-10 dataset, SVHN is highly unbalanced (as shown in Table 5), which can pose additional challenges. In Appendix A.3, we report details on the results and visualize the compression strategies found by Heracles for the different architectures in Fig. 4. ImageNet. Next, we apply our method for moderate pruning of ResNet50 learned on the large-scale dataset ImageNet.
Due to the size of the dataset, we reduce the RL agent’s validation dataset for strategy search to 1 % of the training data.
Sehwag et al.
(2020) demonstrate Hydra’s effectivity for pruning weights of a model learned on ImageNet and we are able to confirm these results in our experiment as shown in Table 3.
[Table]
While for Hydra results with non-uniform compression remains similar, Heracles’s strategies can improve Robust-ADMM in weight pruning.
However, the added value does not suffice to help
Robust-ADMM surpass Hydra.
Moreover, in channel pruning both Hydra and Robust-ADMM show obvious improvements over uniform compression by using our method’s strategies and both methods are then nearly on-par in robustness and performance on benign data.
4.2 Analysis of Heracles’s Strategies
We take weight pruning ( 𝑎 𝑤 = 0 . 1) and channel pruning ( 𝑎 𝑐 = 0 . 5) on CIFAR-10 as an example and inspect the global compression strategies learned by our method.
Fig.
2 visualizes the learned strategies by Heracles for VGG16, ResNet18, and WRN-28-4.
Channel pruning.
The learned strategies for channel pruning (orange lines) consistently preserve more parameters in the first several layers and prune the convolutional layers at the end of the networks more aggressively.
The strategy, however, differs in compression rates of fully connected layers of the network architectures.
For VGG16, it is notable that the RL agent preserves much more of them than the middle convolutional layers.
Interestingly, in residual block based networks (ResNet18 and WRN-28-4) the RL agent discovers pruning potential on the last connected layer.
Weight pruning.
With an overall compression rate of 𝑎 𝑤 = 0 . 1, the learned strategies for weight pruning (blue lines) are more diverse for the individual network architectures.
Networks with residual blocks share parameters which causes a more homogeneous parameter distribution on each layer.
As an example, for ResNet18 the agent does not preserve front layers but prunes layers more homogeneously.
Also for WRN-28-4 the pruning strategy approaches uniformity, which also explains the similarity in results between uniform and non-uniform strategies in Table 2b.
For
VGG16 (a conventional CNN without shortcut layers) in contrast, Heracles particularly preserves layers in the front and prunes layers in the back more distinctively.
[Table] [Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
o n
R a t e
[Figure]
4.3 Pruning Convergence
We have shown the capability of Heracles’s strategies to outperform related work for compressing channels on CIFAR-10, but our evaluation also exposes the lack of significant improvement when pruning weights.
In this section, we take ResNet18 as an example to inspect the RL agent’s searching progress in Fig.
3 to detail the underlying reasons.
Top sub-figures (a and b) refer to moderate compression, bottom ones (c and d) show very aggressive pruning.
Left sub-figures (a and c) belong to channel pruning, whereas right sub-figures (b and d) show weight pruning.
Channel pruning.
Convergence for moderate pruning at 𝑎 𝑐 = 0 .
5 works flawlessly.
After 300 steps the RL agent has successfully determined a strategy that reaches the highest reward.
While high initial exploration leads to large fluctuation, after 350 episodes the reward converges.
At 𝑎 𝑐 = 0 .
1, in turn, model performance is strongly degraded by the highly aggressive pruning.
However, the
RL agent keeps excavating better strategies, yielding good results eventually (cf.
Table 2a).
Weight pruning.
At 𝑎 𝑤 = 0 . 1, the process exhibits a certain instability due to the high sensitivity to the compression rate.
Still, the final stage converges to the overall best reward.
Differently, aggressive pruning at 𝑎 𝑤 = 0 .
01 hinders successful exploration.
The best strategies found, thus, merely realize performance on-par with uniform pruning.
[Figure]
Heracles exhibits three noteworthy limitations in practice: First, due to the randomness inherent to reinforcement learning exploration, our method exhibits variance in the yield results and ultimately cannot guarantee that the RL agent converges to an optimal strategy. Our experiments, however, show that this still succeeds in the majority of the cases, surpassing the state of the art.
Second,
Heracles’s non-uniform strategies are more effective for channel pruning than weight pruning .
This can be seen directly in the reported accuracy, but also the convergence visualization in Fig. 3, where channel pruning (left) is more effective and stable than weight pruning (right)–even at aggressive compression rates. Third, one needs to pay attention to runtime performance and exploration efficiency. Strategy search for pruning weights of the evaluated CIFAR-10 models, for instance, requires 4 .
6 × longer on average than channel pruning on a single NVIDIA RTX-3090 card.
This further emphasizes Heracles’s primary suitability for channel pruning.
Moreover, models with pruned channels will also be more resource-friendly during inference as channel pruning reduces the dimensionality of computations rather than zeroing out single values.
6 Conclusion
Striking a balance between benign accuracy and adversarial robustness during pruning is challenging.
Related work has already shown impressive results using uniform compression strategies.
With Heracles, we present a method that learns a global but layer-specific and thus non-uniform compression strategy, which can be used to further benefit existing, state-of-the-art approaches.
For instance, we increase performance for aggressive channel-pruning ( 𝑎 𝑐 = 0 . 1) with Robust-ADMM by up to 10.80 and 9.78 (C&W ∞ ) percentage points for benign and adversarial accuracy, respectively.
Weight pruning using our compression strategies has shown less distinctive results but still is slightly better or on par with related work. This is founded in the fact that here the best compression strategy often is close to uniformity. If so Heracles also finds these close-to-uniform strategies.
Such flexibility when pruning deep neural networks is crucial in practice to adapt to the model at hand. The results using non-uniform strategies are particularly promising on channel granularity, where Heracles significantly improves related work on adversarially robust pruning.
Acknowledgments
The authors thank the anonymous reviewers for their valuable suggestions, and gratefully acknowledge funding by the Helmholtz Association (HGF) within topic “46.23 Engineering Secure Systems” and by SAP S.E. under project DE-2020-021.
A Appendix
In the appendix, we extend on three aspects of our method: We begin by detailing the process of layer-wise action rate control to reach the preset target compression rate.
After that, we further describe our experimental setup, before we investigate the performance of Heracles’s non-uniform strategy on SVHN.
The action control algorithm helps to constrain the layer’s compression rate to reach the targeted global network size.
In related work (Boyd et al., 2011; Sehwag et al., 2020), the global compression rate is defined as the sparsity rather than the ratio of preserved parameters for both weight and channel pruning.
Weight pruning is apt to be controlled by the sparsity.
However, since sparsity on the channel granularity can not reflect the real network size after pruning with a global nonuniform strategy, sparsity control is not appropriate anymore.
Our approach, thus, considers the influence from layer 𝑙 on its preceding (joint) layer 𝑙 − 1 for channel pruning, such that the pruned network has the exact same size as considered in related work.
We denote determining the global parameter-wise compression rate as the function ConvertRate in Algorithm 2.
Algorithm
[Table]
A.2 Experimental setup
We conduct experiments on three different datasets: CIFAR-10, SVHN, and ImageNet.
In the following, we elaborate on the considered networks and their pre-trained performances as used in our experiment, the action range used by the action-control algorithm, and settings of the RL agent doing the strategy search.
Networks.
In related work, a variety of different deep neural network architectures are used for evaluating pruning approaches.
Some of the supposedly identical networks, however, show subtle differences.
For a fair comparison, we thus center our experiments on the architectures used by
Sehwag et al.
and build the exact same VGG16, ResNet18, and WRN-28-4 as in their opensource implementation 2 .
For our experiments on the large-scale ImageNet dataset, we additionally use ResNet50 as proposed by He et al.
We conduct 90 and 100 epochs for pre-training models on the large-scale and the small-scale datasets, respectively, with a learning rate starting at 0 .
1 while adapting it with a cosine learning-rate schedule (Loshchilov and Hutter, 2016).
For pruning with Heracles’s strategies, we employ the respective default schedules in both Hydra and Robust-ADMM to adapt the learning rate.
Table 4 summarizes the used network architectures and lists their performance after adversarial pre-training.
These models are used for all pruning experiments to ensure an identical starting point for our comparison.
[Table]
Action range. The range of compression rates ( 𝑎 𝑚𝑖𝑛 , 𝑎 𝑚𝑎𝑥 ) considered by the RL agent (the action values) have to be specified upfront, for which a few things need to be followed: For aggressive pruning, sufficient neurons must remain in each layer rather than being pruned entirely.
For moderate pruning, the agent has to be encouraged to explore different possibilities.
Consequently, we set the range as [ 0 . 01 , 0 . 8 ] and [ 0 . 005 , 0 . 5 ] , for weight pruning with compression rates of 𝑎 𝑤 = 0 . 1 and 𝑎 𝑤 = 0 .
01, respectively.
For channel pruning with compression rates of 𝑎 𝑐 = 0 .and
{+𝑎 𝑐 = 0 . 1, in turn, we use [ 0 . 1 , 1 . 0 ] and [ 0 . 05 , 0 . 5 ] . Note that for channel pruning, we maintain a compression rate of 1 . 
+}
(no compression) for the first layer to keep input information intact.
Moreover, for residual blocks, we set the compression rate on the shortcut layers to the same value as for the connected backbone layers, such that networks with residual blocks are processable by the channel-wise pruning strategy. RL agent setting. We use DDPG (Lillicrap et al., 2016) as RL agent for determining the layer’s state and, thus, predict compression rates. In our implementation, the actor network and the critic network are both constructed with two 300 neurons wide, fully connected layers.
Moreover, the size of the replay buffer is set to 200 times the number of prunable layers in the neural network—for instance, for pruning weights of VGG16 the size equals to 200 × 16 = 3 , 200.
The training of the
RL agent is performed with learning rates 0.01 and 0.001 respectively on critic and actor in DDPG. And we use a soft update of 0 . 01 on the target model.
During the 300 episodes long RL search phase, we train the agent for 20 epochs on states sampled from the replay buffer with a batch size of 128. For better exploration, we additionally set the exponential decay 𝛿 to 0 . 99. A.3 Heracles’s performance on SVHN In contrast to CIFAR-10, neither training nor testing data of SVHN is balanced.
While we use the common accuracy to determine robustness and the model’s natural performance in our experiments on CIFAR-10, for SVHN we hence use the balanced accuracy (Brodersen et al., 2010) as the evaluation metric.
Table 5 shows the class-wise distribution of the SVHN dataset.
For both datasets, a completely damaged model with random predictions thus is indicated by a (balanced) accuracy of 10 % as the datasets contain 10 classes each.
Channel pruning.
Table 6a summarizes the experimental results of channel pruning on SVHN.
With the help of Heracles, both Hydra and Robust-ADMM improve robustness and natural accuracy at 𝑎 𝑐 = 0 . 5.
Interestingly, pruning ResNet18 by Robust-ADMM equipped with our strategies even yields up to 9 .
05 percentage points higher robustness against PGD-10 than uniform compression.
However, the variance is rather high, which originates the inherent randomness of reinforcement
[Table]
learning. The accuracy on benign data, in turn, remains stable across all strategies and our method yields a pruned model with at least a performance on-par with the uniform strategy. This clearly shows the positive impact of Heracles in moderate pruning. Aggressive channel pruning ( 𝑎 𝑐 = 0 . 1) is more challenging for SVHN, though.
Neither uniform nor our non-uniform compression can [Table] prevent VGG16 from complete damage (10 %) and the uniform strategy adapts better to WRN-28when pruning with Robust-ADMM. In all other situations, Heracles’s strategies still yield on-par or better performance in aggressive channel pruning as well.
Even in situations where uniform compression yields a completely damaged model (ResNet18 and WRN-28-4 using Hydra), our method selects a strategy that results in meaningful and competitive results. Furthermore, we visualize the strategies learned by Heracles for moderate channel pruning in Fig. 4 (orange lines).
Similar to our results on CIFAR-10, the RL agent discovers high redundancy (and thus pruning potential) in the middle convolutional layers for VGG16. Differently, the compression on ResNet18 is nearly uniform along convolutional layers where the RL agent, however, still acknowledges the higher importance of the final fully connected layer. Weight pruning. In comparison to channel pruning, Heracles is more stable and achieves higher robustness on VGG16 and ResNet18 (cf.
Table 6b). For both moderate ( 𝑎 𝑤 = 0 . 1) and aggressive ( 𝑎 𝑤 = 0 . 01) pruning, state-of-the-art methods can benefit from Heracles’s strategies and improves up to 30.66 and 26.77 percentage points for benign accuracy and FGSM robustness, respectively. However, also here similar instability issues happen for pruning ResNet18 with Robust-ADMM.
In Fig. 4b, we see that the RL agent approaches a strategy close to uniformity for WRN-28-4 and ResNet18.
For the latter, in particular for early layers the variance is larger than for layers further back. Overall, we observe that non-uniform compression is more beneficial for channel pruning than weight pruning. However, for aggressive weight pruning ( 𝑎 𝑤 = 0 . 01) by Robust-ADMM, WRN-28-4’s performance is completely damaged for uniform and non-uniform strategies alike.
[Table]
fc fc fc [Figure]