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: The paper proposes a new method, called feature flow regularization, for structured sparsity and pruning in deep neural networks. In the proposed method, the length and curvature of trajectories connecting the features of adjacent hidden layers are penalized which is claimed to implicitly result in structured sparsity. Experiments compare the method with several state-of-the-art methods on two datasets and architectures.