Feature Flow Regularization: Improving Structured  Sparsity in Deep Neural Networks

1 INTRODUCTION . Deep neural networks ( DNNs ) have achieved huge success in a wide range of applications . Meanwhile , DNNs require considerably more computational cost and storage space as they become deeper in order to achieve higher accuracy . Denil et al . ( 2013 ) demonstrated that there is significant redundancy in the parameterization of DNNs . The Lottery Ticket Hypothesis ( Frankle & Carbin , 2019 ) conjectures that there exist sparse sub-networks that can obtain a comparable accuracy with the original network when trained in isolation . Model compression methods have been proposed to balance accuracy and model complexity , e.g . weight pruning ( Drucker & Le Cun , 1992 ; Hassibi & Stork , 1993 ; Han et al. , 2015 ; Guo et al. , 2016 ; Hu et al. , 2016 ; Li et al. , 2016 ) and quantization ( Gong et al. , 2014 ) , low-rank approximation ( Denton et al. , 2014 ; Jaderberg et al. , 2014 ; Liu et al. , 2015 ) , and sparsity structure learning ( Wen et al. , 2016 ; Alvarez & Salzmann , 2016 ; Zhou et al. , 2016 ; Liu et al. , 2017 ; Louizos et al. , 2018 ; Gao et al. , 2019 ; Yuan et al. , 2020 ) . Weight pruning removes less important parameters in the network . In particular , filter pruning ( Li et al. , 2016 ) removes entire filters in the network together with their related channels , which can compress and accelerate DNNs efficiently . Existing structured pruning methods can be divided into two categories : parameter-based methods ( Li et al. , 2016 ; Molchanov et al. , 2016 ; Liu et al. , 2017 ; He et al. , 2018 ; Lin et al. , 2019 ; He et al. , 2019 ; Zhuang et al. , 2020 ; Liebenwein et al. , 2020 ) that use some criteria to identify unimportant filters and remove them , and feature-based methods ( Luo et al. , 2017 ; He et al. , 2017 ; Zhuang et al. , 2018 ; Ye et al. , 2018 ; Li et al. , 2020 ; Lin et al. , 2020 ; Tang et al. , 2020 ) that select unimportant feature maps and then remove related filters and channels . For example , Li et al . ( 2020 ) incorporated two feature map selections : discovering features with low diversity and removing features that have high similarities with others . In this paper , we propose a new regularization method on the trajectory connecting features of adjacent hidden layers , namely feature flow regularization ( FFR ) . FFR smooths the trajectory of features , which implicitly improves the structured sparsity in DNN . Our motivation is that the trajectory of data along the network reflects the DNN structure . Shorter and straighter trajectory corresponds to an efficient and sparse structure of DNN . An illustration is given in Figure 1b . Our main contributions are : ( 1 ) We propose a new regularization ( FFR ) on the trajectory connecting the features of hidden layers , to improve the structured sparsity in DNN from a perspective of the trajectory of data along the network . This method is different from the existing sparsity structure learning methods , which directly impose regularization or constraints on the parameters . Our method is also different from those pruning methods based on feature maps , which use the information of the feature map individually or in pairs ( for similarity ) without global relationship . ( 2 ) We analyze the effect of FFR applied to convolutional layer and residual layer , and show that FFR encourages DNN to learn a sparse structure during training by penalizing the sparsity of both parameters and features . ( 3 ) Experimental results show that FFR achieves a comparable or even better pruning ratio in terms of parameters and FLOPs than recent state-of-the-art pruning methods . 2 RELATED WORK . Filter pruning . Various criteria for filter selection in pruning have been proposed . Li et al . ( 2016 ) used L1 norm to select unimportant filters and removed the filters whose norm is lower than the given threshold together with their connecting feature maps . Molchanov et al . ( 2016 ) measured the importance of filters based on the change in the cost function induced by pruning . Luo et al . ( 2017 ) ; He et al . ( 2017 ) formulated pruning as a constraint optimization problem and selected most representative neurons based on minimizing the reconstitution error . Lin et al . ( 2019 ) pruned filters as well as other structures by generative adversarial learning . He et al . ( 2019 ) pruned redundant filters utilizing geometric correlation among filters in the same layer . Hu et al . ( 2016 ) ; Zhuang et al . ( 2018 ) ; Li et al . ( 2020 ) ; Lin et al . ( 2020 ) removed filters based on the information , e.g . sparsity , rank or diversity , of feature maps that are generated by the filters . Our method learns sparse DNN during training and adopts magnitude-based pruning scheme after training . Sparsity regularization . Some studies introduced sparsity regularization to find sparse structure of DNN . A commonly used strategy is to impose group Lasso and relaxed L0 regularization ( Zhou et al. , 2016 ; Wen et al. , 2016 ; Alvarez & Salzmann , 2016 ; Louizos et al. , 2018 ) . Liu et al . ( 2017 ) ; Huang & Wang ( 2018 ) associated a scaling factor with feature maps and imposed regularization on these scaling factors during training to automatically identify unimportant channels . The feature maps with small scaling factor values will be pruned . Gao et al . ( 2019 ) imposed a Cross-Layer grouping and a Variance Aware regularization on the parameters to improve the structured sparsity for residual models . Zhuang et al . ( 2020 ) used polarization regularizer on scaling factors . Yuan et al . ( 2020 ) proposed a method to dynamically grow deep networks by continuously sparsifying structured parameter sets . Different from these available methods that directly introduce regularization on the parameters or scaling factors , our FFR imposes regularization on the trajectory connecting features of hidden layers to control the parameters and enforce structured sparsity implicitly . 3 METHOD . 3.1 FEATURE FLOW REGULARIZATION . We define feature flow as the trajectory formed by connecting the output features of adjacent hidden layers . For a DNN , the collection of trajectories with different input data from the training dataset reflect the network structure . We control these trajectories to obtain a sparse network . Consider the forward propagation of a DNN with L layers { xl } l=0,1 , ... , L : xl+1 = hl ( xl , wl ) , ( 1 ) where xl+1 is the output feature of the l-th layer , hl is the mapping in the l-th layer , and wl is the collection of trainable parameters . Introducing a temporal partition : { tl = l/L } Ll=0 with time interval ∆t = 1/L , and regarding xl as the value of a function x ( t ) at time step tl , without considering dimensional consistency , Eq . ( 1 ) can be rewritten as ( He et al. , 2016 ; E , 2017 ; Lu et al. , 2017 ; Chen et al. , 2018 ) x ( tl+1 ) = x ( tl ) + ∆t · ĥl ( x ( tl ) , wl ) , ( 2 ) where ĥl = ( hl − xl ) /∆t . This can be interpreted as a discretization of evolution along a trajectory of the network described by an ordinary differential equation ( E , 2017 ; Lu et al. , 2017 ; Chen et al. , 2018 ) dx ( t ) dt = ĥ ( x ( t ) , w ( t ) , t ) . ( 3 ) The feature flow is the trajectory formed by connecting features of { xl } , and is denoted by Γ { xl } . See Figure 1a for an illustration of the feature flow . We regard the trajectory x ( t ) as well as the feature flow trajectory Γ { xl } as a ” curve ” . Recall that for a curve γ ( t ) : [ 0 , 1 ] → RD with arc length parameter s , its length C ( γ ) and total absolute curvature K ( γ ) are C ( γ ) : = ∫ 1 0 ‖γ′ ( t ) ‖dt , K ( γ ) : = ∫ C ( γ ) 0 |κ ( s ) |ds , ( 4 ) where κ ( s ) = ‖γ′′ ( s ) ‖ is the curvature of the curve . We introduce feature flow regularization ( FFR ) to improve the structured sparsity of DNN , borrowing the definitions of length and total absolute curvature of a curve to the feature flow , i.e. , the trajectory formed by connecting features of hidden layers . For a feature flow associated with hidden layer features { xl } l=0,1 , ... , L , the FFR is R ( x ) : = k1C ( x ) + k2K ( x ) , ( 5 ) where C ( x ) = L−1∑ l=0 ‖xl+1 − xl‖ , K ( x ) = L−1∑ l=1 ‖xl+1 − 2xl + xl−1‖ , ( 6 ) with ‖ · ‖ being the L1 norm , and k1 , k2 > 0 the hyperparameters . Here up to some constant factors , C ( x ) ∼ ∑L−1 l=0 ‖x′ ( tl ) ‖ is an approximation of the total length of the trajectory x ( t ) , and K ( x ) ∼ ∑L−1 l=1 |κ ( tl ) | is an approximation of its total absolute curvature , where { xl = x ( tl ) } is a discretization of trajectory x ( t ) with time partition { tl = l/L } Ll=0 . FFR smooths the feature flow by controlling the length and curvature . Intuitively , the length term in Eq . ( 6 ) makes the feature flow short , and the curvature term in Eq . ( 6 ) keeps the feature flow from bending too much . As a result , DNN trained under FFR has a more sparse structure ; See demonstration in Figure 1b . More quantitative analysis is given in Sec . 4 . We further give an illustration example in Figure 1c , which is a two-dimensional visualization showing the smooth effect of FFR on the feature flow of a ResNet . In this example , the input , features and output are all points in two dimensions , and the feature flows are actual curves in two dimensions . We can see that the feature flow under FFR is shorter and straighter . More detail of this illustration example is given in Appendix Sec . A.1 . The curvature term also encourages more uniform distribution of hidden features on the trajectory from the input to the output , which makes it easier and more accurate to learn with sparse parameters ; see the discussion in appendix A.4 and Figure 7 3.2 FFR APPLIED TO DNN AND PRUNING . For a DNN and training dataset { ( x ( j ) , y ( j ) } N j=1 , using Eq . ( 5 ) , the FFR is : R ( X ( j ) ) = k1 L−1∑ i=0 ‖x ( j ) l+1 − x ( j ) l ‖+ k2 L−1∑ i=1 ‖x ( j ) l+1 − 2x ( j ) l + x ( j ) l−1‖ , ( 7 ) where X ( j ) = { xl , l = 0 , . . . , L } ( j ) denotes the set of output features of hidden layers for the j-th input data . The loss function with FFR in training is : 1 N N∑ j=1 [ J ( x ( j ) , y ( j ) , W ) +R ( X ( j ) ) ] , ( 8 ) where J ( x , y , W ) is the loss function before applying FFR . Note that two hidden states xl , xl+1 may have different dimensions . To fix the dimension mismatch problem , we adopt the same strategy as He et al . ( 2016 ) , i.e , using a linear projection Pl by the shortcut connections to match the dimensions . We replace xl in Eq . ( 7 ) by Plxl , where Pl is the learnt projection matrix and will be treated as learnable parameters in training . In implementation , we first group the features according to the stage ( features dimensions ) : X ( j ) = G⋃ g=1 { xg,1 , xg,2 , . . . , xg , lg } ( j ) , L = G∑ g=1 lg . ( 9 ) Here G is the number of stages in X ( j ) and lg is the number of hidden layers in stage g. Secondly , we use the projection matrix to link different groups since the dimensional mismatch only occurs at the first feature of each stage . Using this method , the FFR of X ( j ) becomes : R ( X ( j ) ) = k1 G∑ g=1 ‖x ( j ) g,1 − Pgx ( j ) g−1 , l ( g−1 ) ‖+ lg−1∑ i=1 ‖x ( j ) g , i+1 − x ( j ) g , i‖ + k2 G∑ g=1 ‖x ( j ) g,2 − 2x ( j ) g,1 + Pgx ( j ) g−1 , l ( g−1 ) ‖+ lg−1∑ i=2 ‖x ( j ) g , i+1 − 2x ( j ) g , i + x ( j ) g , i−1‖ . ( 10 ) In the meantime , the hyperparameters k1 , k2 may vary with the feature dimensions so that FFR can uniformly control the features at different stages . In our experiments , we adjust k1 , k2 to be inversely proportional to the scale of feature maps . Moreover , the FFR process can be generalized to the case where we choose features every several layers . We can denote the selected hidden layers as L = { li , i = 0 , 1 , . . . , m , m = # L } . Then we apply FFR to the feature flow that connects the features in the set X ( j ) = { xl , l ∈ L } . After training , we conduct one-shot filter pruning : removing filters with small magnitude and removing channels in the next layer that convolve with the feature maps generated by the pruned filters . Finally , we fine tune the pruned network for a few epochs . Our FFR training and pruning method is summarized in Algorithm 1 . Algorithm 1 FFR Training and One-shot Pruning Require : training dataset { ( x ( j ) , y ( j ) } N j=1 , a neural network and hyperparameters k1 , k2 . Pre-step 1 : group the features in X = { xl , l ∈ L } according to the stage as in Eq . ( 9 ) , Pre-step 2 : write down FFRR ( X ( j ) ) in Eq . ( 10 ) for each paired data ( x ( j ) , y ( j ) ) . Training step : train the network under loss function with FFR given in Eq . ( 8 ) . Pruning step : remove filters and the corresponding channels in the trained model , fine tune the pruned model for a few epochs . return a compact neural network .

Summary: This paper deals with network pruning from a fresh perspective. It essentially argues that networks that implement simple mappings are more amenable to pruning. The authors support this claim with experimental evidence. Personally I think this idea is intellectually appealing, and may have uses beyond the network pruning setting. Some concerns related to prior art linger.