Split LBI for Deep Learning: Structural Sparsity via Differential Inclusion Paths

1 INTRODUCTION . The expressive power of deep neural networks comes from the millions of parameters , which are optimized by Stochastic Gradient Descent ( SGD ) ( Bottou , 2010 ) and variants like Adam ( Kingma & Ba , 2015 ) . Remarkably , model over-parameterization helps both optimization and generalization . For optimization , over-parameterization may simplify the landscape of empirical risks toward locating global optima efficiently by gradient descent method ( Mei et al. , 2018 ; 2019 ; Venturi et al. , 2018 ; Allen-Zhu et al. , 2018 ; Du et al. , 2018 ) . On the other hand , over-parameterization does not necessarily result in a bad generalization or overfitting ( Zhang et al. , 2017 ) , especially when some weight-size dependent complexities are controlled ( Bartlett , 1997 ; Bartlett et al. , 2017 ; Golowich et al. , 2018 ; Neyshabur et al. , 2019 ) . However , compressive networks are desired in many real world applications , e.g . robotics , selfdriving cars , and augmented reality . Despite that ` 1 regularization has been applied to deep learning to enforce the sparsity on weights toward compact , memory efficient networks , it sacrifices some prediction performance ( Collins & Kohli , 2014 ) . This is because that the weights learned in neural networks are highly correlated , and ` 1 regularization on such weights violates the incoherence or irrepresentable conditions needed for sparse model selection ( Donoho & Huo , 2001 ; Tropp , 2004 ; Zhao & Yu , 2006 ) , leading to spurious selections with poor generalization . On the other hand , ` 2 regularization is often utilized for correlated weights as some low-pass filtering , sometimes in the form of weight decay ( Loshchilov & Hutter , 2019 ) or early stopping ( Yao et al. , 2007 ; Wei et al. , 2017 ) . Furthermore , group sparsity regularization ( Yuan & Lin , 2006 ) has also been applied to neural networks , such as finding optimal number of neuron groups ( Alvarez & Salzmann , 2016 ) and exerting good data locality with structured sparsity ( Wen et al. , 2016 ; Yoon & Hwang , 2017 ) . Yet , without the aid of over-parameterization , directly training a compressive model architecture may meet the obstacle of being trapped in local optima in contemporary experience . Alternatively , researchers in practice typically start from training a big model using common task datasets like ImageNet , and then prune or distill such big models to small ones without sacrificing too much of the performance ( Jaderberg et al. , 2014 ; Han et al. , 2015 ; Zhu et al. , 2017 ; Zhou et al. , 2017 ; Zhang et al. , 2016 ; Li et al. , 2017 ; Abbasi-Asl & Yu , 2017 ; Yang et al. , 2018 ; Arora et al. , 2018 ) . In particular , a recent study ( Frankle & Carbin , 2019 ) created the lottery ticket hypothesis based on empirical observations : “ dense , randomly-initialized , feed-forward networks contain subnetworks ( winning tickets ) that – when trained in isolation – reach test accuracy comparable to the original network in a similar number of iterations '' . How to effectively reduce an over-parameterized model thus becomes the key to compressive deep learning . Yet , Liu et al . ( 2019 ) raised a question , is it necessary to fully train a dense , over-parameterized model before finding important structural sparsity ? In this paper , we provide a novel answer by exploiting a dynamic approach to deep learning with structural sparsity . We are able to establish a family of neural networks , from simple to complex , by following regularization paths as solutions of differential inclusions of inverse scale spaces . Our key idea is to design some dynamics that simultaneously exploit over-parameterized models and structural sparsity . To achieve this goal , the original network parameters are lifted to a coupled pair , with one weight set W of parameters following the standard gradient descend to explore the over-parameterized model space , while the other set of parameters learning structure sparsity in an inverse scale space , i.e. , structural sparsity set Γ . The large-scale important parameters are learned at a fast speed while the small unimportant ones are learned at a slow speed . The two sets of parameters are coupled in an ` 2 regularization . The dynamics enjoys a simple discretization , i.e . the Split Linearized Bregman Iteration ( SplitLBI ) , with provable global convergence guarantee shown in this paper . Here , SplitLBI is a natural extension of SGD with structural sparsity exploration : SplitLBI reduces to the standard gradient descent method when the coupling regularization is weak , while it leads to a sparse mirror descent when the coupling is strong . Critically , SplitLBI enjoys a nice property that important subnet architecture can be rapidly learned via the structural sparsity parameter Γ following the iterative regularization path , without fully training a dense network first . Particularly , the support set of structural sparsity parameter Γ learned in the early stage of this inverse scale space discloses important sparse subnet architectures . Such architectures can be fine-tuned or retrained to achieve comparable test accuracy as the dense , over-parameterized networks . As a result , the structural sparsity parameter Γ may enable us to rapidly find “ winning tickets ” in early training epochs for the “ lottery ” of identifying successful subnetworks that bear comparable test accuracy to the dense ones . This point is empirically validated in our experiments . Historically , the Linearized Bregman Iteration ( LBI ) was firstly proposed in applied mathematics as iterative regularization paths for image reconstruction and compressed sensing ( Osher et al. , 2005 ; Yin et al. , 2008 ) , later applied to logistic regression ( Shi et al. , 2013 ) . The convergence analysis was given for convex problems ( Yin et al. , 2008 ; Cai et al. , 2009 ) , yet remaining open for non-convex problems met in deep learning . Osher et al . ( 2016 ) established statistical model selection consistency for high dimensional linear regression under the same irrepresentable condition as Lasso , later extended to generalized linear models ( Huang & Yao , 2018 ) . To relax such conditions , SplitLBI was proposed by Huang et al . ( 2016 ) to learn structural sparsity in linear models under weaker conditions than generalized Lasso , that was successfully applied in medical image analysis ( Sun et al. , 2017 ) and computer vision ( Zhao et al. , 2018 ) . In this paper , it is the first time that SplitLBI is exploited to train highly non-convex neural networks with structural sparsity , together with a global convergence analysis based on the Kurdyka-Łojasiewicz framework Łojasiewicz ( 1963 ) . Contributions . ( 1 ) SplitLBI , as an extension of SGD , is applied to deep learning by exploring both over-parameterized models and structural sparsity in the inverse scale space . ( 2 ) Global convergence of SplitLBI in such a nonconvex optimization is established based on the KurdykaŁojasiewicz framework , that the whole iterative sequence converges to a critical point of the empirical loss function from arbitrary initializations . ( 3 ) Stochastic variants of SplitLBI demonstrate the comparable and even better performance than other training algorithms on ResNet-18 in large scale training such as ImageNet-2012 , among other datasets , together with additional structural sparsity in successful models for interpretability . ( 4 ) Structural sparsity parameters in SplitLBI provide important information about subnetwork architecture with comparable or even better accuracies than dense models before and after retraining -- SplitLBI with early stopping can provide fast “ winning tickets ” without fully training dense , over-parameterized models . 2 METHODOLOGY . The supervised learning task learns a mapping ΦW : X → Y , from input space X to output space Y , with a parameter W such as weights in neural networks , by minimizing certain loss functions on training samples L̂n ( W ) = 1n ∑n i=1 ` ( yi , ΦW ( xi ) ) . For example , a neural network of l-layer is defined as ΦW ( x ) = σl ( W lσl−1 ( W l−1 · · ·σ1 ( W 1x ) ) , where W = { W i } li=1 , σi is the nonlinear activation function of the i-th layer . Differential Inclusion of Inverse Scale Space . Consider the following dynamics , Ẇt κ = −∇W L̄ ( Wt , Γt ) ( 1a ) V̇t = −∇ΓL̄ ( Wt , Γt ) ( 1b ) Vt ∈ ∂Ω̄ ( Γt ) ( 1c ) where V is a sub-gradient of Ω̄ ( Γ ) : = Ωλ ( Γ ) + 12κ‖Γ‖ 2 for some sparsity-enforced regularization Ωλ ( Γ ) = λΩ1 ( Γ ) ( λ ∈ R+ ) such as Lasso or group Lasso penalties for Ω1 ( Γ ) , κ > 0 is a damping parameter such that the solution path is continuous , and the augmented loss function is L̄ ( W , Γ ) = L̂n ( W ) + 1 2ν ‖W − Γ‖22 , ( 2 ) with ν > 0 controlling the gap admitted between W and Γ . Compared to the original loss function L̂n ( W ) , the L̄ ( W , Γ ) additionally adopt the variable splitting strategy , by lifting the original neural network parameter W to ( W , Γ ) with Γ modeling the structural sparsity of W . For simplicity , we assumed L̄ is differentiable with respect to W here , otherwise the gradient in Eq . ( 1a ) is understood as subgradient and the equation becomes an inclusion . The differential inclusion system ( 1 ) , called Split Inverse Scale Space ( SplitISS ) , can be understood as a gradient descent flow of Wt in the proximity of Γt and a mirror descent flow ( Nemirovski & Yudin , 1983 ) of Γt associated with a sparsity enforcement penalty Ω̄ . In mirror descent flow , gradient descent goes on the dual space consisting of sub-gradients Vt , driving the flow in sparse primal space of Γt . For a large enough ν , it reduces to the gradient descent method for Wt . Yet the solution path of Γt exhibits the following property in the separation of scales : starting at the zero , important parameters of large scale will be learned fast , popping up to be nonzeros early , while unimportant parameters of small scale will be learned slowly , appearing to be nonzeros late . In fact , taking Ωλ ( Γ ) = ‖Γ‖1 and κ→∞ for simplicity , Vt as the subgradient of Ω̄t , undergoes a gradient descent flow before reaching the ` ∞-unit box , which implies that Γt = 0 in this stage . The earlier a component in Vt reaches the ` ∞-unit box , the earlier a corresponding component in Γt becomes nonzero and rapidly evolves toward a critical point of L̄ under gradient flow . On the other hand , the Wt follows the gradient descent with a standard ` 2-regularization . Therefore , Wt closely follows dynamics of Γt whose important parameters are selected . Such a property is called as the inverse scale space in applied mathematics ( Burger et al. , 2006 ) and recently was shown to achieve statistical model selection consistency in high dimensional linear regression ( Osher et al. , 2016 ) and general linear models ( Huang & Yao , 2018 ) , with a reduction of bias as κ increases . In this paper , we shall see that the inverse scale space property still holds empirically for the highly nonconvex neural network training via Eq . ( 1 ) . For example , Fig . 1 shows a LeNet trained on MNIST by the discretized dynamics , where important sparse filters are selected in early epochs while the popular SGD returns dense filters . Compared with directly enforcing a penalty function such as ` 1 or ` 2 regularization min W R̂n ( W ) : = L̂n ( W ) +Ωλ ( W ) , λ ∈ R+ . ( 3 ) SplitISS avoids the parameter correlation problem in over-parameterized models . In fact , a necessary and sufficient condition for Lasso or ` 1-type sparse model selection is the incoherence or irrepresentable conditions ( Tropp ( 2004 ) ; Zhao & Yu ( 2006 ) ) that are violated for highly correlated weight parameters , leading to spurious discoveries . In contrast , Huang et al . ( 2018 ) showed that equipped with such a variable splitting where Γ enjoys an orthogonal design where the restricted Hessian of the augmented loss on Γ is orthogonal , the SplitISS can achieve model selection consistency under weaker conditions than generalized Lasso , relaxing the incoherence or irrepresentable conditions when parameters are highly correlated . For weight parameter W , instead of directly being imposed with ` 1-sparsity , it adopts ` 2-regularization in the proximity of the sparse path of Γ that admits simultaneously exploring highly correlated parameters in over-parameterized models and sparsity regularization . Split Linearized Bregman Iterations . SplitISS admits an extremely simple discrete approximation , using the Euler forward discretization of dynamics ( 1 ) : Wk+1 = Wk − καk · ∇W L̄ ( Wk , Γk ) , ( 4a ) Vk+1 = Vk − αk · ∇ΓL̄ ( Wk , Γk ) , ( 4b ) Γk+1 = κ · ProxΩλ ( Vk+1 ) , ( 4c ) where V0 = Γ0 = 0 , W0 can be small random numbers such as Gaussian distribution in neural networks , for some complex networks it can be initialized as common setting.The proximal map in Eq . ( 4c ) that controls the sparsity of Γ is given by ProxΩλ ( V ) = arg min Γ { 1 2 ‖Γ− V ‖22 + Ωλ ( Γ ) } , ( 5 ) We shall call such an iterative procedure as Split Linearized Bregman Iteration ( SplitLBI ) , that was firstly coined in Huang et al . ( 2016 ) as an iterative regularization path for sparse modeling in high dimensional statistics . In the application to neural networks , the loss becomes highly non-convex , the SplitLBI returns a sequence of sparse models from simple to complex ones whose global convergence condition to be shown below , while solving Eq . ( 3 ) at various levels of λ might not be tractable except for over-parameterized models . The sparsity-enforcement penalty used in convolutional neural networks can be chosen as follows . Our sparsity framework aims at regularizing the groups of weight parameters using group Lasso penalty ( Yuan & Lin , 2006 ) , Ω1 ( Γ ) = ∑ g ‖Γg‖2 , where ‖Γg‖2 = √∑|Γg| i=1 ( Γ g i ) 2 and |Γg| is the number of weights in Γg . Thus Eq . ( 4c ) has a closed form solution Γg = κ ·max ( 0 , 1− 1/‖V g‖2 ) V g for the g-th filter . We treat convolutional and fully connected layers in different ways . ( 1 ) For a convolutional layer , Γg = Γg ( cin , cout , size ) denote the convolutional filters where size denotes the kernel size and cin and cout denote the numbers of input channels and output channels , respectively . When we regard each group as each convolutional filter , g = cout ; otherwise for weight sparsity , g can be every element in the filter that reduces to the Lasso . ( 2 ) For a fully connected layer , Γ = Γ ( cin , cout ) where cin and cout denote the numbers of inputs and outputs of the fully connected layer . Each group g corresponds to each element ( i , j ) , and the group Lasso penalty degenerates to the Lasso penalty .

Summary: This paper studies the problem of finding important structural sparsity of over-parameterized neural networks. Compared to the traditional pruning and distillation based compression methods, the author propose a novel optimization based algorithm called SplitLBI, which can efficiently and effectively find the important subnetwork architecture. In addition, the authors provide the convergence guarantee of the proposed method for solving nonconvex optimization problems under certain conditions. Thorough experiments and ablation studies validate the advantage of the proposed method.