Initialization and Regularization of Factorized Neural Layers

1 INTRODUCTION . Most neural network layers consist of matrix-parameterized functions followed by simple operations such as activation or normalization . These layers are the main sources of model expressivity , but also the biggest contributors to computation and memory cost ; thus modifying these layers to improve computational performance while maintaining performance is highly desirable . We study the approach of factorizing layers , i.e . reparameterizing them so that their weights are defined as products of two or more matrices . When these are smaller than the original matrix , the resulting networks are more efficient for both training and inference ( Denil et al. , 2013 ; Moczulski et al. , 2015 ; Ioannou et al. , 2016 ; Tai et al. , 2016 ) , resulting in model compression . On the other hand , if training cost is not a concern , one can increase the width or depth of the factors to over-parameterize models ( Guo et al. , 2020 ; Cao et al. , 2020 ) , improving learning without increasing inference-time cost . This can be seen as a simple , teacher-free form of knowledge distillation . Factorized layers also arise implicitly , such as in the case of multi-head attention ( MHA ) ( Vaswani et al. , 2017 ) . Despite such appealing properties , networks with factorized neural layers are non-trivial to train from scratch , requiring custom initialization , regularization , and optimization schemes . In this paper we focus on initialization , regularization , and how they interact with gradient-based optimization of factorized layers . We first study spectral initialization ( SI ) , which initializes factors using singular value decomposition ( SVD ) so that their product approximates the target un-factorized matrix . Then , we study Frobenius decay ( FD ) , which regularizes the product of matrices in a factorized layer rather than its individual terms . Both are motivated by matching the training regimen of the analogous un-factorized optimization . Note that SI has been previously considered in the context of model compression , albeit usually for factorizing pre-trained models ( Nakkiran et al. , 2015 ; Yaguchi et al. , 2019 ; Yang et al. , 2020 ) rather than low-rank initialization for end-to-end training ; FD has been used in model compression using an uncompressed teacher ( Idelbayev & Carreira-Perpiñán , 2020 ) . We formalize and study the justifications of SI and FD from both the classical perspective— matching the un-factorized objective and scaling—and in the presence of BatchNorm ( Ioffe & Szegedy , 2015 ) , where this does not apply . Extending recent studies of weight-decay ( Zhang et al. , 2019 ) , we argue that the effective step-size at spectral initialization is controlled by the factorization ’ s Frobenius norm and show convincing evidence that weight-decay penalizes the nuclear norm . We then turn to applications , starting with low-memory training , which is dominated by unstructured sparsity methods—i.e . guessing “ lottery tickets ” ( Frankle & Carbin , 2019 ) —with a prevailing trend of viewing low-rank methods as uncompetitive for compression ( Blalock et al. , 2020 ; Zhang et al. , 2020 ; Idelbayev & Carreira-Perpiñán , 2020 ; Su et al. , 2020 ) . Here we show that , without tuning , factorized neural layers outperform all structured sparsity methods on ResNet architectures ( He et al. , 2016 ) , despite lagging on VGG ( Simonyan & Zisserman , 2015 ) . Through ablations , we show that this result is due to using both SI and FD on the factorized layers . We further compare to a recent evaluation of tensor-decomposition approaches for compressed WideResNet training ( Zagoruyko & Komodakis , 2016 ; Gray et al. , 2019 ) , showing that ( a ) low-rank approaches with SI and FD can outperform them and ( b ) they are themselves helped by tensor-variants of SI and FD . We also study a fledgling subfield we term overcomplete knowledge distillation ( Arora et al. , 2018 ; Guo et al. , 2020 ; Cao et al. , 2020 ) in which model weights are over-parameterized as overcomplete factorizations ; after training the factors are multiplied to obtain a compact representation of the same network . We show that FD leads to significant improvements , e.g . we outperform ResNet110 with an overcomplete ResNet56 that takes 1.5x less time to train and has 2x fewer parameters at test-time . Finally , we study Transformer architectures , starting by showing that FD improves translation performance when applied to MHA . We also show that SI is critical for low-rank training of the model ’ s linear layers . In an application to BERT pre-training ( Devlin et al. , 2019 ) , we construct a Frobeniusregularized variant—FLAMBé—of the LAMB method ( You et al. , 2020 ) , and show that , much like for transformers , it improves performance both for full-rank and low-rank MHA layers . To summarize , our main contributions are ( 1 ) motivating the study of training factorized layers via both the usual setting ( model compression ) and recent applications ( distillation , multi-head attention ) , ( 2 ) justifying the use of SI and FD mathematically and experimentally , and ( 3 ) demonstrating their effectiveness by providing strong baselines and novel advances in many settings . Code to reproduce our results is available here : https : //github.com/microsoft/fnl_paper . 1.1 RELATED WORK . We are not the first study gradient descent on factorized layers ; in particular , deep linear nets are well-studied in theory ( Saxe et al. , 2014 ; Gunasekar et al. , 2019 ) . Apart from Bernacchia et al . ( 2018 ) these largely examine existing algorithms , although Arora et al . ( 2018 ) do effectively propose overcomplete knowledge distillation . Rather than the descent method , we focus on the initialization and regularization . For the former , several papers use SI after training ( Nakkiran et al. , 2015 ; Yaguchi et al. , 2019 ; Yang et al. , 2020 ) , while Ioannou et al . ( 2016 ) argue for initializing factors as though they were single layers , which we find inferior to SI in some cases . Outside deep learning , spectral methods have also been shown to yield better initializations for certain matrix and tensor problems ( Keshavan et al. , 2010 ; Chi et al. , 2019 ; Cai et al. , 2019 ) . For regularization , Gray et al . ( 2019 ) suggest compression-rate scaling ( CRS ) , which scales weight-decay using the reduction in parameter count ; this is justified via the usual Bayesian understanding of ` 2-regularization ( Murphy , 2012 ) . However , we find that FD is superior to any tuning of regular weight-decay , which subsumes CRS . Our own analysis is based on recent work suggesting that the function of weight-decay is to aid optimization by preventing the effective step-size from becoming too small ( Zhang et al. , 2019 ) . 2 PRELIMINARIES ON FACTORIZED NEURAL LAYERS . In the training phase of ( self- ) supervised ML , we often solve optimization problems of the form minθPΘ 1 |S| ř px , yqPS ` pfθpxq , yq ` Ωpθq , where fθ : X ÞÑ Y is a function from input domain X to output domain Y parameterized by elements θ P Θ , ` : Y ˆY ÞÑ R is a scalar-valued loss function , Ω : Θ ÞÑ R is a scalar-valued regularizer , and S Ă X ˆY is a finite set of ( self- ) supervised training examples . We study the setting where fθ is a neural network , an L-layer function whose parameters θ consist of L matrices Wi P Rmiˆni and whose output fθpxq given input x is defined recursively using L functions gi via the formula xi “ gipWi , xi´1q , with x0 “ x and fθpxq “ xL . The standard approach to training fθ is to specify the regularizer Ω , ( randomly ) pick an initialization in Θ , and iteratively update the parameters using some first-order algorithm such as SGD to optimize the objective above until some stopping criterion is met . However , in many cases we instead optimize over factorized variants of these networks , in which some or all of the matrices Wi P Rmiˆni are re-parameterized as a product Wi “ Uip śdi j “ 1 MijqV Ti for some inner depth di ě 0 and matrices Ui P Rmiˆri , Vi P Rniˆri , and Mij P Rriˆri @ j . As discussed in the following examples , this can be done to obtain better generalization , improve optimization , or satisfy practical computational or memory constraints during training or inference . For simplicity , we drop the subscript i whenever re-parameterizing only one layer and only consider the cases when inner depth d is 0 or 1 . 2.1 FULLY-CONNECTED LAYERS . A fully-connected layer takes an n-dimensional input xi´1 and outputs an m-dimensional vector xi “ σpWxi´1q , where σ : Rm ÞÑ Rm is an element-wise activation function . Here , decomposing W P Rmˆn into the product UV T , where U P Rmˆr , V P Rnˆr , and setting r ! mintm , nu reduces computation and memory costs from Opmnq to Opmr ` nrq . We refer to this setting as model compression . Standard learning theory suggests that a small rank r also improves generalization , e.g . for a factorized fully-connected ReLU network , applying } W } 2F { } W } 22 ď rankpW q to Neyshabur et al . ( 2018 , Theorem 1 ) and substituting Wi “ UiV Ti gives a w.h.p . margin-bound Õp a mr { |S|q suggesting that generalization error varies with the square root of the rank ( see Corollary A.1 ) . Alternatively , by setting r ě mintm , nu and/or including an inner matrix M P Rrˆr , we can attempt to take advantage of improved optimization due to increased width ( Du & Hu , 2019 ) and/or increased depth ( Arora et al. , 2018 ) . Crucially , this does not increase inference costs because we can recompose the matrix after training and just use the product . As the goal is to obtain a better small model by first training a large one , we refer to this setting as overcomplete knowledge distillation ; of course , unlike regular distillation it is much simpler since there is no student-teacher training stage . 2.2 CONVOLUTIONAL LAYERS . A 2d convolutional layer takes an hˆ w ˆ ci´1-dimensional input xi´1 and outputs a hˆ w ˆ cidimensional output xi defined by convolving ci different k ˆ k filters over each of ci´1 input channels . Often the result is passed through a nonlinearity . 2d convolutional layers are parameterized by ciˆ ci´1ˆ kˆ k tensors and require Opk2cici´1q memory and compute . A straightforward way of factorizing this tensor without using tensor decomposition is to reshape it into a cik ˆ ci´1k matrix W , which can then be decomposed as W “ UV T for U P Rcikˆr , V P Rci´1kˆr and some rank r ą 0 . As in the fully-connected case , we can either set the rank r to be small in order to reduce the number of parameters or alternatively increase the width ( r ) or the depth ( d ) of the factorization to do overcomplete knowledge distillation . Note that in the low-rank case a naive approach does not save computation since we must first multiply U and V T , reshape the product UV T , and then use the resulting tensor in a regular 2d convolution of the original size and complexity . However , as shown by Tai et al . ( 2016 ) , applying the 2d k ˆ k convolution with ci´1 input channels and ci output channels obtained by reshaping UV T is equivalent to a composition of two 1d convolutions : the first defined by V T P Rrˆci´1k consists of r output channels and filters of size k along one input dimension and the second defined by U P Rcikˆr consists of ci output channels and filters of size k along the other input dimension . Together the two 1d convolutions require Opkrpci ` ci´1qq memory and computation , which is significantly better than the Opk2cici´1q cost of the unfactorized case if r ! kmintci , ci´1u .

Summary: This paper studies initialization and regularization in factorized neural networks (reparameterize a weight matrix by the product of several weight matrices). The authors proposed spectral initialization, that is to initialize the factorized matrices using the SVD of the un-factorized matrix. The authors also proposed Frobenius decay that is to regularize the Frobenius norm of the product of the factorized weight matrices. The motivation is to simulate the routines for non-decomposed counterparts. The authors empirically showed the effectiveness of spectral initialization and Frobenius decay in different applications: compressed model training, knowledge distillation, and multi-head self-training.