Sparsity Meets Robustness: Channel Pruning for the Feynman-Kac Formalism Principled Robust Deep Neural Nets

1 INTRODUCTION . Robust deep neural nets ( DNNs ) compression is a fundamental problem for secure AI applications in resource-constrained environments such as biometric verification and facial login on mobile devices , and computer vision tasks for the internet of things ( IoT ) ( Cheng et al. , 2017 ; Yao et al. , 2017 ; Mohammadi et al. , 2018 ) . Though compression and robustness have been separately addressed in recent years , much less is studied when both players are present and must be satisfied . To date , many successful techniques have been developed to compress naturally trained DNNs , including neural architecture re-design or searching ( Howard et al. , 2017 ; Zhang et al. , 2018b ) , pruning including structured ( weights sparsification ) ( Han et al. , 2015 ; Srinivas & Babu , 2015 ) and unstructured ( channel- , filter- , layer-wise sparsification ) ( Yang et al. , 2019 ; He et al. , 2017 ) , quantization ( Zhou et al. , 2017 ; Yin et al. , 2019 ; Courbariaux et al. , 2016 ) , low-rank approximation ( Denil et al. , 2013 ) , knowledge distillation ( Polino et al. , 2018 ) , and many more ( Alvarez & Salzmann , 2016 ; Liu et al. , 2015 ) . The adversarially trained ( AT ) DNN is more robust than the naturally trained ( NT ) DNN to adversarial attacks ( Madry et al. , 2018 ; Athalye et al. , 2018 ) . However , adversarial training ( denoted as AT if no ambiguity arises , and the same for NT ) also dramatically reduces the sparsity of the trained DNN ’ s weights . As shown in Fig . 1 , start from the same default initialization in PyTorch , the NT ResNet20 ’ s weights are much sparser than that of the AT counterpart , for instance , the percent of weights that have magnitude less than 10−3 for NT and AT ResNet20 are 8.66 % and 3.64 % ( averaged over 10 trials ) , resp . This observation motivates us to consider the following two questions : • 1 . Can we re-design the neural architecture with minimal change on top of the existing one such that the new DNN has sparser weights and better robustness and accuracy than the existing one ? • 2 . Can we develop efficient compression algorithms to compress the AT DNNs with minimal robustness and accuracy degradations ? We note that under the AT , the recently proposed Feynman-Kac formalism principled ResNet ensemble ( Wang et al. , 2019a ) has much sparser weights than the standard ResNet , which gives a natural answer to the first question above . To answer the second question , we leverage state-of-the-art relaxed augmented Lagrangian based sparsification algorithms ( Dinh & Xin , 2018 ; Yang et al. , 2019 ) to perform both structured and unstructured pruning for the AT DNNs . We focus on unstructured and channel pruning in this work . 1.1 NOTATION . Throughout this paper we use bold upper-case letters A , B to denote matrices , bold lower-case letters x , y to denote vectors , and lower case letters x , y and α , β to denote scalars . For vector x = ( x1 , . . . , xd ) > , we use ‖x‖ = ‖x‖2 = √ x21 + · · ·+ x2d to represent its ` 2-norm ; ‖x‖1 =∑d i=1 |xi| to represent its ` 1-norm ; and ‖x‖0 = ∑d i=1 χ { xi 6=0 } to represent its ` 0-norm . For a function f : Rd → R , we use ∇f ( · ) to denote its gradient . Generally , wt represents the set of all parameters of the network being discussed at iteration t , e.g . wt = ( wt1 , w t 2 , ... , w t M ) , where w t j is the weight on the jth layer of the network at the t-th iteration . Similarly , ut = ( ut1 , u t 2 , ... , u t M ) is a set of weights with the same dimension as wt , whose value depends on wt and will be defined below . We use N ( 0 , Id×d ) to represent the d-dimensional Gaussian , and use notation O ( · ) to hide only absolute constants which do not depend on any problem parameter . 1.2 ORGANIZATION . This paper is organized in the following way : In section 2 , we list the most related work to this paper . In section 3 , we show that the weights of the recently proposed Feynman-Kac formalism principled ResNet ensemble are much sparser than that of the baseline ResNet , providing greater efficiency for compression . In section 4 , we present relaxed augmented Lagrangian-based algorithms along with theoretical analysis for both unstructured and channeling pruning of AT DNNs . The numerical results are presented in section 5 , followed by concluding remarks . Technical proofs and more related results are provided in the appendix . 2 RELATED WORK . Compression of AT DNNs : Gui et al . ( 2019 ) considered a low-rank form of the DNN weight matrix with ` 0 constraints on the matrix factors in the AT setting . Their training algorithm is a projected gradient descent ( PGD ) ( Madry et al. , 2018 ) based on the worst adversary . In their paper , the sparsity in matrix factors are unstructured and require additional memory . Sparsity and Robustness : Guo et al . ( 2018 ) shows that there is a relationship between the sparsity of weights in the DNN and its adversarial robustness . They showed that under certain conditions , sparsity can improve the DNN ’ s adversarial robustness . The connection between sparsity and robustness has also been studied recently by Ye et al . ( 2019 ) , Rakin et al . ( 2019 ) , and et al . In our paper , we focus on designing efficient pruning algorithms integrated with sparse neural architectures to advance DNNs ’ sparsity , accuracy , and robustness . Feynman-Kac formalism principled Robust DNNs : Neural ordinary differential equations ( ODEs ) ( Chen et al. , 2018 ) are a class of DNNs that use an ODE to describe the data flow of each input data . Instead of focusing on modeling the data flow of each individual input data , Wang et al . ( 2019a ; 2018a ) ; Li & Shi ( 2017 ) use a transport equation ( TE ) to model the flow for the whole input distribution . In particular , from the TE viewpoint , Wang et al . ( 2019a ) modeled training ResNet ( He et al. , 2016 ) as finding the optimal control of the following TE ∂u ∂t ( x , t ) +G ( x , w ( t ) ) · ∇u ( x , t ) = 0 , x ∈ R d , u ( x , 1 ) = g ( x ) , x ∈ Rd , u ( xi , 0 ) = yi , xi ∈ T , with T being the training set . ( 1 ) where G ( x , w ( t ) ) encodes the architecture and weights of the underlying ResNet , u ( x , 0 ) serves as the classifier , g ( x ) is the output activation of ResNet , and yi is the label of xi . Wang et al . ( 2019a ) interpreted adversarial vulnerability of ResNet as arising from the irregularity of u ( x , 0 ) of the above TE . To enhance u ( x , 0 ) ’ s regularity , they added a diffusion term , 12σ 2∆u ( x , t ) , to the governing equation of ( 1 ) which resulting in the convection-diffusion equation ( CDE ) . By the Feynman-Kac formula , u ( x , 0 ) of the CDE can be approximated by the following two steps : • Modify ResNet by injecting Gaussian noise to each residual mapping . • Average the output of n jointly trained modified ResNets , and denote it as EnnResNet . Wang et al . ( 2019a ) have noticed that EnResNet can improve both natural and robust accuracies of the AT DNNs . In this work , we leverage the sparsity advantage of EnResNet to push the sparsity limit of the AT DNNs . 3 REGULARITY AND SPARSITY OF THE FEYNMAN-KAC FORMALISM PRINCIPLED ROBUST DNNS ’ WEIGHTS . From a partial differential equation ( PDE ) viewpoint , a diffusion term to the governing equation ( 1 ) not only smooths u ( x , 0 ) , but can also enhance regularity of the velocity field G ( x , w ( t ) ) ( Ladyženskaja et al. , 1988 ) . As a DNN counterpart , we expect that when we plot the weights of EnResNet and ResNet at a randomly select layer , the pattern of the former one will look smoother than the latter one . To validate this , we follow the same AT with the same parameters as that used in ( Wang et al. , 2019a ) to train En5ResNet20 and ResNet20 , resp . After the above two robust models are trained , we randomly select and plot the weights of a convolutional layer of ResNet20 whose shape is 64×64×3×3 and plot the weights at the same layer of the first ResNet20 in En5ResNet20 . As shown in Fig . 2 ( a ) and ( b ) , most of En5ResNet20 ’ s weights are close to 0 and they are more regularly distributed in the sense that the neighboring weights are closer to each other than ResNet20 ’ s weights . The complete visualization of this randomly selected layer ’ s weights is shown in the appendix . As shown in Fig . 2 ( c ) and ( d ) , the weights of En5ResNet20 are more concentrated at zero than that of ResNet20 , and most of the En5ResNet20 ’ s weights are close to zero . 4 UNSTRUCTURED AND CHANNEL PRUNING WITH AT . 4.1 ALGORITHMS . In this subsection , we introduce relaxed , augmented Lagrangian-based , pruning algorithms to sparsify the AT DNNs . The algorithms of interest are the Relaxed Variable-Splitting Method ( RVSM ) ( Dinh & Xin , 2018 ) for weight pruning ( Algorithm 1 ) , and its variation , the Relaxed Group-wise Splitting Method ( RGSM ) ( Yang et al. , 2019 ) for channel pruning ( Algorithm 2 ) . Our approach is to apply the RVSM/RGSM algorithm together with robust PGD training to train and sparsify the model from scratch . Namely , at each iteration , we apply a PGD attack to generate adversarial images x′ , which are then used in the forward-propagation process to generate preditions y′ . The back-propagation process will then compute the appropriate loss function and apply RVSM/RGSM to update the model . Previous work on RVSM mainly focused on a one-hidden layer setting ; In this paper , we extend this result to the general setting . To the best of our knowledge , this is the first result that uses RVSM/RGSM in an adversarial training scenario . To explain our choice of algorithm , we discuss a classical algorithm to promote sparsity of the target weights , the alternating direction of multiplier method ( ADMM ) ( Boyd et al. , 2011 ; Goldstein & Osher , 2009 ) . In ADMM , instead of minimizing the original loss function f ( w ) , we seek to minimize the ` 1 regularized loss function , f ( w ) + λ‖u‖1 , by considering the following augmented Lagrangian L ( w , u , z ) = f ( w ) + λ ‖u‖1 + 〈z , w − u〉+ β 2 ‖w − u‖2 , λ , β ≥ 0 . ( 2 ) which can be easily solved by applying the following iterations wt+1 ← arg minw Lβ ( w , ut , zt ) ut+1 ← arg minu Lβ ( wt+1 , u , zt ) zt+1 ← zt + β ( wt+1 − ut+1 ) ( 3 ) Although widely used in practice , ADMM has several drawbacks when it is used to regularize DNN ’ s weights . First , one can improve the sparsity of the final learned weights by replacing ‖u‖1 with ‖u‖0 ; but ‖ · ‖0 is not differentiable , thus current theory of optimization does not apply ( Wang et al. , 2018b ) . Second , the update wt+1 ← arg minw Lβ ( w , ut , zt ) is not a reasonable step in practice , as one has to fully know how f ( w ) behaves . In most ADMM adaptation on DNN , this step is replaced by a simple gradient descent . Third , the Lagrange multiplier term , 〈z , w − u〉 , seeks to close the gap between wt and ut , and this in turn reduces sparsity of ut . The RVSM we will implement is a relaxation of ADMM . RVSM drops the Lagrangian multiplier , and replaces λ‖u‖1 with λ‖u‖0 , and resulting in the following relaxed augmented Lagrangian Lβ ( w , u ) = f ( w ) + λ ‖u‖0 + β 2 ‖w − u‖2 . ( 4 ) The above relaxed augmented Lagrangian can be solved efficiently by the iteration in Algorithm 1 . RVSM can resolve all the three issues associated with ADMM listed above in training robust DNNs with sparse weights : First , by removing the linear term 〈z , w − u〉 , one has a closed form formula for the update of ut without requiring ‖u‖0 to be differentiable . Explicitly , ut = H√ 2λ/β ( wt ) = ( wt1χ { |w1| > √ 2λ/β } , ... , w t dχ { |w1| > √ 2λ/β } ) , where Hα ( · ) is the hard- thresholding operator with parameter α . Second , the update of wt is a gradient descent step itself , so the theoretical guarantees will not deviate from practice . Third , without the Lagrange multiplier term zt , there will be a gap between wt and ut at the limit ( finally trained DNNs ) . However , the limit of ut is much sparser than that in the case of ADMM . At the end of each training epoch , we replace wt by ut for the validation process . Numerical results in Section 5 will show that the AT DNN with parameters ut usually outperforms the traditional ADMM in both accuracy and robustness . Algorithm 1 RVSM Input : η , β , λ , maxepoch , maxbatch Initialization : w0 Define : u0 = H√ 2λ/β ( w0 ) for t = 0 , 1 , 2 , ... , maxepoch do for batch = 1 , 2 , ... , maxbatch do wt+1 ← wt − η∇f ( wt ) − ηβ ( wt − ut ) ut+1 ← arg minu Lβ ( u , wt ) = H√2λ/β ( w t ) end for end for Algorithm 2 RGSM Input : η , β , λ1 , λ2 , maxepoch , maxbatch Objective : f̃ ( w ) = f ( w ) + λ2‖w‖GL Initialization : Initialize w0 , define u0 for g = 1 , 2 , ... , G do u0g = Proxλ1 ( w 0 g ) end for for t = 0 , 1 , 2 , ... , maxepoch do for batch = 1 , 2 , ... , maxbatch do wt+1 = wt−η∇f̃ ( wt ) −ηβ ( wt−ut ) for g = 1 , 2 , ... , G do ut+1g = Proxλ1 ( w t g ) end for end for end for RGSM is a method that generalizes RVSM to structured pruning , in particular , channel pruning . Let w = { w1 , ... , wg , ... , wG } be the grouped weights of convolutional layers of a DNN , whereG is the total number of groups . Let Ig be the indices of w in group g. The group Lasso ( GLasso ) penalty and group- ` 0 penalty ( Yuan & Lin , 2007 ) are defined as ‖w‖GL : = G∑ g=1 ‖wg‖2 , ‖w‖G ` 0 : = G∑ g=1 1‖wg‖2 6=0 ( 5 ) and the corresponding Proximal ( projection ) operators are ProxGL , λ ( wg ) : = sgn ( wg ) max ( ‖wg‖2 − λ , 0 ) , ProxG ` 0 , λ ( wg ) : = wg1‖wg‖2 6= √ 2λ ( 6 ) where sgn ( wg ) : = wg/‖wg‖2 . The RGSM method is described in Algorithm 2 , which improves on adding group Lasso penalty directly in the objective function ( Wen et al. , 2016 ) for natural DNN training ( Yang et al. , 2019 ) .

Summary: This paper focuses on the study of sparse neural architectures and efficient DNN compression algorithms for the robust and accurate deep learning. The authors apply relaxed augmented Lagrangian based sparsification algorithms to perform both unstructured and channel pruning for the AT DNNs. By combing with the sparsity merit of EnResNet, their sparsification algorithms further boost the sparsity limit of the AT DNNs, leading to much better robustness and accuracy. Their approach demonstrates superior natural and robust accuracies under several benchmark attacks.