{"id":734179,"date":"2021-03-24T14:22:38","date_gmt":"2021-03-24T21:22:38","guid":{"rendered":"https:\/\/cm-edgetun.pages.dev\/en-us\/research\/?p=734179"},"modified":"2021-03-28T23:50:27","modified_gmt":"2021-03-29T06:50:27","slug":"factorized-layers-revisited-compressing-deep-networks-without-playing-the-lottery","status":"publish","type":"post","link":"https:\/\/cm-edgetun.pages.dev\/en-us\/research\/blog\/factorized-layers-revisited-compressing-deep-networks-without-playing-the-lottery\/","title":{"rendered":"Factorized layers revisited: Compressing deep networks without playing the lottery"},"content":{"rendered":"\n
\"\"\/<\/figure>\n\n\n\n

From BiT (928 million parameters (opens in new tab)<\/span><\/a>) to GPT-3 (175 billion parameters (opens in new tab)<\/span><\/a>), state-of-the-art machine learning models are rapidly growing in size. With the greater expressivity and easier trainability of these models come skyrocketing training costs, deployment difficulties, and even climate impact<\/a>. As a result, we\u2019re witnessing exciting and emerging research into compressing these models to make them less expensive, small enough to store on any device, and more energy efficient. Perhaps the most popular approach to model compression is pruning, in which redundant model parameters are removed, leaving only a small subset of parameters, or a subnetwork<\/em>. A major drawback of pruning, though, is it requires training a large model first, which is expensive and resource intensive.<\/p>\n\n\n\n

Recent research around pruning has focused on the lottery ticket hypothesis (opens in new tab)<\/span><\/a>, which suggests that most parameters of large models are redundant even during training <\/em>and that there exists a subnetwork responsible for most of the model\u2019s performance. In other words, you can train this subnetwork alone and still obtain the same accuracy as if you had trained all the parameters\u2014provided you are able to identify the subnetwork first. While such subnetworks\u2014called lottery tickets<\/em>\u2014have been shown to exist by pruning large networks after <\/em>training, intensive efforts to design methods that guess them before <\/em>training, such as GraSP (opens in new tab)<\/span><\/a>and hybrid tickets (opens in new tab)<\/span><\/a>, haven\u2019t achieved the same performance.<\/p>\n\n\n\n

Are sparsity-based methods such as pruning and guessing lottery tickets the right way to obtain compressed models? As part of our paper \u201cInitialization and Regularization of Factorized Neural Layers<\/a>,\u201d <\/span>which we\u2019re presenting at the International Conference on Learning Representations<\/a> (ICLR 2021), we revisit the alternative compression approach of factorized neural layers (opens in new tab)<\/span><\/a>. While sparsity-based approaches remove parameters from the weight matrices of a network incrementally, factorization replaces the matrices with products of smaller matrices that are more efficient to store and compute. Despite the fact that factorized layers are more amenable to deployment on deep learning accelerators such as GPUs and software frameworks such as PyTorch, sparsity-based methods remain more popular because they\u2019re seen as better at maintaining high accuracy at high compression rates (opens in new tab)<\/span><\/a> (for example, 10 percent or less parameters remaining).<\/p>\n\n\n\n

Our results contradict this position: we show that if we use the right initialization scheme (spectral initialization,<\/em> or SI<\/em>) and the right regularization penalty (Frobenius decay<\/em>, or FD<\/em>), we can achieve higher accuracy on three benchmark datasets by training a factorized ResNet from scratch than by pruning or guessing lottery tickets. The key principle underlying these two natural methods, neither of which requires extra hyperparameters, is that the training behavior of a factorized model should mimic that of the original (unfactorized) network.<\/p>\n\n\n\n

We further demonstrate the usefulness of these schemes in two settings beyond model compression where factorized neural layers are applied. The first is an exciting new area of knowledge distillation in which an overcomplete factorization is used to replace the complicated and expensive student-teacher training phase with a single matrix multiplication at each layer. The second is for training Transformer-based architectures such as BERT (opens in new tab)<\/span><\/a>, which are popular models for learning over sequences like text and genomic data and whose multi-head self-attention mechanisms are also factorized neural layers.<\/p>\n\n\n\n

Our work is part of Microsoft Research New England\u2019s<\/a> AutoML research efforts<\/a>, which seek to make the exploration and deployment of state-of-the-art machine learning easier through the development of models that help automate the complex processes involved. The code to our work is available at our GitHub repository (opens in new tab)<\/span><\/a>.<\/p>\n\n\n\n

Factorized neural layers and where to find them<\/h2>\n\n\n\n

Deep networks are function approximators in which inputs are passed through a sequence of nonlinear transformations \\(g(W,x)\\), each modifying the previous layer\u2019s output \\(x\\) using some mapping specified by a weight matrix \\(W\\)\u2208 \\(\\mathbb{R}\\)\\(^m\\)\\(^x\\)\\(^n\\). The space and time complexity of each layer is typically tied directly to the dimensionality of this weight matrix; for example, standard implementations of fully connected and convolutional layers require \\(O(mn)\\) operations to compute \\(g(W, x)\\).<\/p>\n\n\n\n

\"Two<\/a>
Figure 1: An example of how to factorize a two-dimensional convolution layer. Here the weight tensor of a regular k <\/em>\u00d7 k<\/em> convolution (left) is reshaped into a matrix and factorized into two smaller weight matrices, each corresponding to a one-dimensional convolution performed along one dimension of an image (right).<\/figcaption><\/figure><\/div>\n\n\n\n

Factorized neural layers replace the matrix with a product of matrices; in the simplest case, we have the standard low-rank decomposition \\(W\\)=\\(UV\\)\\(^T\\) for matrices \\(U\\)\u2208 \\(\\mathbb{R}\\)\\(^m\\)\\(^x\\)\\(^r\\), \\(V\\)\u2208 \\(\\mathbb{R}\\)\\(^n\\)\\(^x\\)\\(^r\\). In the model compression case, we can set \\(r\\)\u226a\\(m,n\\) to improve the complexity of fully connected and convolutional layers from \\(O(mn)\\) to \\(O(r(m+n))\\); as shown in Figure 1, for two-dimensional convolutions, this speedup is achieved using two one-dimensional convolutions, one along each spatial dimension. In the case of Transformers, which take as input a sequence \\(x\\) with hidden dimension \\(d\\), we can express the action of a multi-head attention (MHA) operation parameterized by \\(H\\) query \\((Q_h)\\), key \\((K_h)\\), value \\((V_h)\\), and output \\((O_h)\\) projection matrices as the following summation over its heads \\(h\\)=1,\u2026,\\(H\\):<\/p>\n\n\n\n

\"\"<\/figure>\n\n\n\n

This formulation makes clear that MHA is an aggregation over \\(2H\\) factorized neural layers with rank \\(r=d\/H\\), half of them parameterized by \\(Q_h\\) \\(Kh^T\\) and the other half by \\(V_h\\) \\(O_h^T\\).<\/p>\n\n\n\n

How should training routines handle factorized neural layers?<\/h2>\n\n\n\n

It\u2019s straightforward to apply standard deep network training algorithms such as stochastic gradient descent (SGD) to networks with factorized layers. However, modern optimization procedures have many critical aspects, such as initialization and regularization, that significantly influence the performance of the final model. When we factorize a network, the effect of these components can change substantially, so we often can\u2019t use the same settings as for the unfactorized model. For example, in the unfactorized case, the standard weight decay regularization penalizes the squared Frobenius norm (opens in new tab)<\/span><\/a> ||\\(W||_{Fro}^{2}\\) of each weight matrix in the model; however, if we directly apply this to the factors \\(U\\) and \\(V\\) in a factorized layer, we end up penalizing an upper bound on twice the nuclear norm (opens in new tab)<\/span><\/a> of their product:<\/p>\n\n\n\n

\"\"<\/figure>\n\n\n\n

As indicated in the figure below, this upper bound is tight when training a factorized ResNet using regular weight decay, showing that it effectively penalizes the nuclear norm rather than the Frobenius norm.<\/p>\n\n\n\n

\"Three<\/a>
Figure 2: Comparison of nuclear norm (solid line) and nuclear norm upper bound penalized by weight decay on individual factors (dotted line) during the training of ResNet20 on CIFAR-10, showing that for most of training, weight decay is effectively penalizing the nuclear norm.<\/figcaption><\/figure>\n\n\n\n

This example demonstrates how factorizing the model without modifying the training routine can lead to substantial changes in training behavior. Thus, following the principle that the training behavior of a factorized model should mimic that of the original model to recover the latter\u2019s performance, we argue in favor of using spectral initialization (SI) and Frobenius decay (FD) when training factorized neural layers. SI initializes \\(U\\) and \\(V\\) by applying a rank-deficient singular value decomposition (SVD) to the standard random initialization of the corresponding full-rank weight matrix \\(W\\) in the original model, while FD replaces weight decay by penalizing the squared Frobenius norm \u2016\\(UV\\)\\(^T\\)\u2016\\(_{Fro}^{2}\\) rather than the sum of squared norms of the individual factors.<\/p>\n\n\n\n

In our paper, we extend a recent approximate analysis of weight decay (opens in new tab)<\/span><\/a> to show that these two schemes allow SGD to maintain a high effective step size when training factorized neural layers, mimicking the role of regular weight decay applied to the original models. More informally, SI follows the mimicking principle since the SVD provably returns the best factorized approximation of initialization used by the original model. FD follows the mimicking principle by regularizing the squared Frobenius norm of the product, as is done by weight decay applied to the original model; in contrast, applying weight decay to the individual factors implicitly regularizes the nuclear norm, as shown in the figure above. Together, these two simple, efficient modifications lead to substantial improvements in model compression, knowledge distillation, and the training of Transformer-based architectures.<\/p>\n\n\n\n

Factorization vs. sparsity-based methods for model compression<\/h2>\n\n\n\n

As mentioned above, sparsity-based approaches such as pruning and guessing lottery tickets represent the prevailing trend in model compression. However, in this section, we show that low-rank factorization is competitive with these methods and can even attain better accuracy at the same parameter count. In particular, the figure below demonstrates that, in the case of low-memory ResNets, factorized neural layers outperform not only sparse (low-memory) training methods like guessing lottery tickets and dynamic sparsity (opens in new tab)<\/span><\/a> but often even full model training followed by pruning. The model\u2014a modified ResNet32\u2014and datasets\u2014CIFAR-10, CIFAR-100, and Tiny-ImageNet\u2014we use reflect the same standard setup considered by the lottery ticket papers GraSP and hybrid tickets.<\/p>\n\n\n\n

\"Three<\/a>
Figure 3: Comparison of the accuracy of several training methods: full model training (black); the sparsity-based approaches full model pruning, in which the full model is trained and then pruned (blue with black dots), and sparse training, in which a sparse model is trained (blue); and low-rank factorized layers (orange). When properly initialized and regularized, low-rank training, particularly with spectral initialization and Frobenius decay, achieves better performance at 10 percent compression than sparsity-based methods.<\/figcaption><\/figure>\n\n\n\n

This superior performance depends critically on the use of spectral initialization and Frobenius decay in tandem; interestingly, in the paper, we find that both schemes decrease <\/em>accuracy when used independently. While factorization isn\u2019t always best at very high compression rates, the figure above shows that it\u2019s clearly superior in the standard lottery ticket regime, when the compressed model accuracy is close to that of the original model (usually this means the compressed model has 10 percent or more of the original parameters). Notably, our approach doesn\u2019t require any additional tuning, as the decay coefficient used for the uncompressed model is the same one used by FD. Thus, factorized neural layers serve as a strong, simple baseline regardless of whether we\u2019re targeting memory savings or fast computation.<\/p>\n\n\n\n

Teacher-free teaching with overcomplete knowledge distillation<\/h2>\n\n\n\n

While our motivating application is model compression, factorized neural layers can also be found in other applications, such as knowledge distillation, a field that studies how to obtain a better small model by \u201cteaching\u201d it using a more powerful large model. Recent work on using overparameterization to train compact networks (opens in new tab)<\/span><\/a> suggests that factorized neural layers can help us avoid the (expensive) two-stage student-teacher training process of knowledge distillation by using an overcomplete <\/em>factorization of each of the student\u2019s weight matrices \\(W\\)=\\(UV\\)\\(^T\\), where matrix factors \\(U\\)\u2208 \\(\\mathbb{R}\\)\\(^m\\)\\(^x\\)\\(^r\\),\\(V\\)\u2208 \\(\\mathbb{R}\\)\\(^n\\)\\(^x\\)\\(^r\\) have inner dimension \\(r\\)\u2265\\(m,n\\), making the factors wider than the original weight matrix, or even by using a deep factorization \\(W\\)=\\(UMV\\)\\(^T\\). The goal is to take advantage of the training benefits of deeper or wider networks suggested by recent theoretical work (opens in new tab)<\/span><\/a> while not actually increasing the model capacity. While this increases the parameter count during training, afterward we can directly obtain a network of the original size by multiplying the factors together. We refer to this technique as overcomplete knowledge distillation<\/em>.<\/p>\n\n\n\n

\"Amn$ parameters). An arrow labeled \u201cmultiply back before inference\u201d points to a square labeled \u201cUV superscript T\u201d with the caption \u201c(mn parameters).\u201d\" class=\"wp-image-735718\"\/><\/a>
Figure 4: Comparison between the model compression (top) and knowledge distillation (bottom) pipelines. While both train a factorized model, in overcomplete knowledge distillation, the factors are remultiplied after training to bring the parameter count back to that of a normal model. This is a teacher-free way of taking advantage of the training benefits of larger models without suffering their larger inference cost.<\/figcaption><\/figure>\n\n\n\n

Since spectral initialization isn\u2019t applicable when the decomposition is overcomplete, we investigate the effect of Frobenius decay alone. We find that this regularization is critical for overcomplete knowledge distillation using ResNets trained on CIFAR data; in fact, overparameterizing and applying regular weight decay decreases <\/em>model accuracy. On the other hand, using FD, we can train an overparameterized ResNet56 (1.7 million parameters during training\/850,000 parameters at inference time) that matches the performance of ResNet110 (1.7 million parameters at both training and inference). Furthermore, training the overcomplete ResNet56 is 1.5 times faster than training the regular ResNet110. These results are the first successful application of overcomplete knowledge distillation for large-depth neural networks on CIFAR.<\/p>\n\n\n\n

Factorization-aware training of Transformers<\/h2>\n\n\n\n

Finally, we apply spectral initialization and Frobenius decay in the MHA module of Transformer-based architectures. While we can show that this indeed helps when training such models using regular SGD, large-scale unsupervised models, such as BERT, are usually trained using variants of adaptive algorithms such as Adam (opens in new tab)<\/span><\/a>or LAMB (opens in new tab)<\/span><\/a>that \u201cdecouple (opens in new tab)<\/span><\/a>\u201d weight decay by defining it as an operation that subtracts a small constant times the parameter from itself before taking a gradient step. While the equivalence breaks for adaptive methods, in the case of SGD, this is equivalent to adding a squared Frobenius penalty term to the objective. Thus for adaptive algorithms, we can devise a similar \u201cdecoupled\u201d FD scheme that subtracts a small constant times \\(UV\\)\\(^T\\) \\(V\\)=\u2207\\(_U\\)\\(_{2}^{1\/}\\) \u2016\\(UV\\)\\(^T\\) \u2016 \\(_{Fro}^{2}\\) and \\(VU\\)\\(^T\\) \\(U\\)=\u2207\\(_V\\)\\(_{2}^{1\/}\\) \u2016\\(UV\\)\\(^T\\) \u2016\\(_{Fro}^{2}\\) from the factors \\(U\\) and \\(V\\), respectively.<\/p>\n\n\n\n

\"diagram\"<\/a>
Figure 5: A diagram of how multi-head self-attention implicitly consists of 2H<\/em> factorized neural layers. Specifically, multi-head attention is a sum over H<\/em> attention heads (orange), each a matrix product of two terms: the first produced by the low-rank bilinear form Qh<\/sub> Kh<\/sub>T<\/sup><\/em> and the second by the low-rank linear transform Vh<\/sub> Oh<\/sub>T<\/sup><\/em>. <\/figcaption><\/figure><\/div>\n\n\n\n

Applying FD in this manner to the MHA module when pretraining BERT-Base (110 million parameters) on unsupervised data yields a better downstream evaluation on the SQuAD task. When the MHA embedding dimension is halved (14.3 million fewer parameters), the advantage of FD continues to hold. On BERT-Large (340 million parameters), we\u2019re able to halve the MHA embedding dimension (44.2 million fewer parameters) while losing less than a point in terms of F1-score.<\/p>\n\n\n\n

Next steps in efficient model training<\/h2>\n\n\n\n

By studying how a standard initialization scheme and standard regularization scheme need to be modified to handle factorized layers, we were able to obtain better algorithms for learning efficient models, knowledge distillation, and training Transformers. We hope these results spur more investigation into how different components of standard training pipelines behave when training efficient models from scratch. For example:<\/p>\n\n\n\n