Sparse Flows: Pruning Continuous-depth Models

1 Introduction The continuous analog of normalizing flows ( CNFs ) ( Chen et al. , 2018 ) efficiently ( Grathwohl et al. , 2019 ) maps a latent space to data by ordinary differential equations ( ODEs ) , relaxing the strong constraints over discrete normalizing flows ( Dinh et al. , 2016 , Durkan et al. , 2019 , Huang et al. , 2020 , Kingma and Dhariwal , 2018 , Papamakarios et al. , 2017 , Rezende and Mohamed , 2015 ) . CNFs enable learning flexible models by unconstrained neural networks . While recent works investigated ways to improve CNFs ’ efficiency ( Finlay et al. , 2020 , Grathwohl et al. , 2019 , Li et al. , 2020 ) , regularize the flows ( Onken et al. , 2020 , Yang and Karniadakis , 2020 ) , or solving their shortcomings such as crossing trajectories ( Dupont et al. , 2019 , Massaroli et al. , 2020 ) , less is understood about their inner dynamics during and post training . In this paper , we set out to use standard pruning algorithms to investigate generalization properties of sparse neural ODEs and continuous normalizing flows . In particular , we investigate how the inner dynamics and the modeling performance of a continuous flow varies if we methodologically prune its neural network architecture . Reducing unnecessary weights of a neural network ( pruning ) ( Han et al. , 2015b , Hassibi and Stork , 1993 , LeCun et al. , 1990 , Li et al. , 2016 ) without loss of accuracy results in smaller network size ( Hinton et al. , 2015 , Liebenwein et al. , 2021a , b ) , computational efficiency ( Luo et al. , 2017 , Molchanov et al. , 2019 , Yang et al. , 2017 ) , faster inference ( Frankle and Carbin , 2019 ) , and enhanced interpretability ( Baykal et al. , 2019a , b , Lechner et al. , 2020a , Liebenwein et al. , 2020 ) . Here , our main objective is to better understand CNFs ’ dynamics in density estimation tasks as we increase network sparsity and to show that pruning can improve generalization in neural ODEs . ∗denotes authors with equal contributions . Code : https : //github.com/lucaslie/torchprune 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Pruning improves generalization in neural ODEs . Our results consistently suggest that a certain ratio of pruning of fully connected neural ODEs leads to lower empirical risk in density estimation tasks , thus obtaining better generalization . We validate this observation on a large series of experiments with increasing dimensionality . See an example here in Figure 2 . Pruning flattens the loss surface of neural ODEs . Additionally , we conduct a Hessian-based empirical investigation on the objective function of the flows-under-test in density estimation tasks to better understand why pruning results in better generalization . We find that for Neural ODEs , pruning decreases the value of the Hessian ’ s eigenvalues , and as a result , flattens the loss which leads to better generalization , c.f. , Keskar et al . ( 2017 ) ( Figure 3 ) . Pruning helps avoiding mode-collapse in generative modeling . In a series of multi-modal density estimation tasks , we observe that densely connected CNFs often get stuck in a sharp local minimum ( See Figure 3 ) and as a result , can not properly distinguish different modes of data . This phenomena is known as mode-collapse . Once we sparsify the flows , the quality of the density estimation task increases significantly and consequently mode-collapse does not occur . Pruning finds minimal and efficient neural ODE representations . Our framework finds highly optimized and efficient neural ODE architectures via pruning . In many instances , we can reduce the parameter count by 70-98 % ( 6x-50x compression rate ) . Notably , one can not directly train such sparse and efficient continuous-depth models from scratch . 2 Background . In this section , we describe the necessary background to construct our framework . We show how to perform generative modeling by continuous depth models using the change of variables formula . Generative modeling via change of variables . The change of variables formula uses an invertible mapping f : RD → RD , to wrap a normalized base distribution pz ( z ) , to specify a more complex distribution . In particular , given z ∼ pz ( z ) , a random variable , the log density for function f ( z ) = x can be computed by Dinh et al . ( 2015 ) : log px ( x ) = log pz ( z ) − log det ∣∣∣∣∂f ( z ) ∂z ∣∣∣∣ , ( 1 ) where ∂f ( z ) ∂z is the Jacobian of f . While theoretically Eq . 1 demonstrates a simple way to finding the log density , from a practical standpoint computation of the Jacobian determinant has a time complexity of O ( D3 ) . Restricting network architectures can make its computation more tractable . Examples include designing normalizing flows ( Berg et al. , 2018 , Papamakarios et al. , 2017 , Rezende and Mohamed , 2015 ) , autoregressive transformations ( Durkan et al. , 2019 , Jaini et al. , 2019 , Kingma et al. , 2016 , Müller et al. , 2019 , Oliva et al. , 2018a , Wehenkel and Louppe , 2019 ) , partitioned transformations ( Dinh et al. , 2016 , Kingma and Dhariwal , 2018 ) , universal flows ( Kong and Chaudhuri , 2020 , Teshima et al. , 2020 ) , and the use of optimal transport theorem ( Huang et al. , 2020 ) . Alternative to these discrete transformation algorithms , one can construct a generative model similar to ( 1 ) , and declare f by a continuous-time dynamics ( Chen et al. , 2018 , Grathwohl et al. , 2019 , Lechner et al. , 2020b ) . Given a sample from the base distribution , one can parametrize an ordinary differential equations ( ODEs ) by a function f ( z ( t ) , t , θ ) , and solve the ODE to obtain the observable data . When f is a neural network , the system is called a neural ODE ( Chen et al. , 2018 ) . Neural ODEs . More formally , a neural ODE is defined by finding the solution to the initial value problem ( IVP ) : ∂z ( t ) ∂t = f ( z ( t ) , t , θ ) , z ( t0 ) = z0 , with z0 ∼ pz0 ( z0 ) , to get z ( tn ) the desired output observations at a terminal integration step n ( Chen et al. , 2018 ) .2 Continuous Normalizing Flows . If z ( tn ) is set to our observable data , given samples from the base distribution z0 ∼ pz0 ( z0 ) , the neural ODE described above forms a continuous normalizing flow ( CNF ) . CNFs modify the change in log density by the left hand-side differential equation and as a result the total change in log-density by the right hand-side equation ( Chen et al. , 2018 , Grathwohl et al. , 2019 ) : ∂ log p ( z ( t ) ) ∂t = −Tr ( ∂f ∂z ( t ) ) , log p ( z ( tn ) ) = log p ( z ( t0 ) ) − ∫ t1 t0 Tr ( ∂f ∂z ( t ) ) dt . ( 2 ) The system of two differential equations ( the neural ODE ( ∂z ( t ) ∂t = f ( z ( t ) , t , θ ) ) and ( 2 ) can then be solved by automatic differentiation algorithms such as backpropagation through time ( Hasani et al. , 2021 , Rumelhart et al. , 1986 ) or the adjoint sensitivity method ( Chen et al. , 2018 , Pontryagin , 2018 ) . Computation of Tr ( ∂f ∂z ( t ) ) costs O ( D2 ) . A method called the Free-form Jacobian of Reversible Dynamics ( FFJORD ) ( Grathwohl et al. , 2019 ) improved the cost to O ( D ) by using the Hutchinson ’ s trace estimator ( Adams et al. , 2018 , Hutchinson , 1989 ) . Thus , the trace of the Jacobian can be estimated by : Tr ( ∂f ∂z ( t ) ) = Ep ( ε ) [ εT ∂f∂z ( t ) ε ] , where p ( ε ) is typically set to a Gaussian or Rademacher distribution ( Grathwohl et al. , 2019 ) . Throughout the paper , we investigate the properties of FFJORD CNFs by pruning their neural network architectures . 3 Pruning Neural ODEs . We enable sparsity in neural ODEs and CNFs by removing , i.e. , pruning , redundant weights from the underlying neural network architecture during training . Pruning can tremendously improve the parameter efficiency of neural networks across numerous tasks , such as computer vision ( Liebenwein et al. , 2020 ) and natural language processing ( Maalouf et al. , 2021 ) . 3.1 A General Framework for Training Sparse Flows . Our approach to training Sparse Flows is inspired by iterative learning rate rewinding , a recently proposed and broadly adopted pruning framework as used by Liebenwein et al . ( 2020 ) , Renda et al . ( 2020 ) among others . In short , our pruning framework proceeds by first training an unpruned , i.e. , dense network to obtain a warm initialization for pruning . Subsequently , we proceed by iteratively pruning and retraining the network until we either obtain the desired level of sparsity , i.e. , prune ratio or when the loss for a pre-specified hold-out dataset ( validation loss ) starts to deteriorate ( early stopping ) . We note that our framework is readily applicable to any continuous-depth model and not restricted to FFJORD-like models . Moreover , we can account for various types of pruning , i.e. , unstructured pruning of weights and structured pruning of neurons or filters . An overview of the framework is provided in Algorithm 1 and we provide more details below . 3.2 From Dense to Sparse Flows . TRAIN a dense flow for a warm initialization . To initiate the training process , we first train a densely-connected network to obtain a warm initialization ( Line 2 of Algorithm 1 ) . We use Adam with a fixed step learning decay schedule and weight decay in some instances . Based on the warm initialization , we start pruning the network . PRUNE for Sparse Flow . For the prune step ( Line 5 of Algorithm 1 ) we either consider unstructured or structured pruning , i.e. , weight or neuron/filter pruning , respectively . At a fundamental level , 2One can design a more expressive representation ( Hasani et al. , 2020 , Vorbach et al. , 2021 ) of continuousdepth models by using the second-order approximation of the neural ODE formulation ( Hasani et al. , 2021 ) . This representation might give rise to a better neural flows which will be the focus of our continued effort . Algorithm 1 SPARSEFLOW ( f , Φtrain , PR , e ) Input : f : neural ODE model with parameter set θ ; Φtrain : hyper-parameters for training ; PR : relative prune ratio ; e : number of training epochs per prune-cycle . Output : f ( · ; θ̂ ) : Sparse Flow ; m : sparse connection pattern . 1 : θ0 ← RANDOMINIT ( ) 2 : θ ← TRAIN ( θ0 , Φtrain , e ) . Initial training stage with dense neural ODE ( “ warm start ” ) . 3 : m← 1|θ0| . Initialize binary mask indicating neural connection pattern . 4 : while validation loss of Sparse Flow decreases do 5 : m← PRUNE ( m θ , PR ) . Prune PR % of the remaining parameters and update mask . 6 : θ ← TRAIN ( m θ , Φtrain , e ) . Restart training with updated connection pattern . 7 : end while 8 : θ̂ ← m θ , and return f ( · ; θ̂ ) , m unstructured pruning aims at inducing sparsity into the parameters of the flow while structured pruning enables reducing the dimensionality of each flow layer . For unstructured pruning , we use magnitude pruning ( Han et al. , 2015a ) ) , where we prune weights across all layers ( global ) with magnitudes below a pre-defined threshold . For structured pruning , we use the ` 1-norm of the weights associated with the neuron/filter and prune the structures with lowest norm for constant per-layer prune ratio ( local ) as proposed by Li et al . ( 2016 ) . See Table 1 for an overview . TRAIN the Sparse Flow . Following the pruning step , we re-initiate the training with the new sparsity pattern and the unpruned weights ( Line 6 of Algorithm 1 ) . Note that we do not change the training hyperparameters between the different stages of training . ITERATE for increased sparsity and performance . Naturally , we can iteratively repeat the PRUNE and TRAIN step to further sparsify the flow ( Lines 4-7 of Algorithm 1 ) . Moreover , the resulting sparsity-performance trade-off is affected by the total number of iterations , the relative prune ratio PR per iteration , and the amount of training between PRUNE steps . We find that a good trade-off is to keep the amount of training constant across all iterations and tune it such that the initial , dense flow is essentially trained to ( or close to ) convergence . Depending on the difficulty of the task and the available compute resources we can then adapt the per-iteration prune ratio PR . Note that the overall relative sparsity after n iterations is given by ( 1− PR ) n. Detailed hyperparameters and more explanations for each experiment are provided in the supplementary material .

Summary: **Summary**  The paper describes the effect of using pruning during neural ODEs training. Namely, the authors propose to train neural ODEs iteratively by alternating training (fine-tuning) and pruning steps. They demonstrate the benefits of applying the proposed training procedure to the density estimation tasks. Also, the paper considers a simple classification task to compare the robustness of decision boundaries of pruned ODE-based flows.