GradMax: Growing Neural Networks using Gradient Information

1 INTRODUCTION . The architecture of deep learning models influences a model ’ s inductive biases and has been shown to have a crucial effect on both the training speed and generalization ( dÁscoli et al. , 2019 ; Neyshabur , 2020 ) . Searching for the best architecture for a given task is an active research area with diverse approaches , including neural architecture search ( NAS ) ( Elsken et al. , 2019 ) , pruning ( Liu et al. , 2018 ) , and evolutionary algorithms ( Stanley & Miikkulainen , 2002 ) . Most of these approaches are costly , as they require large search spaces or large architectures to start with . In this work we consider an alternative approach : Can we start with a small network and learn an efficient architecture without ever needing a large network or exceeding the size of our final architecture ? The idea of incrementally increasing the size of a model has been used in many settings such as boosting ( Friedman , 2001 ) , continual learning ( Rusu et al. , 2016 ) , architecture search ( Elsken et al. , 2017 ) , optimization ( Fukumizu & Amari , 2000 ) , and reinforcement learning ( Berner et al. , 2019 ) . Despite the wide range of applications of growing neural networks , the initialization of newly grown neurons is rarely studied . Existing work on growing neural networks either adds new neurons randomly ( Chen et al. , 2016 ; Berner et al. , 2019 ) or chooses them with the aim of decreasing the training loss ( Bengio et al. , 2006 ; Liu et al. , 2019 ; Wu et al. , 2020a ) . In this work we take a different approach . Instead of improving the training objective immediately through growing , we focus on improving the subsequent training dynamics . As we will show , this has a longer-lasting effect than the greedy approach of improving loss during growing . Our main contribution is a new growing method that maximizes the gradient norm of newly added neurons , which we call GradMax . We show that this gradient maximization problem can be solved directly using the singular value decomposition ( SVD ) and initializing the new neurons using the top k singular vectors . We show that the increased gradient norm persists during future steps , which in return yields faster training . 2 GROWING NEURAL NETWORKS . Neural network learning is often formalized as the following optimization problem : argmin f∈S E ( x , y ) ∼D [ L ( f ( x ) , y ) ] , ( 1 ) where S is a set of neural networks , L is a loss function andD is a data set consisting of inputs x and outputs y ) . Most often the set S is constrained to a single architecture ( e.g. , ResNet-18 ( Zagoruyko & Komodakis , 2016 ) ) and only the parameters of the model are learned . However , hand-crafted Before ℓℓ-1 ℓ+1 Before ℓℓ-1 ℓ+1 Auxiliary Connections Adding new layers with identitiy in between Before ℓℓ-1 ℓ+1 Auxiliary Connections New Layers ` +1 weights for the new neuron . GradMax sets incoming weights to zero ( dashed lines ) in order to keep the output unchanged , and initializes outgoing weights using SVD ( equation 11 ) . This maximizes the gradients on the incoming weights with the aim of accelerating training . architectures are often suboptimal and various approaches aim to optimize the architecture itself together with its parameters . Methods such as random search ( Bergstra & Bengio , 2012 ) , NAS , and pruning search a larger space that contains a variety of architectures . Growing neural networks has the same goal : It learns the architecture and weights jointly . It aims to achieve this by incrementally increasing the size of the model . Since this approach never requires training a large model from scratch it needs less memory and compute than pruning methods ( Yuan et al. , 2021 ) . Reducing the cost of finding efficient architectures is important given the ever-growing size of architectures and the accompanying increase in environmental footprint ( Thompson et al. , 2021 ) . Growing neural networks : When , where and how ? Algorithms for growing neural networks often start training with a small seed architecture . Then over the course of the training new neurons are added to the seed architecture , either increasing the width of the existing layers or creating new layers . Algorithms for growing neural networks should address the following questions : 1 . When to add new neurons ? For instance , some methods require the training loss to plateau before growing , whereas others grow using a predefined schedule . 2 . Where to add new capacity ? We can add new neurons to the existing layers or create new layers among the existing ones . 3 . How to initialize the new capacity ? In this work we mainly focus on the question of how and introduce a new initialization method for the new neurons . Our approach can also be used to guide when and where to grow new neurons . However , in order to make our comparison with other initialization methods fair we keep the growing schedule ( where and when ) fixed . How to grow new neurons ? Suppose that during training we would like to grow k new neurons at layer ` as depicted in Figure 1 . After the growth , new neurons are appended to the existing weight matrices as follows : W+ ` = [ W ` W new ` ] ( 2 ) W+ ` +1 = [ W ` +1 W new ` +1 ] . ( 3 ) A desirable property of growing algorithms is to preserve the information that the network has learned . Therefore the initialization of new neurons should ensure our neural network has the same outputs before and after growth . We describe two main approaches to growing : splitting and adding . Splitting ( Chen et al. , 2016 ; Liu et al. , 2019 ) duplicates existing neurons and adjusts outgoing weights so that the output in the next layer is unchanged . However , splitting has some limitations : ( 1 ) It creates redundant neurons and small changes required to break symmetry cause changes in the output distribution ; ( 2 ) It can ’ t be used for growing new layers as it requires existing neurons to begin with . Another line of work focuses on adding new neurons while ensuring that the output does not change by setting either W new ` or W new ` +1 to zero . This approach doesn ’ t have the limitations mentioned above and often results in better performance ( Wu et al. , 2020a ) . Below we summarize two methods that use this approach : Random Berner et al . ( 2019 ) initialize new neurons randomly when growing new layers . Firefly Wu et al . ( 2020a ) combine splitting and adding random neurons . Random neurons are trained for a few steps to reduce loss directly and the most promising neurons are selected for growing . Previous work focused either on solely keeping the output of the neural network the same ( Random ) or on immediately reducing the training objective ( Firefly ) . In this work we take a different approach and focus on improving the training dynamics . We initialize the weights so that the network ’ s output remains unchanged while trying to improve the training dynamics by maximizing the gradient norm . The hypothesis is that this accelerates training and leads to a better model in the long run . 3 GRADMAX . This section describes our main contribution : GradMax , a method that maximizes the gradient norm of the new weights . We constrain the norm of the new weights to avoid the trivial solution where the weight norm tends to infinity . We also require that the network output is unchanged when the neuron is added . Note that , due the latter , gradients of the existing weights are unchanged during growth and therefore we maximize the gradients on the new weights introduced . argmax W new ` , W new ` +1 ∥∥∥∥ED [ ∂L∂W new ` ] ∥∥∥∥2 F + ∥∥∥∥ED [ ∂L∂W new ` +1 ] ∥∥∥∥2 F s.t . { ‖W new ` ‖F , ∥∥W new ` +1∥∥F ≤ c W new ` +1h new ` = 0 ( 4 ) Motivation : Large gradients lead to large objective decrease Consider running gradient descent on a function that is differentiable with a β-Lipschitz gradient . A classical result for this class of functions is that the decrease in objective function after one step of gradient descent with step-size 1/β is upper bounded by L ( W ` ) − β2 ‖∇L ( W ` ) ‖ 2 ( Nesterov , 2003 ) . This upper-bound decreases as the norm of the gradient increases . This observation implies that the larger the gradient , the larger the expected decrease ( assuming a constant Lipschitz constant ) . Hence , interleaving gradient maximizing steps within training will lead to an improved decrease in later iterations . We use this observation in the context of growing neural networks and propose a method that interleaves the normal training process with steps that maximize the gradient norm by adding new units , which we now describe in detail . W ` L ( W ` ) − β2 ‖∇L ( W ` ) ‖ 2 L ( W ` ) GradMax The general maximization problem in eq . ( 4 ) is non-trivial to solve . However , we show that with some simplifying assumptions we can find an approximate solution using a singular value decomposition ( SVD ) . We begin by introducing the notation . Consider 3 fully connected layers denoted with indices ` − 1 , ` and ` + 1 and the following recursive definition : z ` = W ` h ` −1 ( 5 ) h ` = f ( z ` ) , ( 6 ) where subscripts denote layer-indices . Let M ` denote the number of units in layer ` , and f the activation function . The vectors z ` , h ` ∈ RM ` are pre-activations and activations , with h0 = x an input sampled from the dataset D. The entries of W ` ∈ RM ` −1×M ` are the parameters of the layer . When growing k neurons the weight matrices W ` and W ` +1 are replaced by W+ ` and W + ` +1 as defined in eqs . ( 2 ) and ( 3 ) . We denote the pre-activations and activations of the new neurons with znew ` and h new ` . The gradients of the new weights can be derived using the chain rule : ∂L ∂W new ` = ( f ′ ( znew ` ) W new , > ` +1 ∂L ∂z ` +1 ) h > ` −1 ( 7 ) ∂L ∂W new ` +1 = ∂L ∂z ` +1 hnew , > ` . ( 8 ) The simplifying assumptions that we will now make are that W new ` = 0 and that f ( 0 ) = 0 with gradient f ′ ( 0 ) = 1 . Note that this guarantees that W new ` +1h new ` = 0 , independent of the training data . Moreover , it simplifies the gradients to ∂L ∂W new ` = W new , > ` +1 ∂L ∂z ` +1 h > ` −1 ( 9 ) ∂L ∂W new ` +1 = 0 , ( 10 ) which reduces our problem to argmax W new ` +1 ∥∥∥∥W new , > ` +1 ED [ ∂L∂z ` +1h > ` −1 ] ∥∥∥∥2 F , s.t . ∥∥W new ` +1∥∥F ≤ c . ( 11 ) The solution to this maximization problem is letting the columns of W new ` +1 be the top k left-singular vectors of the matrix ED [ ∂L ∂z ` +1 h > ` −1 ] and scaling them by c‖ ( σ1 , ... , σk ) ‖ ( where σi is the i-th largest singular value ) . In order to make a fair comparison between different methods we scale each initialization such that their norm is equal to the mean norm of the existing neurons . Note that it is feasible to also use the singular values to guide where and when to grow . For example , neurons could be added when the singular values meet a certain threshold , and layers to grow could be chosen depending on which have the largest singular values . We leave this for future work . Non-linearities and normalization Note that our assumptions that f ( 0 ) = 0 and f ′ ( 0 ) = 1 apply to common activation functions such as rectified linear units ( if zero is included in the linear part ) and the hyperbolic tangent . Other functions could be adapted to fit this definition as well ( e.g. , by shifting the output of the logistic function by − 12 ) . In fact our derivation holds for f ′ ( 0 ) = a , a 6= 0 . This means that batch normalization can also be used , since it meets these assumptions . However , care must be taken with activation functions ( or their derivatives ) that are unstable around zero . For example , batch normalization scales the gradients by 1ε when all inputs are zero , which can lead to large jumps in parameter space ( Lewkowycz et al. , 2020 , and Appendix C ) . GradMax-Optimized ( GradMaxOpt ) One can also directly optimize eq . ( 11 ) using an iterative method such as projected gradient descent . Our experiments show that this approach generally does not work as well as using the SVD , highlighting the benefit of having a closed-form solution . However , if the outgoing weights are set to zero ( W new ` +1 = 0 ) instead of the incoming weights then the solution can no longer be found using SVD and direct optimization of eq . ( 4 ) could provide a solution . This could be preferable in some situations since it removes the constraints on the activation function . Moreover , it can avoid the unstable behaviour of functions such as batch normalization . Full gradient estimation Solving eq . ( 11 ) requires calculating the gradient ∂L∂z ` +1h > ` −1 over the full dataset . This can be expensive so in practical applications we propose using a large minibatch instead . If the minibatch gradient is sufficiently close to the full gradient then the top singular vectors are likely to be close as well ( Stewart , 1998 ) . We validate this experimentally in Section 4.2 .

Summary: The paper proposes a method to grow neural networks by adding neurons in a specific way. While requiring that the forward pass stays unchanged, the authors want to initialize the new weights such that the gradient wrt to these weights is maximized. They show small but consistent improvements compared to the baselines although the proposed method seems considerably more complex than baselines, as far as I can see.