Powerpropagation: A sparsity inducing weight reparameterisation

1 Introduction . Deep learning models are emerging as the dominant choice across several domains , from language [ e.g . 1 , 2 ] to vision [ e.g . 3 , 4 ] to RL [ e.g . 5 , 6 ] . One particular characteristic of these architectures is that they perform optimally in the overparameterised regime . In fact , their size seems to be mostly limited by hardware constraints . While this is potentially counter-intuitive , given a classical view on overfitting , the current understanding is that model size tends to have a dual role : It leads to better behaved loss surfaces , making optimisation easy , but also acts as a regulariser . This gives rise to the double descent phenomena , where test error initially behaves as expected , growing with model size due to overfitting , but then decreases again as the model keeps growing , and asymptotes as the model size goes to infinity to a better performance than obtained in the classical regime [ 7 , 8 ] . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . As a consequence , many of the state of the art models tend to be prohibitively large , making them inaccessible to large portions of the research community despite scale still being a driving force in obtaining better performance . To address the high computational cost of inference , a growing body of work has been exploring ways to compress these models . As highlighted by several works [ e.g . 9 , 10 , 11 ] , size is only used as a crutch during the optimisation process , while the final solution requires a fraction of the capacity of the model . A typical approach therefore is to sparsify or prune the neural network after training by eliminating parameters that do not play a vital role in the functional behaviour of the model . Furthermore , there is a growing interest in sparse training [ e.g . 12 , 13 , 14 ] , where the model is regularly pruned or sparsified during training in order to reduce the computational burden . Compressed or sparse representations are not merely useful to reduce computation . Continual learning , for example , focuses on learning algorithms that operate on non-iid data [ e.g . 15 , 16 , 17 ] . In this setup , training proceeds sequentially on a set of tasks . The system is expected to accelerate learning on subsequent tasks as well as using newly acquired knowledge to potentially improve on previous problems , all of this while maintaining low memory and computational footprint . This is difficult due to the well studied problem of catastrophic forgetting , where performance on previous tasks deteriorates rapidly when new knowledge is incorporated . Many approaches to this problem require identifying the underlying set of parameters needed to encode the solution to a task in order to freeze or protect them via some form of regularisation . In such a scenario , given constraints on model size , computation and memory , it is advantageous that each learned task occupies as little capacity as possible . In both scenarios , typically , the share of capacity needed to encode the solution is determined by the learning process itself , with no explicit force to impose frugality . This is in contrast with earlier works on L0 regularisation that explicitly restrict the learning process to result in sparse and compressible representations [ 18 ] . The focus of our work , similar to the L0 literature , is on how to encourage the learning process to be frugal in terms of parameter usage . However instead of achieving this by adding an explicit penalty , we enhance the “ rich get richer ” nature of gradient descent . In particular we propose a new parameterisation that ensures steps taken by gradient descent are proportional to the magnitude of the parameters . In other words , parameters with larger magnitudes are allowed to adapt faster in order to represent the required features to solve the task , while smaller magnitude parameters are restricted , making it more likely that they will be irrelevant in representing the learned solution . 2 Powerpropagation . The desired proportionality of updates to weight magnitudes can be achieved in a surprisingly simple fashion : In the forward pass of a neural networks , raise the parameters of your model ( element-wise ) to the α-th power ( where α > 1 ) while preserving the sign . It is easy to see that due to the chain rule of calculus the magnitude of the parameters ( raised to α− 1 ) will appear in the gradient computation , scaling the usual update . Therefore , small magnitude parameters receive smaller gradient updates , while larger magnitude parameters receive larger updates , leading to the aforementioned “ rich get richer ” phenomenon . This simple intuition leads to the name of our method . More formally , we enforce sparsity through an implicit regulariser that results from reparameterising the model similar to [ 19 , 20 ] . This line of research builds on previous work on matrix factorisation [ e.g . 21 , 22 ] . In [ 19 ] a parameterisation of the form w = v v − u u is used to induce sparsity , where stands for element-wise multiplication and we need both v and u to be able to represent negative values , since the parameters are effectively squared . In our work we rely on a simpler formulation where w = v|v|α−1 , for any arbitrary power α ≥ 1 ( as compared to fixing α to 2 ) , which , since we preserved the sign of v , can represent both negative and positive values . For α = 1 this recovers the standard setting . If we denote by Θ = RM the original parameter space or manifold embedded inRM , our reparameterisation can be understood through an invertible map Ψ , where its inverse Ψ−1 projects θ ∈ Θ into φ ∈ Φ , where Φ is a new parameter space also embedded in RM , i.e . Φ = RM . The map is defined by applying the function Ψ : R → R , Ψ ( x ) = x|x|α−1 element-wise , where by abuse of notation we refer to both the vector and element level function by Ψ . This new parameter space or manifold Φ has a curvature ( or metric ) that depends on the Jacobian of Ψ . Similar constructions have been previously used in optimisation , as for example in the case of the widely known Mirror Descent algorithm [ 23 ] , where the invertible map Ψ is the link function . For deep learning , Natural Neural Networks [ 24 ] rely on a reparameterisation of the model such that in the new parameter space , at least initially , the curvature is close to the identity matrix , making a gradient descent step similar to a second order step . Warp Gradient Descent [ 25 ] relies on a meta-learning framework to learn a nonlinear projection of the parameters with a similar focus of improving efficiency of learning . In contrast , our focus is to construct a parameterisation that leads to an implicit regularisation towards sparse representation , following [ 19 ] , rather then improving convergence . Given the form of our mapping Ψ , in the new parameterisation the original weight θi will be replaced by φi , where Ψ ( φi ) = φi|φi|α−1 = θi and i indexes over the dimensionality of the parameters . Note that we apply this transformation only to the weights of a neural network , leaving other parameters untouched . Given the reparameterised loss L ( · , Ψ ( φ ) ) , the gradient wrt . to φ becomes ∂L ( · , Ψ ( φ ) ) ∂φ = ∂L ∂Ψ ( φ ) ∂Ψ ( φ ) ∂φ = ∂L ∂Ψ ( φ ) diag ( α|φ|◦α−1 ) . ( 1 ) Note that diag indicates a diagonal matrix , and |φ|◦α−1 indicates raising element-wise the entries of vector |φ| to the power α − 1 . ∂L∂Ψ ( φ ) is the derivative wrt . to the original weight θ = φ|φ| ◦α−1 which is the gradient in the original parameterisation of the model . This is additionally multiplied ( element-wise ) by the factor α|φi|α−1 , which will scale the step taken proportionally to the magnitude of each entry . Finally , for clarity , this update is different from simply scaling the gradients in the original parameterisation by the magnitude of the parameter ( raised at α − 1 ) , since the update is applied to φ not θ , and is scaled by φ . The update rule ( 1 ) has the following properties : ( i ) 0 is a critical point for the dynamics of any weight φi , if α > 1 . This is easy to see as ∂L ∂φi = 0 whenever φi = 0 due to the α|φi|α−1 factor . ( ii ) In addition , 0 is surrounded by a plateau and hence weights are less likely to change sign ( gradients become vanishingly small in the neighbourhood of 0 due to the scaling ) . This should negatively affect initialisations that allow for both negative and positive values , but it might have bigger implications for biases . ( iii ) This update is naturally obtained by the Backpropagation algorithm . This comes from the fact that Backpropagation implies applying the chain-rule from the output towards the variable of interest , and our reparameterisation simply adds another composition ( step ) in the chain before the variable of interest . At this point the perceptive reader might be concerned about the effect of equation ( 1 ) on established practises in the training of deep neural networks . Firstly , an important aspect of reliable training is the initialisation ( [ e.g . 26 , 27 , 28 ] ) or even normalisation layers such as batch-norm [ 29 ] or layernorm [ 30 ] . We argue that our reparameterisation preserves all properties of typical initialisation schemes as it does not change the function . Specifically , let θi ∼ p ( θ ) where p ( θ ) is any distribution of choice . Then our reparameterisation involves initialising φi ← sign ( θi ) · α √ |θi| , ensuring the neural network and all intermediary layers are functionally the same . This implies that hidden activations will have similar variance and mean as in the original parameterisation , which is what initialisation and normalisation focus on . Secondly , one valid question is the impact of modern optimisation algorithms ( [ e.g . 31 , 32 , 33 ] ) on our reparameterisation . These approaches correct the gradient step by some approximation of the curvature , typically given by the square root of a running average of squared gradients . This quantity will be proportional ( at least approximately ) to diag ( α|φ|◦α−1 ) 1 . This reflects the fact that our projection relies on making the space more curved and implicitly optimisation harder , which is what these optimisation algorithms aim to fix . Therefore , a naive use with Powerprop . would result in a reduction of the “ rich get richer ” effect . On the other hand , avoiding such optimisers completely can considerably harm convergence and performance . The reason for this is that they do not only 1To see this assume the weights do not change from iteration to iteration . Then each gradient is scaled by the same value diag ( α|φ|◦2 ( α−1 ) ) which factors out in the summation of gradients squared , hence the correction from the optimiser will undo this scaling . In practice φ changes over time , though slowly , hence approximately this will still hold . correct for the curvature induced by our reparameterisation , but also the intrinsic curvature of the problem being solved . Removing the second effect can make optimisation very difficult . We provide empirical evidence of this effect in the Appendix . To mitigate this issue and make Powerprop . straightforward to use in any setting , we take inspiration from the target propagation literature [ 34 , 35 , 36 , 37 ] which proposes an alternative way of thinking about the Backpropagation algorithm . In our case , we pretend that the exponentiated parameters are the de-facto parameters of the model and compute an update wrt . to them using our optimiser of choice . The updated exponentiated parameters are then seen as a virtual target , and we take a gradient descent step on φ towards these virtual targets . This will result in a descent step , which , while it relies on modern optimisers to correct for the curvature of the problem , does not correct for the curvature introduced by our parameterisation . Namely if we denote optim : RM → RM as the function that implements the correction to the gradient done by some modern optimiser , our update becomes ∆φ = optim ( ∂L ∂Ψ ( φ ) ) diag ( α|φ|◦α−1 ) . The proof that our update is indeed correct follows the typical steps taken in the target propagation literature . From a first order Taylor expansion of L ( φ−η∆φ ) , we have that in order for ∆φ to reduce the loss , the following needs to hold : 〈∆φ , ∂L∂φ 〉 > 0 . But we know that 〈 optim ( ∂L ∂Ψ ( φ ) ) , ∂L∂Ψ ( φ ) 〉 > 0 as this was a valid step on Ψ ( φ ) . Because diag ( α|φ|◦α−1 ) is positive definite ( diagonal matrix with all entries positive ) , we can multiply it on both sides , proving that 〈∆φ , ∂L∂φ 〉 > 0 . We provide more details in the Appendix . We will rely on this formulation in our empirical evaluation .

Summary: The paper proposes a method to parameterize weights such that the resulting model is inherently sparse. This is achieved by dynamically updating the parameters proportional to their magnitude i.e., the higher magnitude weights receive larger updates while the low magnitude weights remain unaffected. The idea is novel and useful for training sparse networks geared towards applications in resource-constrained environments. Additionally, as elaborated in the paper, it is also useful in the cases of continual learning where a single model is used to learn multiple tasks. The claims are supported by experimental evaluation.