I NITIALIZATION AND R EGULARIZATION OF F ACTORIZED N EURAL L AYERS
Anonymous authors
Paper under double-blind review
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, self-taught 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, self-taught (i.e. 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 self-taught knowledge distillation (Arora et al., 2018; Guo et al., 2020; Cao et al., 2020) in which a model’s weights are over-parameterized as over-complete factorizations; the factors are then multiplied after training to obtain the compact representation of the same network.
We show that FD leads to signiﬁcant improvements, demonstrating self-taught training of a ResNet56 model that outperforms a regularly-trained ResNet110.
Finally, we study Transformer architectures, starting by showing that FD improves translation per formance 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 (knowledge distillation, multihead 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.
1. R ELATED W ORK
We are not the ﬁrst study at 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 propose what is effectively self-taught knowledge distillation.
Rather than the descent method, we focus on the initialization and regularization of factorized layers.
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 useful on regular ResNets but inferior to SI on wider networks.
Outside deep learning, spectral methods have also been shown to yield better gradient descent initializations for certain matrix and tensor problems (Chi et al., 2019;
Cai et al., 2019).
For regularization, Gray et al.
(2019) suggest compression-rate scaling , which scales weight-decay using the reduction in parameter count; this is justiﬁed via the standard 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 this approach.
Our own analysis is based on recent work suggesting that the function of weight-decay in modern deep nets 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 .
Low-rank decomposition also decreases the model complexity, which standard learning theory suggests can improve generalization error of the model; indeed, some recent results bound the generalization error of a neural network by the stable rank (and thus also the rank) of its weight matrices (Neyshabur et al., 2018, Theorem 1).
We refer to this setting as model compression .
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 self-taught 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
[Equation]
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 self-taught 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
P R r (cid:2) c i (cid:1) 1 k 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 ‚ h (cid:16)
Softmax (cid:3) xQ
h K Th x T a d { H (cid:11)
xV h
O Th (1)
MHA combines 2 H quadratic forms Q h K Th ,
V h O Th of rank r (cid:16) d
{ H , each a product of matrices, i.e. a factorized layer.
We refer to the ﬁrst form as “Query-Key” and the second as “Output
[Table]
Value.”
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
[Equation]
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
[Equation]
T , the easiest approach to implement is to replace each λ 2 } W } 2 F term in Ω p θ q by
[Equation]
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)
U p – dj (cid:16) 1 M j q V T (cid:15)(cid:15)(cid:15)
F .
By substituting the factorization into the original regularized empirical risk minimization, FD makes the minimal possible change to the objective.
In particular, any rank- r optimum of the non-factorized problem will also minimize the factorized one.
Notably, the squared Frobenius norm 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’ weight-norms do not affect inference or capacity, regular weightdecay remains useful for optimizing deep nets.
Recently, Zhang et al.
(2019) extended the analysis of Hoffer et al.
(2018) to argue that the effective learning rate 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 learning rate 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.
While often not precisely true due to stability offset terms, it otherwise 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 A):
Claim 3.1.
Let g be a normalized layer of a differentiable model f θ : X (cid:222)(cid:209) Y parameterized by the product UV T of matrices U P R m (cid:2) r , V T P R n (cid:2) r , r ¥ 1 .
Suppose at step t its parameters are updated by SGD with sufﬁciently small step-size η ¡ 0 via the rules U t (cid:0) 1 — U t (cid:1) η ∇ U t and V t (cid:0) 1 — V t (cid:1) η ∇ V t , where for parameter P we let ∇ P (cid:16) 1 | B | (cid:176)
p x,y qP B ∇ P (cid:96) p f θ t p x q , y q be the gradient at θ t w.r.t.
P using batch B (cid:128) p X , Y q .
Then deﬁning ˆ ∇ 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 has an update satisfying
[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
[Equation]
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]
S ELF -T AUGHT K NOWLEDGE D ISTILLATION
In contrast to compressed model training
[Table]
in knowledge distillation we do have the capacity to train a large model
[Table]
but still want to deploy a small one.
We consider what we call self-taught knowledge distillation, in which a network is over-parameterized in such a way that an equivalent small model can be directly recovered.
Previously, this has been studied for small models (Arora et al., 2018;
Guo et al., 2020) or using convolution-speciﬁc methods (Cao et al., 2020).
We take a straightforward 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.
We study 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 of ResNet and train the self-taught networks using the same routine as for the base model, except when replacing weight-decay by FD.
We do not study SI here, simply using the default initialization for U , V T and M (cid:18) N p I m , 0 .
01 2 q .
Note that the approach does seem sensitive to initialization, so we report the median of three trials.
Table 3 shows the strength of this approach: on both CIFAR datasets we are able to train a ResNet that beats ResNet56 and a ResNet56 that beats ResNet110, with both cases needing only half the parameters at inference.
Often we can even do so while using about as many parameters at training time as the larger model.
The experiments also show that FD is critical: without it, self-taught distillation performs worse than the baseline model, and especially poorly in the “deep” setting.
A detailed visualization of this result, 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.
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.
[Figure]
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 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 to improve optimization.
For example, SI as deﬁned can be applied to any other 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 the individual factors.
The success of SI and FD for both low-rank and tensor decomposition also suggests that these schemes may be useful when using other types of factorized neural layers such as ACDC (Moczulski et al., 2015) and
K-matrices (Dao et al., 2020), both deﬁned as products of structured matrices.
A 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]
2 η Tr p ˆ W Tt ˆ ∇ t q (cid:0)
O p η 2 q ; we can then take the square root of both sides and use a Tayor expansion to obtain ρ t (cid:0)
[Equation]
Then, starting from Equation 4 divided by ρ t (cid:0) 1 , substituting Equation 5, and applying a Taylor expansion yields
[Equation]
[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.
B 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.
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.
C 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.
[Table] [Table]
[Table] [Table]
test param.
(M) 0.
[Table]
depth decay (
[Table]
(cid:16)
[Table]
WD WD
FD
WD
FD
WD
FD test accuracy (%) 71.
training param.
(M) 1. test param.
(M) 1.
test accuracy (%) 70.
training param.
(M) 0.
test param.
(M) 0.
test accuracy (%) 68.
training param.
(M) 0.
test param.
(M) 0.