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: Shapeshifter reparameterizes dense kernels into a product of two lower-rank matrices. Shapeshifter introduces reshapes and transposes to make this reparameterization more efficient. The paper provides a theoretical justification for its expressivity and strong experimental results in machine translation.