I NITIALIZATION AND R EGULARIZATION OF F ACTORIZED N EURAL L AYERS
A BSTRACT
Factorized layers—operations parameterized by products of two or more matrices—occur in a variety of deep learning contexts, including compressed model training, certain types of knowledge distillation, and multi-head selfattention architectures.
We study how to initialize and regularize deep nets containing such layers, examining two simple, understudied schemes, spectral initialization and Frobenius decay , for improving their performance.
The guiding insight is to design optimization routines for these networks that are as close as possible to that of their well-tuned, non-decomposed counterparts; we back this intuition with an analysis of how the initialization and regularization schemes impact training with gradient descent, drawing on modern attempts to understand the interplay of weight-decay and batch-normalization.
Empirically, we highlight the beneﬁts of spectral initialization and Frobenius decay across a variety of settings.
In model compression, we show that they enable low-rank methods to signiﬁcantly outperform both unstructured sparsity and tensor methods on the task of training low-memory residual networks; analogs of the schemes also improve the performance of tensor decomposition techniques.
For knowledge distillation, Frobenius decay enables a simple, overcomplete baseline that yields a compact model from over-parameterized training without requiring retraining with or pruning a teacher network.
Finally, we show how both schemes applied to multi-head attention lead to improved performance on both translation and unsupervised pre-training.
I NTRODUCTION
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 deﬁned as products of two or more matrices.
When these are smaller than the original matrix, the resulting networks are more efﬁcient 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 ﬁrst 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-Perpi˜n´an, 2020).
We formalize and study the justiﬁcations 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-Perpi˜n´an, 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 ﬂedgling subﬁeld 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 signiﬁcant 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— FLAMB´e —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. R ELATED W ORK
We are not the ﬁrst 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.
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 ﬁnd 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 justiﬁed via the usual Bayesian understanding of (cid:96) 2 -regularization (Murphy, 2012).
However, we ﬁnd 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).
P RELIMINARIES ON F ACTORIZED N EURAL L AYERS
In the training phase of (self-)supervised ML, we often solve optimization problems of the form min θ P Θ 1 | S | (cid:176) p x , y qP S (cid:96) p f θ p x q , y q (cid:0)
Ω p θ q , where f θ : X (cid:222)(cid:209) Y is a function from input domain X to output domain Y parameterized by elements θ P Θ , (cid:96) :
Y (cid:2) Y (cid:222)(cid:209) R is a scalar-valued loss function, Ω : Θ (cid:222)(cid:209) R is a scalar-valued regularizer, and S (cid:128) X (cid:2)
Y is a ﬁnite 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 W i P R m i (cid:2) n i and whose output f θ p x q given input x is deﬁned recursively using L functions g i via the formula x
i (cid:16) g i p W i , x i (cid:1) 1 q , with x 0 (cid:16) x and f θ p x q (cid:16) x L .
The standard approach to training f θ is to specify the regularizer Ω , (randomly) pick an initialization in Θ , and iteratively update the parameters using some ﬁrst-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 W i
P
R m i (cid:2) n i are re-parameterized as a product W i (cid:16) U i p – d i j (cid:16) 1 M ij q V Ti for some inner depth
d i ¥ 0 and matrices
U i P
R m i (cid:2) r i , V i P R n i (cid:2)
r i , and M ij P R r i (cid:2) r i @ 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. F ULLY -C ONNECTED L AYERS
A fully-connected layer takes an n -dimensional input x i (cid:1) 1 and outputs an m -dimensional vector x i (cid:16) σ p W x
i (cid:1) 1 q , where σ : R m (cid:222)(cid:209) R m is an element-wise activation function.
Here, decomposing W
P R m (cid:2) n into the product UV T , where U P R m (cid:2) r , V P R n (cid:2) r , and setting r !
min t m, n u reduces computation and memory costs from O p mn q to O p mr (cid:0) nr q .
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 }
F {} W } 22 ⁄ rank p W q to Neyshabur et al.
(2018, Theorem 1) and substituting W i (cid:16)
U
i V T i gives a w.h.p.
margin-bound ˜ O pa mr {| S |q suggesting that generalization error varies with the square root of the rank (see Corollary A.1).
Alternatively, by setting r ¥ min t m, n u and/or including an inner matrix M P R r (cid:2) 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 ﬁrst 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. C ONVOLUTIONAL L AYERS
A 2d convolutional layer takes an h (cid:2) w (cid:2)
c i (cid:1) 1 -dimensional input x i (cid:1) 1 and outputs a h (cid:2) w (cid:2) c
i - dimensional output x i deﬁned by convolving c i different k (cid:2) k ﬁlters over each of c
i (cid:1) 1 input channels.
Often the result is passed through a nonlinearity.
2d convolutional layers are parameterized by c i (cid:2)
c i (cid:1) 1 (cid:2) k (cid:2) k tensors and require O p k 2 c
i c i (cid:1) 1 q memory and compute.
A straightforward way of factorizing this tensor without using tensor decomposition is to reshape it into a c i k (cid:2)
c i (cid:1) 1 k matrix W , which can then be decomposed as W (cid:16) UV T for U P R c
i k (cid:2) r , V P R c i (cid:1) 1 k (cid:2) 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 ﬁrst 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 (cid:2) k convolution with c
i (cid:1) 1 input channels and c i output channels obtained by reshaping UV T is equivalent to a composition of two 1d convolutions: the ﬁrst deﬁned by V T
[Equation]
consists of r output channels and ﬁlters of size k along one input dimension and the second deﬁned by U P R c i k (cid:2) r consists of c i output channels and ﬁlters of size k along the other input dimension.
Together the two 1d convolutions require O p kr p
c i (cid:0) c i (cid:1) 1 qq memory and computation, which is signiﬁcantly better than the O p k 2 c
i c i (cid:1) 1 q cost of the unfactorized case if r !
k min t c
i , c i (cid:1) 1 u .
2. M ULTI -H EAD A TTENTION
An MHA layer (Vaswani et al., 2017) with H attention heads and hidden dimension d can be expressed as being parameterized with 4 H matrices: one each of Q h , K h , V h , O h P R d (cid:2) d { H for each head h .
Then for a length- T input x P R T (cid:2) d it outputs
H ‚
[Equation]
Softmax
[Equation]
MHA combines 2 H quadratic forms Q
I NITIALIZATION AND R EGULARIZATION
We now deﬁne the initialization and regularization schemes we study as natural extensions of techniques for non-factorized models.
They thus require no tuning when an existing training implementation of a non-factorized deep net is available.
We later discuss how to justify these schemes when layers are normalized, e.g. using BatchNorm (Ioffe & Szegedy, 2015).
In all experiments with convolutional models in this and subsequent sections we factorize all layers except the ﬁrst and last, which are small, and determine layer ranks by multiplying a uniform scaling factor by the product of a layer’s output channels and kernel width.
This rank-scale can be varied to attain the desired number of parameters.
Note that our approach may be further improved via a more sophisticated or adaptive rank-assignment scheme (Idelbayev & Carreira-Perpi ˜ n ´ an, 2020).
3. S PECTRAL I NITIALIZATION
Initialization is a major focus of deep learning research (He et al., 2015; Mishkin & Matas, 2016;
Yang & Schoenholz, 2017).
A common approach is to prevent compounding changes in the norms of the intermediate representations across layers caused by repeated multiplication by the weight matrices.
The spectral initialization scheme for initializing low-rank factorized layers attempts to inherit, when possible, this property from an existing initialization by using SVD to ensure that the resulting product matrix is as close as possible to the original parameter:
Deﬁnition 3.1.
Let W P R m (cid:2) n be a parameter of an unfactorized layer.
For r ⁄ min t m, n u the spectral initialization (SI) of the factors U P R m (cid:2) r , V P R n (cid:2) r of the corresponding factorized layer sets U (cid:16) ˜
U ?
Σ and V (cid:16)
˜ V ?
Σ , for ˜ U , Σ , ˜ V (cid:16) SVD r p W q given by the rank- r SVD of W .
SI preserves the largest singular value of W , so if the original scheme did not suffer from a compounding increase in representation norm then neither will spectral initialization.
On the other hand, while low-rank layers impose a nullspace, SI aligns it with the directions minimized by W .
3. F ROBENIUS D ECAY
Weight-decay is a common regularizer for deep nets, often implemented explicitly by adding Ω p
[Equation]
W P θ } W } 2 F to the objective for some λ ¥ 0 .
Classically, it is thought to improve generalization by constraining model capacity.
When training factorized layers parameterized by
U p –
[Equation]
j q V T , the easiest approach to implement is to replace each λ 2 } W } 2 F term in Ω p θ q by
[Equation]
(cid:9) . However, this yields a very different optimization problem: for example, if we consider the case of d (cid:16)
then this regularizer is in fact an upper bound on the nuclear norm of the recomposed matrix UV T (Srebro & Shraibman, 2005, Lemma 1):
[Equation]
In fact, Figure 1 shows that for weight-decay this upper bound is tight throughout the training of factorized ResNet20 across a variety of ranks and initializations, suggesting that the naive approach is indeed regularizing the nuclear norm rather than the Frobenius norm.
Since compression already constrains capacity, in the low-rank case one might favor just reducing regularization, e.g. multiplying λ by the compression rate (Gray et al., 2019).
However, Figureshows that this can lead to worse performance, and the approach still penalizes the nuclear norm.
Frobenius decay avoids this issue by simply penalizing the squared norm of the entire factorization:
Deﬁnition 3.2.
For λ ¥ 0 let λ 2 } W } 2 F be the contribution of an unfactorized layer parameterized by
W P R m (cid:2) n to the penalty term.
Then the Frobenius decay (FD) penalty on matrices U P R m (cid:2) r ,
V P R n (cid:2) r , and M j P R r (cid:2) r of the corresponding factorized layer is λ(cid:15)(cid:15)(cid:15)
[Equation]
[Equation]
(cid:16)
[Equation]
(cid:15)(cid:15)(cid:15)
F .
By substituting the factorization directly, FD makes the least change to the problem:
rank- r optima of the non-factorized objective will also minimize the factorized one.
We can also bound generalization error of the ReLU net from Section 2.1 by a term ˜
O
[Equation]
varying directly with the quantity penalized by FD (see Corollary A.2).
Notably, FD is a stronger penalty than the nuclear norm implicitly regularized by weight-decay, yet it still yields better models.
3. I NITIALIZATION AND R EGULARIZATION IN THE P RESENCE OF N ORMALIZATION
The use of spectral initialization and Frobenius decay is largely motivated by the norms of the recomposed matrices: SI prevents them increasing feature vector norms across layers, while FD constrains model capacity via parameter norms.
However, normalization layers like BatchNorm (Ioffe & Szegedy, 2015) and others (Ba et al., 2016) largely negate the forward-pass and modelcapacity effects of the norms of the weights parameterizing the layers they follow.
Thus for most modern models we need a different explanation for the effectiveness of SI and FD.
Despite the fact that most layers’ norms do not affect inference or capacity, weight-decay remains useful for optimizing deep nets.
Recently, Zhang et al.
(2019) extended the analysis of Hoffer et al.
(2018) to argue that the effective step-size of the weight direction ˆ W is roughly η {} W } 2 F , where η is the SGD step-size.
Thus by preventing the norm of W from growing too large weight-decay maintains a large-enough effective step-size during training.
We draw on this analysis to explore the effect of SI and FD on factorized models.
For simplicity, we deﬁne a normalized layer to be one that does not depend on the scale of its parameter.
Ignoring stability offset terms, this deﬁnition roughly holds for normalized linear layers, convolutional layers in ResNets, and the Output-Value quadratic form in Transformers if the residual connection is added after rather than before normalization.
Deﬁnition 3.3.
A normalized layer g p W , x q parameterized by W P R m (cid:2) n is one that satisﬁes g p W , x q (cid:16) g p ρ W , x q for all W P
R m (cid:2) n and all positive scalars ρ .
Because the output does not depend on the magnitude of UV T , what matters is the direction of the composed matrix.
During an SGD step this direction is updated as follows (proof in Appendix B):
Claim 3.1.
At all steps t ¥ 0 let g be a normalized layer of a differentiable model f θ t : X (cid:222)(cid:209) Y parameterized by U t V T t for U t P
R m (cid:2) r , V Tt P R n (cid:2) r , r ¥ 1 .
Suppose we update P (cid:16) U and
P (cid:16) V by SGD, setting P t (cid:0) 1 — P t (cid:1) η ∇ P t using a gradient ∇ P t (cid:16) 1 | B | (cid:176) p x , y qP B ∇ P t (cid:96) p f θ t p x q , y q over batch B (cid:128) p X , Y q and η ¡ 0 sufﬁciently small.
Then for ˆ ∇ t (cid:16) ∇ ˆ W t V t V Tt (cid:0)
U t U
Tt ∇ ˆ W t we have that the vectorized direction ˆ w t (cid:16) vec p
ˆ W t q of W t (cid:16)
U t V Tt is updated as
[Equation]
Note that (1) we ignore decay because λ (cid:16) O p η q so any resulting term is O p η 2 q and (2) the update rule is almost the same as that obtained for the unfactorized case by Zhang et al.
(2019), except they have ˆ ∇ t as the true gradient of the direction.
Thus, apart from a rank-one correction, ˆ W t is approximately updated with step-size η {} W t } 2 F multiplying a linear transformation of its gradient.
To understand the nature of this transformation, note that at spectral initialization we have that
V 0 V T 0 (cid:16) U
[Equation]
Σ r are diagonal matrices of singular values of the full-rank initialization W ; furthermore, if W is a Gaussian ensemble with scale 1 {? n , which roughly aligns with common initialization schemes (He et al., 2015), then its singular values are roughly distributed around 1 and supported on r 0 , 2 s (Bai & Yin, 1993).
Since ˆ ∇ t (cid:16) ∇ ˆ W t V t V Tt (cid:0)
U t U
Tt ∇ ˆ W t , this suggests that, at spectral initialization, an effective learning rate of η {} W 0 } 2 F is a reasonable approximation for the factorized update.
This points to the role of SI being to initialize the factorization at an appropriate scale and perhaps also to make the ﬁrst update more aligned with the gradient w.r.t.
ˆ W 0 .
As in the unfactorized case, our analysis suggests that, the main role of decay may be to maintain a large effective learning rate η {} W } 2 F ; furthermore, FD may be more effective than regular decay because it provides stronger regularization and directly penalizes the quantity of interest.
We support this hypothesis using experiments analogous to Zhang et al.
by comparing training a low-rank
ResNet-20 with FD to training it with no decay but at each step normalizing all BatchNormed layers to have the same Frobenius norm as the FD run.
In Figure 3 we see that the latter scheme closely tracks the former in terms of both the training loss and the test accuracy.
Figure 3 also shows that
FD maintains a higher effective step-size than regular weight-decay throughout training.
C OMPRESSED M ODEL T RAINING : L OW -R ANK , S PARSE , AND T ENSORIAL
We ﬁrst study SI and FD for training factorized models with low-rank convolutional layers, comparing against the dominant approaches to low-memory training: sparse layers and tensor decomposition.
For direct comparison we evaluate models with near-identical parameter counts; note that by normalizing by memory we disadvantage the low-rank approach, which often has a comparative advantage for speed.
All models are trained for 200 epochs with the same optimizer settings as for the unfactorized models; the weight-decay coefﬁcient is left unchanged when replacing by FD.
Low-Rank and Sparse Training: We train modiﬁed ResNet32 and VGG19 used in the lotteryticket-guessing literature (Wang et al., 2020; Su et al., 2020).
Motivated by Frankle & Carbin (2019), such methods ﬁx a sparsity pattern at initialization and train only the unpruned weights.
While they achieve high parameter savings, their computational advantages over full models are less clear, as software and accelerators often do not efﬁciently implement arbitrary sparsity (Paszke et al., 2019;
NVIDIA, 2020).
In contrast, we do see acceleration via low-rank convolutions, almost halving the time of ResNet’s forward and backward pass at the highest compression.
For completeness, we also show methods that vary the sparsity pattern (dynamic), prune trained models (pruning), and prune trained models and retrain (lottery); note that the latter two require training an uncompressed model.
In Table 1 we see that the low-rank approach, with SI & FD, dominates at the higher memory settings of ResNet across all three datasets considered, often outperforming even approaches that train an uncompressed model ﬁrst.
It is also close to the best compressed training approach in the lowest memory setting for CIFAR-100 (Krizhevksy, 2009) and Tiny-ImageNet (Deng et al., 2009).
[Table] [Table]
On the other hand, the low-rank approach is substantially worse for VGG; nevertheless, ResNet is both smaller and more accurate, so if the goal is an accurate compressed model learned from scratch then one should prefer the low-rank approach.
Our results demonstrate a strong, simple baseline not frequently compared to in the low-memory training (Frankle & Carbin, 2019;
Wang et al., 2020;
Su et al., 2020).
In fact, since it preserves the top singular components of the original weights,
SI can itself be considered a type of (spectral) magnitude pruning.
Finally, Table 1 highlights the complementary nature of SI and FD, which outperform regular low-rank training on both models, although interestingly they consistently decrease ResNet performance when used separately.
Matrix and Tensor Decomposition: We next compare against tensor decompositions, another common approach to small model training.
A recent evaluation of tensors and other related approaches by Gray et al.
(2019) found that Tensor-Train decomposition (Oseledets, 2011) obtained the best memory-accuracy trade-off on WideResNet; we thus compare directly to this approach.
Note that, while we normalize according to memory, Tensor-Train must be expanded to the full tensor prior to convolution and thus increases the required compute, unlike low-rank factorization.
In Table 2 we show that at 6.7% of the original parameters the low-rank approach with FD and SI signiﬁcantly outperforms Tensor-Train.
Tensor-Train excels at the highly compressed 1.7% setting but is greatly improved by leveraging tensor analogs of SI—decomposing a random initialization rather than randomly initializing tensor cores directly—and FD—penalizing the squared Frobenius norm of the full tensor.
We also compare to CRS (Gray et al., 2019), which scales regular weightdecay by the compression rate.
It is roughly as beneﬁcial as FD across different evaluations of
Tensor-Train, but FD is signiﬁcantly better for low-rank factorized neural layers.
[Table]
O VERCOMPLETE K NOWLEDGE D ISTILLATION
In contrast to compressed model training, in knowledge distillation (KD) we have the capacity to train a large model but want to deploy a small one.
We study what we call overcomplete KD, in which a network is over-parameterized in such a way that an equivalent small model can be directly recovered.
Similar approaches have been previously studied only with small models (Arora et al., 2018; Guo et al., 2020) or using convolution-speciﬁc methods (Ding et al., 2019; Cao et al., 2020).
We take a simple factorization approach in which we decompose weights W P R m (cid:2) n to be products of 2 or 3 matrices while increasing the parameter count, either via depth or width.
Here we consider three cases: the full (rank) setting where W (cid:16) UV T for U P R m (cid:2) m , V P R n (cid:2) m , the deep setting where W (cid:16) UMV T for U , M P R m (cid:2) m , V P R n (cid:2) m , and the wide setting where W (cid:16) UV T for
U P R m (cid:2) 3 m , V P R n (cid:2) 3 m .
As before, we factorize all but the ﬁrst and last layer and train factorized networks using the same routine as for the base model, except when replacing weight-decay by FD.
We do not study SI here, using the default initialization for U , V T and setting M (cid:16) I m .
Table 3 shows the strength of this approach: we can train ResNet32 to beat ResNet56 and ResNet to beat ResNet110 while needing about the same number of parameters during training and only half as many at inference. Note that training ResNet56 using the “full” setting is also 1.5x faster than training ResNet110 (see Table 8 for timings).
Furthermore, we improve substantially upon DOConv (Cao et al., 2020), which obtains much smaller improvements over the unfactorized baseline. In fact, Table 9 suggests our approach compares favorably even with regular KD methods; combined with its simplicity and efﬁciency this suggests that our overcomplete method can be considered as a baseline in the ﬁeld.
Finally, we also show that Frobenius decay is critical: without it, overcomplete KD performs worse than regular training.
A detailed visualization of this, comparing the CIFAR performance of factorized ResNets across a range of rank-scale settings covering both the current high-rank (distillation) case and the previous low-rank (compression) case, can be found in Figure 2.
M ULTI -H EAD A TTENTION AS F ACTORIZED Q UADRATIC F ORMS
Our ﬁnal setting is multi-head attention (Transformer) architectures (Vaswani et al., 2017).
As discussed in Section 2, the MHA component is already a factorized layer, consisting of an aggregation over the output of two quadratic forms per head: the “Query-Key” (QK) forms passed into the softmax and the “Output-Value” (OV) forms that multiply the resulting attention scores (c.f. Equation 1).
Transformers also contain large linear layers, which can also be factorized for efﬁciency.
Improving Transformer Training and Compression: We start with the original Transformer architecture on the IWSLT-14 translation task (Cettolo et al., 2014) but use an SGD-based training routine as a baseline (Gehring et al., 2017).
As there is no weight-decay by default, we ﬁrst tune both this and FD on the non-factorized model; here FD is applied to the implicitly factorized MHA layers. ters as
[Table] In Figure 4 we show that this alone yields an improvement: whereas the effect of weight-decay is either negligible or negative, tuning FD does improve the BLEU score.
Furthermore, we see that tuning just the OV form in MHA is more robust at higher regularization levels than tuning both OV and QK.
We conclude by examining both SI and FD when reducing the number of parameters by (1) factorizing all linear and embedding layers and (2) scaling down the embedding dimension in
MHA.
In Figure 4 we see that the beneﬁt of Frobenius decay disappears when compressing; on the other hand, SI provides a strong boost under both types of decay, and is in fact necessary for FD to work at all.
Note that the major effect here is for the factorized linear layers—we found that SI has minimal effect when applied to MHA, likely because those initializations have already been tuned.
FLAMB´e for Unsupervised BERT Pre-Training: Lastly, we examine BERT (Devlin et al., 2019), a large transformer trained on a massive unsupervised text corpus and evaluated on downstream language tasks.
The state-of-the-art training approach is via the LAMB optimizer (You et al., 2020) using weight-decay based on the AdamW algorithm of Loshchilov & Hutter (2019), in which λη times each parameter is subtracted from itself; this is equivalent to (cid:96) 2 -regularization for SGD but not for adaptive methods.
We can deﬁne a similar Frobenius alternative by subtracting λη times the Frobenius gradients UV T V (cid:16) ∇ U 12 } UV T } 2 F and V U T U (cid:16) ∇ V 12 } UV T } 2 F from U and V , respectively; when used with the LAMB optimizer we call this method FLAMB ´ e.
We see in Tablethat FLAMB ´ e outperforms the simple FD modiﬁcation of LAMB and, as with IWSLT, leads to an improvement in downstream task performance without changing model size.
For BERT, however, applying FD via FLAMB´e also leads to better downstream performance when scaling down the
MHA embedding dimension by half.
Besides achieving a better compressed model, the success of
FLAMB´e also shows the potential of new types of decay schemes for adaptive methods.
C ONCLUSION
In this paper we studied the design of training algorithms for deep nets containing factorized layers, demonstrating that two simple specializations of standard initialization and regularization schemes from the unfactorized case lead to strong improvements for model compression, knowledge distillation, and MHA-based architectures.
While we largely focused on the case where the unfactorized model uses Gaussian initialization and (cid:96) 2 -regularization, we believe our work provides guidance for the many cases where other schemes are used to enforce alternative priors or improve optimization.
For example, SI as deﬁned can be applied to any random initialization, while our FD results suggest that regularizers such as penalty terms and DropConnect (Wan et al., 2013) should be applied on the product matrix rather than directly on individual factors.
The success of SI and FD for both low-rank and tensor decomposition also suggests that these schemes may be useful for other types of factorized neural layers, such as ACDC (Moczulski et al., 2015) or K-matrices (Dao et al., 2020).
A G ENERALIZATION OF F ACTORIZED L AYERS
In this section we brieﬂy discuss how to apply the generalization bounds in Neyshabur et al.
(2018) to factorized models.
Deﬁnition A.1.
For any γ ¡ 0 and distribution D over classiﬁcation data X (cid:2) Y the γ -margin-loss of a model f θ : X (cid:222)(cid:209) Y is deﬁned as (cid:96)
[Equation]
Note that (cid:96)
D is the expected classiﬁcation loss over the distribution and (cid:96) γ Uniform p S q is the empirical γ -margin-loss.
We have the following corollaries:
Corollary A.1.
Let f θ be a neural network with L (cid:1) 1 factorized fully-connected ReLU layers
g i p
U
i V Ti , x i (cid:1)
[Equation]
q (cid:16) max
[Equation]
V Ti
[Equation]
1 , 0 m ( of hidden dimension m and one factorized classiﬁcation layer g L p U L V TL , x L (cid:1) [-
-]1 q (cid:16) U L V TL x L (cid:1) 1 .
Let D be a distribution over B -bounded classiﬁcation data X (cid:2) Y and suppose we have a ﬁnite set S of i.i.d.
samples from it.
If rank p U
i V Ti q ⁄ r
and } U
i V Ti } 2 u ⁄ σ for all
i then for any δ ¡
we have w.p. 1 (cid:1) δ that
[Equation]
O
[Equation]
Proof.
Apply the inequality } W i } 2 F {} W i } 22 ⁄ rank
p W i q to Neyshabur et al.
(2018, Theorem 1) and substitute W i (cid:16) U
i V T i .
Corollary A.2.
Let f θ be a neural network with L (cid:1) 1 factorized fully-connected ReLU layers
g i
[Equation]
of hidden dimension m
and one factorized classiﬁcation layer
[Equation]
Let D be a distribution over B -bounded classiﬁcation data X (cid:2) Y from which we have a ﬁnite set S of i.i.d.
samples.
If } U i –
[Equation]
1 M ij V Ti } 2 u ⁄ σ for all i then for any δ ¡
we have w.p. 1 (cid:1) δ that
[Equation]
Proof.
Apply the equality
[Equation]
} 22 to Neyshabur et al.
(2018, Theorem 1) and substitute W i (cid:16) U i p
[Equation]
B P ROOF OF C LAIM 3.
Proof.
Let ρ t (cid:16) } W t } F be the Frobenius norm of the composed matrix at time t .
Applying the update rules for U t and V t and using the fact that ρ t ∇ W t (cid:16) ∇ ˆ W t yields
[Equation]
Taking the squared norm of both sides yields ρ 2 t (cid:0) 1 (cid:16) ρ 2 t
[Equation]
η [Equation] p ˆ
[Equation]
ˆ ∇
[Equation] q [Equation] O p η 2 q ; we can then take the square root of both sides and use a Tayor expansion to obtain
[Equation]
Then, starting from Equation 6 divided by ρ t (cid:0) 1 , substituting Equation 7, and applying a Taylor expansion yields
[Equation]
Vectorizing and observing that p ˆ w Tt vec p ˆ ∇ t qq ˆ w t (cid:16) ˆ w t ˆ w Tt vec p ˆ ∇ t q yields the result.
C E XPERIMENTAL D ETAILS FOR T RAINING C ONVOLUTIONAL N ETWORKS
For experiments with regular ResNets on CIFAR we used code provided here: https:// github.com/akamaster/pytorch_resnet_cifar10 .
All hyperparameter settings are the same, except initialization and regularization as appropriate, with the exception that we use a warmup epoch with a 10 times smaller learning rate for ResNet56 for stability (this is already done by default for ResNet110).
For comparisons with sparse model training of ResNet32 x2 and VGG19 NFC we use code by Wang et al.
(2019, https://github.com/alecwangcq/EigenDamage-Pytorch ), which is closely related to that of the lottery-ticket-guessing paper by Wang et al.
All hyperparameter settings are the same, except (1) initialization and regularization as appropriate and (2) for
Tiny-ImageNet we only train for 200 epochs instead of 300.
For comparisons with tensor decomposition training of WideResNet we use code by Gray et al.
(2019, https://github.com/BayesWatch/deficient-efficient ).
All hyperparameter settings are the same, except initialization and regularization as appropriate.
D E XPERIMENTAL D ETAILS FOR T RAINING T RANSFORMER M ODELS
For experiments with Transformer models on machine translation we used code provided here: https://github.com/StillKeepTry/Transformer-PyTorch .
All hyperparameter settings are the same, except initialization and regularization as appropriate.
For comparisons with LAMB optimization of BERT, we use an implementation provided by
NVIDIA:
https://github.com/NVIDIA/DeepLearningExamples/tree/master/
PyTorch/LanguageModeling/BERT .
All hyperparameter settings are the same, except initialization and regularization as appropriate.
For ﬁne-tuning on SQuAD we apply the same optimization routine to all pre-trained models. E P
AST W ORK ON K NOWLEDGE D ISTILLATION
We ﬁrst brieﬂy summarize past work on overcomplete KD.
While the use of factorized neural layers for improving training has theoretical roots (Arora et al., 2018; Du & Hu, 2019), we are aware of two other works focusing on experimental practicality: ExpandNets (Guo et al., 2020) and DOConv (Cao et al., 2020).
As the former focuses on small student networks, numerically we compare directly to the latter, showing much better improvement due to distillation for both ResNet56 and ResNet110 on both CIFAR-10 and CIFAR-100 in Table 3. Note that we are also aware of one paper proposing a related method that also trains an over-parameterized model without increasing expressivity that can then be collapsed into a smaller “original” model (Ding et al., 2019); their approach, called ACNet, passes the input to each layer through differently-shaped kernels that can be composed additively. Note that it is unclear how to express this method as a factorization and it may not be easy to generalize to non-convolutional networks, so we do not view it as an overcomplete
KD approach.
We now conduct a brief comparison with both these works and more standard approaches to KD.
Direct comparison with past work is made difﬁcult by the wide variety of training routines, teacher models, student models, and evaluation procedures employed by the community. Comparing our speciﬁc approach is made even more difﬁcult by the fact that we have no teacher network, or at least not one that is a standard model used in computer vision. Nevertheless, in Table 9 we collect an array of existing results that can be plausibly compared to our own overcomplete distillation of ResNets. Even here, note that absolute numbers vary signiﬁcantly, so we focus on changes in accuracy.
As can be seen from the results, our overcomplete approach yields the largest improvements for ResNet and ResNet110 on both CIFAR-10 and CIFAR-100. For ResNet32, the Snapshot Distillation method of Xu & Liu (2019) outperforms our own, although it does not do so for ResNet110 and is not evaluated for ResNet56. On CIFAR-10 the additive ACNet approach also has a larger performance improvement for ResNet32. Nevertheless, our method is still fairly close in these cases, so the results in Table 9 together with the simplicity and short (single-stage) distillation routine of our overcomplete approach suggest that it should be a standard baseline for KD.
[Table]
[Table] [Table] [Table] [Table]
[Table] [Table] [Table]
[Table]
(Kim et al., 2018) ResNet 71.
+2.
Snapshot (Yang et al., 2019) self 70.
+0.
DO-Conv (Cao et al., 2020) overcomplete 70.
+0.
wide factorization w. FD (ours) overcomplete 70.
+2. SD (Xu & Liu, 2019) self 73.
+1.
Snapshot (Yang et al., 2019) self 71.
+1.
DO-Conv (Cao et al., 2020) overcomplete 72.
+0.
deep factorization w. FD (ours) overcomplete 72.
+2.