Orthogonalising gradients to speedup neural network optimisation

1 INTRODUCTION Neural network layers are made up of several identical , but differently parametrised , components , e.g . filters in a convolutional layer , or heads in a multi-headed attention layer . Layers consist of several components so that they can provide a diverse set of intermediary representations to the next layer , however , there is no constraint or bias , other than the implicit bias from the cost function , to learning different parametrisations . We introduce this diversification bias in the form of orthogonalised gradients and find a resultant speed-up in learning and sometimes improved performance , see fig . 1 . Our novel contributions include this new optimisation method , thorough testing on CIFAR10 and ImageNet , additional testing on a semi-supervised learning method , and experiments to support our hypothesis . In section 2 we detail the method and results to give an understanding of how this method works and its capabilities . Then , in section 3 , we provide experimental justifications and supporting experiments for this method along with finer details of the implementation and limitations . 2 OVERVIEW OF NEW METHOD AND RESULTS . 2.1 RELATED WORKS . Gradient orthogonalisation has been explored in the domain of multi-task learning ( Yu et al. , 2020 ) to keep the different tasks separate and relevant . However in this work we focus on orthogonalisation for improving single task performance . Weight orthogonalisation has been extensively explored with both empirical ( Bansal et al. , 2018 ; Jia et al. , 2017 ) and theoretical ( Jia et al. , 2019 ) justifications . However , modifying the weights during training is unstable , and , in addition , it limits the weights to a tiny subspace . Deep learning is know to work well despite the immense size of the weight space , and as such we do not view this as an advantage . Xie et al . ( 2017 ) obtain improved performance over Stochastic Gradient Descent ( SGD ) via weight orthogonalisation and allows them to train very deep networks , we aim to achieve the same thing while being more flexible with model and optimisation method choice . We do this by orthogonalising the gradients before they are used by an optimisation method rather than modifying the weights themselves . 2.2 ORTHOGONALISING GRADIENTS . Given a neural network , f , with L layers made from components , c , f = ◦Li=1 ( fi ) , ( 1 ) fl ( x ) = [ cl1 ( x ) , cl2 ( x ) , . . . , clNl ( x ) ] , ( 2 ) where ◦ is the composition operator , Nl is the number of components in layer l , cl : RSl−1×Nl−1 → RSl is a parametrised function and cli denotes cl parametrised with θli ∈ RPl giving fl : RSl−1×Nl−1 → RSl×Nl parametrised by θl ∈ RPl×Nl . Let Gl = [ ∇cl1 , ∇cl2 , . . . , ∇clNl ] , ( 3 ) be the Pl ×Nl matrix of the components ’ gradients . Then the nearest orthonormal matrix , i.e . the orthonormal matrix , Ol , that minimises the Frobenius norm of its difference from Gl min Ol ‖Ol −Gl‖ subject to ∀i , j : 〈 Oli , Olj 〉 = δij , where δij is the Kronecker delta function , is the product of the left and right singular vector matricies from the Singular Value Decomposition ( SVD ) of Gl ( Trefethen & Bau III , 1997 ) , Gl = UlΣlV T l , ( 4 ) Ol = UlV T l . ( 5 ) Thus , we can adjust a first-order gradient descent method , such as Stochastic Gradient Descent with Momentum ( SGDM ) ( Polyak , 1964 ) , to make steps where the components are pushed in orthogonal directions , v ( t+1 ) l = γv ( t ) l + ηO ( t ) l , and ( 6 ) θ ( t+1 ) l = θ ( t ) l − v ( t+1 ) l , ( 7 ) where vl is the velocity matrix , t ∈ Z0+ is the time , γ is the momentum decay term , and η is the step size . We call this method Orthogonal Stochastic Gradient Descent with Momentum ( OrthogonalSGDM ) . This modification can clearly be applied to any first-order optimisation algorithm by replacing the gradients with O ( t ) l before the calculation of the next iterate . Code for creating orthogonal optimisers in PyTorch is provided at https : //anonymous . 4open.science/r/Orthogonal-Optimisers . And code for the experiments in this work is provided at https : //anonymous.4open.science/r/ Orthogonalised-Gradients 2.3 RESULTS . 2.3.1 CIFAR-10 . We trained a suite of models on the CIFAR-10 ( Krizhevsky et al. , 2009 ) data set with a mini-batch size of 1024 , learning rate of 10−2 , momentum of 0.9 , and a weight decay of 5 × 10−4 for 100 epochs . We then repeated this using Orthogonal-SGDM instead of SGDM and plot the results in figs . 2 and 3 and table 1 . Orthogonal-SGDM is more efficient and achieves better test accuracy than SGDM for every model we trained on CIFAR-10 without hyper-parameter tuning . The validation curves follow the training curves , figs . 4 and 5 , and have the same patterns , this means that Orthogonal-SGDM exhibits the same generalisation performance as SGDM . More importantly though , we can see that the model 1As described in appendix B.1 2Model same as in He et al . ( 2015 ) 3From https : //pytorch.org/vision/0.9/models.html learns much faster at the beginning of training , as shown by fig . 2 , this means that we do not need as many epochs to get to a well-performing network . This is especially good in light of the large data sets that new models are being trained on , where they are trained for only a few epochs , or even less ( Brown et al. , 2020 ) . For SGDM the performance of the residual networks designed for ImageNet ( Deng et al. , 2009 ) ( 18 , 34 , 50 ) get worse as the models get bigger . The original ResNet authors , He et al . ( 2015 ) , note that unnecessarily large networks may over-fit on a small data set such as CIFAR-10 . However , when trained with Orthogonal-SGDM , these models do not suffer from this over-parametrisation problem and even slightly improve in performance as the models get bigger , in clear contrast to SGDM . This agnosticism to over-parametrisation helps alleviate the need for the practitioner to tune a model ’ s architecture to the task at hand to achieve a reasonable performance . 2.3.2 MATCHING RESNET ’ S PERFORMANCE . Having shown that Orthogonal-SGDM speeds up learning with non-optimised hyper-parameters , we now aim to show that it can achieve state-of-the-art results . To do this we use the same hyperparameters as the original ResNet paper ( He et al. , 2015 ) , which have been painstakingly tuned to benefit SGDM , to train using Orthogonal-SGDM . This also tests the efficacy of Orthogonal-SGDM as a drop-in replacement for SGDM . OrthogonalSGDM gets close to the original results , table 2 , even though the hyper-parameters are perfected for SGDM . It is the authors ’ belief that with enough hyper-parameter tuning SGD or SGDM will be the best optimisation method ; however , this experiment shows that Orthogonal-SGDM is robust to hyper-parameter choice and can easily replace SGDM in existing projects . Unfortunately , the authors do not have the compute-power to extensively hyper-parameter tune a residual network for Orthogonal-SGDM , however , it is exceedingly likely that better results would be achieved by doing so . 2.3.3 IMAGENET . Orthogonal-SGDM also works on a large data set such as ImageNet ( Deng et al. , 2009 ) — fig . 6 . Using a resnet34 , mini-batch size of 1024 , learning rate of 10−2 , momentum of 0.9 , and a weight decay of 5×10−4 , for 100 epochs . SGDM achieves a test accuracy of 61.9 % and a test loss of 1.565 while Orthogonal-SGDM achieves 67.5 % and 1.383 respectively . While these results are a way off the capabilities of the model they still demonstrate a significant speed-up and improvement from using Orthogonal-SGDM , especially at the start of learning , and further reinforces how a dearth of hyper-parameter tuning impedes performance . 2.3.4 BARLOW TWINS . Barlow Twins ( Zbontar et al. , 2021 ) is a semi-supervised method that uses “ the cross-correlation matrix between the outputs of two identical networks fed with distorted versions of a sample ” to avoid collapsing to trivial solutions . While the authors do provide code , we could not replicate their results by running it . To train within our compute limitations we used a mini-batch size of 1024 instead of 2048 however this should not affect the results since “ Barlow Twins does not require large batches ” ( Zbontar et al. , 2021 ) . Additionally , Barlow Twins uses the Layer-wise Adaptive Rate Scaling ( LARS ) algorithm ( You et al. , 2017 ) , which is designed to adjust the learning rate based on the ratio between the magnitudes of the gradients and weights , there should be no significant slowdown , or speed-up , in learning due to the mini-batch size . We do not orthogonalise the gradients for the dense layers ( see section 3.5 ) . Comparing our own runs , we establish that orthogonalising the gradients before the LARS algorithm does speed up learning as shown in fig . 7 , in agreement with previous experiments . This is evidence that orthogonalising gradients is also beneficial for semi-supervised learning and , moreover , that optimisation algorithms other than SGDM can be improved in this way . 2.3.5 IN COMPARISON TO ADAM . We compare our method to the Adam optimiser ( Kingma & Ba , 2014 ) . Adam has found its place as a reliable optimiser that works over a wide variety of hyper-parameter sets , yielding consistent performance with little fine-tuning needed . We see our optimisation method occupying the same space as Adam . Our method outperforms Adam on all but one hyper-parameter set — table 3 . In addition since we can apply our method to any previous optimisation method , we also test Orthogonal-Adam and find that it outperforms Adam too including at high learning rates where Adam suffers from blow-ups . See figs . 13 and 14 for the training plots . 3 DISCUSSION OF PROBLEM AND METHOD . 3.1 NORMALISATION . When we perform SVD on the reshaped gradient tensor , we obtain an orthonormal matrix , since this changes the magnitude of the gradient we look at the effect of this normalisation . Normalised SGDM ( N-SGDM ) ( Nesterov , 2003 ) provides an improvement in non-convex optimisation since it is difficult to get stuck in a local minimum as the step size is not dependent on the gradient magnitude . However , it hinders convergence to a global minimum since there is no way of shortening the step size ; deep learning is highly non-convex and is unlikely to be optimised to a global minimum . Therefore , it stands to reason that normalising the gradient would speed up the optimisation of deep networks . We compare N-SGDM to normalising the gradients per component — i.e . normalising the columns of Gl , eq . ( 3 ) , instead of orthognormalising it — Component Normalised SGDM ( CN-SGDM ) , as well as to SGDM and Orthogonal-SGDM . N-SGDM improves over SGDM , and CN-SGDM improves over N-SGDM except from the oft case where it diverges . Finally , Orthogonal-SGDM obtains the best solutions while remaining stable on all the models .

Summary: This paper proposes to orthogonalize the gradients of each neuron (within a layer) in order to improve training dynamics. They build on the assumption that neurons in a layer tend to learn similar features early in the training, and thus pushing them in orthogonal directions would yield to better results. The method requires calculating SVD on gradients at every iteration, thus unlikely to be practical (no FLOPs/run-time discussion). Training curves on Cifar10 and Imagenet shows faster convergence and better generalization.