Shapeshifter: a Parameter-efficient Transformer using Factorized Reshaped Matrices

1 Introduction . Natural language models involve large number of parameters . A single encoder-decoder Transformer [ 1 ] in its base variant has about 44 million parameters , not counting the word embedding matrix , which adds another 10 million or more , depending on the chosen vocabulary or tokenization scheme . Base variant of encoder-only BERT [ 2 ] , including the embedding , has about 108 million parameters . GPT-3 [ 3 ] has about 175 billion parameters , and the largest of the Switch Transformer [ 4 ] models has 1.5 trillion . This explosion in the model size has led to increased interest in approaches for reducing the number of parameters in the model . Models with high-dimensional parameter space have much lower intrinsic dimension [ 5 ] , that is , training trajectory can be successfully restricted to a random , smaller-dimensional subspace , even ∗Work performed while at Virginia Commonwealth University . †Corresponding author . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . though training a small-parameter architecture is often less successful . These observations have been recently extended to the fine-tuning trajectories of language models [ 6 ] . Lottery ticket hypothesis [ 7 , 8 ] , recently demonstrated to hold also for language model fine-tuning [ 9 , 10 ] , shows that smaller subnetworks can be selected from a large model , re-trained in isolation , and perform as well as the large model ; what those subnetworks are is not known a prior , in absence of the trained large model , though . Some approaches for reducing model size build on this observation to train a smaller model based on an existing large model , for example a 66 million parameter DistillBERT [ 11 ] student has been distilled from 108 million parameter BERT-base [ 2 ] teacher with little loss in quality . Alternatively , distillation has been used , for example in the context of machine translation , to increase the quality of the models while maintaining their size [ 12 , 13 ] . Other approaches train a reduced-parameter model de novo , without relying on an already trained large model . For example , DeLighT [ 14 ] uses an alternative parameterization of the multi-headed self-attention based on group linear transform to reduce a 62 million parameter Transformer to 22 million . One simple way to reduce model size involves factorized matrix representations . ALBERT [ 15 ] employs a rank r decomposition of a d×nvocab embedding matrix storing d-dimensional embedding vectors for each of the nvocab tokens by using a stack of two linear layers , d× r on top of r× nvocab . Similar low-rank decomposition is also used implicitly in the multi-headed self-attention in the generic Transformer [ 1 ] with hidden dimension d and nheads self-attention heads . In each Transformer head , a r = d/nheads-rank factorized representation involving d × d/nheads key ( K ) and query ( Q ) matrices are used , with the pairwise self-attention scores for sequence x calculated using xTKTQx , instead of xTWx involving a full general attention d× d trainable matrix W as originally considered in trainable-attention encoder-decoder LSTM models [ 16 ] . In both cases , the models are trained from a random initialization of the factorized parameter space , instead of attempting to find the lowest-error factorized representation of an already trained original model . We explore here a Transformer model that uses an alternative way to decompose a matrix into two smaller matrices . Instead of standard low-rank factorization as above , it involves reshaping and reordering matrix dimensions prior to the decomposition , and is equivalent to a sum of Kronecker products with an efficient implementation . The compact model based on the smaller matrices is trained de novo , and matrices in the compact model are not aimed to be approximations of the matrices in the original model . For non-square model matrices , the approach allows for increased reduction in parameters for the same decomposition rank . For square matrices , the benefits come from increased expressiveness of the decomposition for the same rank , allowing for reducing the rank needed to preserve model accuracy , and thus reducing the model size . Our main contribution is proving that stacking multiple linear layers decomposed this way increases the expressiveness of the network , unlike stacking multiple low-rank layers factorized in the standard way , which can only map into a subspace of dimensionality equal to the rank . Empirically , we show that the decomposition can reduce the size of a simple encoder-decoder Transformer to as little as 4 million parameters , including 2 million for the model and 2 million for the embeddings . The technique can be employed automatically to any matrix , including embedding , dense , attention , and output layers , without requiring any modification of the model architecture . 2 Related Work . The two approaches most closely related to ours are word2ket embeddings [ 17 ] and parameterized hypercomplex multiplication ( PHM ) linear layers [ 18 ] . Word2ket represents input embedding matrix in language models via a sum of Kronecker products , each involving two smaller matrices . It achieves above 100-fold reduction in the number of trainable embedding parameters ; however , their approach relies on lazy tensors to efficiently extract embeddings for a limited number of words in the input sequence , and does not extend to output embedding matrix or other model matrices that perform general linear transformations of their input vectors . PHM linear layers [ 18 ] arise from generalizing hypercomplex number layers [ 19 ] to arbitrary dimensionality r , with the arithmetic over the numbers learned during training . For a given dimensionality r , the arithmetic takes form of a sum of r Kronecker products . Irrespective of the original matrix size n×m , each Kronecker product involves a small r × r matrix and a larger n/r × m/r matrix , which together have r3 + nm/r trainable parameters . For r used in practice , the second term dominates , resulting in r-fold reduction in size . PHM layers with r as large as 16 have been demonstrated to be effective , leading to a reduction of a 44 million parameter Transformer to a 2.9 million parameter PHM-Transformer . These numbers do not include the embedding matrix for token embeddings which is not reduced in size by the PHM approach ; and which for example for 32 , 000 BPE tokens requires about 16 million parameters . Both of the Kronecker-based compact language model approaches [ 17 , 18 ] , as well as Kroneckerbased convolutional [ 20 ] and recurrent networks [ 21 ] , demonstrate that a sum of Kronecker products leads to very compact representations . Here , we advance theoretical understanding of why this low-rank representation has advantages in terms of expressivity compared to a standard matrix factorization with the same rank . We also provide a more efficient , more general Kronecker-based representation . Unlike word2ket , it can be used in linear transformations anywhere in the model , an unlike PHM it also applies to embedding matrices . To represent an n×m matrix , our approach requires 2r √ mn instead of r3 + nm/r parameters required by PHM . For example , for r = 16 , the highest used by PHM , for a Transformer with hidden dimension of 512 , each 512×2048 feed-forward matrix in the self-attention head would be reduced from 1 million parameters to 69,632 parameters by PHM , but to 16,384 by our approach . Kronecker product of two matrices is a concrete-basis representation of a finite-dimensional linear operator defined by a tensor product of two underlying linear operators . Thus , Kronecker-productbased representations , including PHM layers and ours , can be seen as special case of a more general family of higher-order tensor-product-based representations . Kronecker-based representation unfolds a n×m matrix from two smaller matrices , n1 ×m1 and n2 ×m2 , with m = m1m2 and n = n1n2 . In tensor-based representations , this is generalized to m = ∏o j=1mj and n = ∏o j=1 nj for some tensor order o , leading to unfolding the n×m matrix from an o-way tensor that is then decomposed using low rank leading to reduction in parameters . For example , the tensor-train representation [ 22 ] represents a matrix using a series of core tensors with matching ranks . Tensor decomposition has been used successfully in convolutional networks [ 23 , 24 ] and recurrent networks [ 25 ] . In the context of language models , tensor-train representation has been used to construct tensorised embeddings [ 26 ] , which use order-three tensor representation to obtain 60-fold reduction in embedding in a sequenceto-sequence Transformer , and also use order-six tensors to achieve almost 400-fold reduction in an LSTM model for sentiment analysis . Similarly , word2ket [ 17 ] uses order-four tensor representation to achieve more than 90,000-fold reduction in the embedding size for question-answering LSTM model . Beyond embeddings , Tensorized Transformer [ 27 ] uses block-term tensor decomposition – a combination of CP and Tucker decompositions – to reduce the multi-headed self-attention layer with nheads by a factor 1/nheads , leaving other layers intact ; this allowed for reducing a 52M Transformer down to 21M trainable parameters , with little drop in BLEU scores on WMT ’ 16 En-De translation task compared to original Transformer . So far , none of the Kronecker- or tensor-productbased approaches have been applied comprehensively to all components of a Transformer model – embeddings , attention , and feed-forward layers . 3 Limitations of Factorized Matrices in Deep Networks . A simple approach for reducing the size and complexity of a model involves replacing each linear transformation n×mmatrixW with its factorized , low-rank representation involving an n×r matrix A and an r ×m matrix B , with the rank r m , n . During training and inference , the product AB is used instead of W , essentially replacing W with a two linear layers , A stacked on B . In this way , the number of parameters reduces from nm needed to represent W to r ( n+m ) needed to represent A and B . One potential approach to use the factorized representation is to take an already trained model , and factorize each matrix Wj into its low-rank approximation Ŵj = AjBj using a method such as SVD . However , analyzing an idealized deep model with linear activations shows that errors from the layers add up [ 28 ] , ‖ ∏ jWj − ∏ j Ŵj‖ = ∑ j‖Wj − Ŵj‖ . Instead , a typical approach , used also here , is to train a model de novo in its factorized form , using randomly initialized matrices Aj and Bj , without attempting to approximate individual linear transformations of the non-factorized model . Representing an n × m matrix W as a low-rank product AB limits the linear transformation – the output is constrained to be a low-dimensional subspace of Rn . Stacking rank-r factorized linear transformations into an idealized deep network with linear activations does not improve the expressiveness , the dimensionality of the space of possible outputs is still constrained by the smallest of the decomposed layers rank . While this may have regularizing effect for medium values of the rank , it may prevent from using very small ranks and thus from creating compact models . 4 Stacked Kronecker Product-based Representations . For n = m , the standard factorized ( n× r ) ( r ×m ) → n × m representation is optimal among pairwise-product-based representations in the sense that all elements of a column of A are multiplied with each element of a row of B . For non-square matrices , for example matrix A is smaller than B , a more efficient all-pairs product representation can be made by increasing the size of A and decreasing the size of B . In Shapeshifter , instead of the standard ( n× r ) ( r ×m ) → n ×m factorized representation , we rearrange the mn parameters of the original matrix W into a √ mn× √ mn matrix3 W ′ , and then use a factorized representation of W ′ as a √ mn× r matrix A and an r × √ mn matrix B : ( √ nm× r ) ( r × √ nm ) multiply−−−−→ √nm×√nm reshape−−−−→ √n×√m×√n×√m transpose−−−−−→ √ n× √ n× √ m× √ m reshape−−−−→ n×m This representation uses 2r √ mn parameters instead of nm . We apply this representation to all matrices in attention and feed-forward layers in all multi-head attention blocks , and to the embedding matrices . For non-square matrices , the representation results in parameter saving compared to standard factorization with the same rank r. For square matrices , no parameter saving is achieved . However , as we show below , the way the matrix is decomposed allows for increased expressive power . For r = 1 , the n×m matrix W arises from taking products of each element from √ nm× 1 matrix A with each element from 1 × √ nm matrix B . It is known to be equivalent [ 29 ] to a Kronecker product A⊗B involving √ n× √ m matrices A , B resulting from reshaping the single-column and single-row matrices , respectively . Beyond r = 1 , the representation is equivalent to representing W as sum of r Kronecker products , W = ∑r j=1Aj ⊗ Bj . Compared to the equivalent sum of Kronecker-products representation , using the matrix factorization representation avoids the explicit summation of r terms . In current libraries , it results in a simpler implementation involving only matrix multiplication , reshaping and transposing , and is more efficient compared to performing individual Kronecker products followed by an explicit summation along the rank dimension .

Summary: This paper presents an approach called ShapeShifter to compress large language models. The paper follows the approach of reducing number of parameters by using factorized matrix representations. The primary idea is to represent the matrix as a sum of Kronecker products, instead of standard low-rank factorizations. The paper theoretically shows that for non-square matrices, the proposed approach allows for increased reduction in parameters for the same decomposition rank. For square matrices, it allows for increased expressiveness of the decomposition for the same rank. Experimental results on machine translation tasks show that the proposed approach performs much better than the state-of-the-art model compression approaches and is just slightly worse than the large models.