summaries: 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. This paper proposes a method to reduce the number of parameters by weight factorization. Unlike traditional low-rank factorization (n x r and r x m), this work decomposes the matrix into two sqrt(mn) x r matrices. The algorithm includes the reordering of elements to recover the desired shape. The authors find connections to Kronecker products and theoretically support the expressiveness of the proposed method. Experiments on NMT tasks show that Shapeshifter can achieve better parameter efficiency (fewer parameters with comparable performance). The results are compared to two similar works: PHM and DeLighT. 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. A sizable effort has been made in recent history to obtain the accuracy of very large language models at a fraction of the parameter of the large models. Broadly these methods fall into various categories like model compression, factorized representation, distillation etc. This paper underlines another such attempt at obtaining a factorized representation of matrices to reduce the parameter space. The authors justify the use of Kronecker products to factorize parameter space, by showing that stacked layers of low rank matrices increase expressiveness. They also show the effectiveness of their proposed method by achieving a 4 to 14-fold reduction in parameters without sacrificing accuracy. This paper proposes an algorithm they call ‘Shapeshifter’ to reduce the number of parameters needed for Transformer-based models while preserving expressiveness and performance. The proposed algorithm is based on the idea of low-rank factorizations of weight matrices, but uses a different factorization method involving sums of Kronecker products that leads to greater expressiveness. They show theoretical analysis to prove their factorization can accurately represent any $n \times m$ matrix or stack of matrices (in an idealized case), and experimental results using transformer-based models for translation on English-German, English-French and English-Romanian. They demonstrate superior or comparable performance with vastly fewer parameters compared to some other model compression techniques.
abstractText: Language models employ a very large number of trainable parameters. Despite being highly overparameterized, these networks often achieve good out-of-sample test performance on the original task and easily fine-tune to related tasks. Recent observations involving, for example, intrinsic dimension of the objective landscape and the lottery ticket hypothesis, indicate that often training actively involves only a small fraction of the parameter space. Thus, a question remains how large a parameter space needs to be in the first place — the evidence from recent work on model compression, parameter sharing, factorized representations, and knowledge distillation increasingly shows that models can be made much smaller and still perform well. Here, we focus on factorized representations of matrices that underpin dense, embedding, and self-attention layers. We use low-rank factorized representation of a reshaped and rearranged original matrix to achieve space efficient and expressive linear layers. We prove that stacking such low-rank layers increases their expressiveness, providing theoretical understanding for their effectiveness in deep networks. In Transformer models, our approach leads to more than tenfold reduction in the number of total trainable parameters, including embedding, attention, and feed-forward layers, with little degradation in on-task performance. The approach operates out-of-the-box, replacing each parameter matrix with its compact equivalent while maintaining the architecture of the network.
authors.name: Aliakbar Panahi
authors.name: Seyran Saeedi
authors.name: Tom Arodz
references.authors: Ashish Vaswani Noam Shazeer Niki Parmar Jakob Uszkoreit Llion Jones Aidan N Gomez Łukasz Kaiser Illia Polosukhin
references.title: Attention is all you need
references.venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,
references.authors: Jacob Devlin Ming-Wei Chang Kenton Lee Kristina Toutanova
references.title: BERT: pre-training of deep bidirectional transformers for language understanding
references.venue: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies,
references.authors: Tom Brown Benjamin Mann Nick Ryder Melanie Subbiah Jared D Kaplan Prafulla Dhariwal Arvind Neelakantan Pranav Shyam Girish Sastry Amanda Askell Sandhini Agarwal Ariel Herbert-Voss Gretchen Krueger Tom Henighan Rewon Child Aditya Ramesh Daniel Ziegler Jeffrey Wu Clemens Winter Chris Hesse Mark Chen Eric Sigler Mateusz Litwin Scott Gray Benjamin Chess Jack Clark Christopher Berner Sam McCandlish Alec Radford Ilya Sutskever Dario Amodei
references.title: Language models are few-shot learners
references.venue: Advances in Neural Information Processing Systems,
references.authors: William Fedus Barret Zoph Noam Shazeer
references.title: Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity
references.venue: arXiv preprint arXiv:2101.03961,
references.authors: Chunyuan Li Heerad Farkhoor Rosanne Liu Jason Yosinski
references.title: Measuring the intrinsic dimension of objective landscapes
references.venue: In International Conference on Learning Representations,
references.authors: Armen Aghajanyan Sonal Gupta Luke Zettlemoyer
references.title: Intrinsic dimensionality explains the effectiveness of language model fine-tuning
references.venue: In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers),
references.authors: Jonathan Frankle Michael Carbin
references.title: The lottery ticket hypothesis: Finding sparse, trainable neural networks
references.venue: In International Conference on Learning Representations,
references.authors: Jonathan Frankle Gintare Karolina Dziugaite Daniel Roy Michael Carbin
references.title: Linear mode connectivity and the lottery ticket hypothesis
references.venue: In International Conference on Machine Learning,
references.authors: Tianlong Chen Jonathan Frankle Shiyu Chang Sijia Liu Yang Zhang Zhangyang Wang Michael Carbin
references.title: The lottery ticket hypothesis for pre-trained BERT networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Sai Prasanna Anna Rogers Anna Rumshisky
references.title: When BERT plays the lottery, all tickets are winning
references.venue: In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing (EMNLP),
references.authors: Victor Sanh Lysandre Debut Julien Chaumond Thomas Wolf
references.title: DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter
references.authors: Bei Li Ziyang Wang Hui Liu Quan Du Tong Xiao Chunliang Zhang Jingbo Zhu
references.title: Learning light-weight translation models from deep transformer
references.venue: Proceedings of the AAAI Conference on Artificial Intelligence,
references.authors: Fusheng Wang Jianhao Yan Fandong Meng Jie Zhou
references.title: Selective knowledge distillation for neural machine translation
references.venue: In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers),
references.authors: Sachin Mehta Marjan Ghazvininejad Srinivasan Iyer Luke Zettlemoyer Hannaneh Hajishirzi
references.title: DeLighT: Very deep and light-weight transformer
references.venue: In International Conference on Learning Representations ICLR’2021,
references.authors: Zhenzhong Lan Mingda Chen Sebastian Goodman Kevin Gimpel Piyush Sharma Radu Soricut
references.title: ALBERT: A lite BERT for self-supervised learning of language representations
references.venue: In International Conference on Learning Representations
references.authors: Minh-Thang Luong Hieu Pham Christopher D Manning
references.title: Effective approaches to attentionbased neural machine translation
references.venue: In Proceedings of the 2015 Conference on Empirical Methods in Natural Language Processing,
references.authors: Aliakbar Panahi Seyran Saeedi Tom Arodz
references.title: word2ket: Space-efficient word embeddings inspired by quantum entanglement
references.venue: In International Conference on Learning Representations
references.authors: Aston Zhang Yi Tay SHUAI Zhang Alvin Chan Anh Tuan Luu Siu Hui Jie Fu
references.title: Beyond fully-connected layers with quaternions: Parameterization of hypercomplex multiplications with 1/n parameters
references.venue: In International Conference on Learning Representations ICLR’2021,
references.authors: Yi Tay Aston Zhang Anh Tuan Luu Jinfeng Rao Shuai Zhang Shuohang Wang Jie Fu Siu Cheung Hui
references.title: Lightweight and efficient neural natural language processing with quaternion networks
references.venue: In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics,
references.authors: Shuchang Zhou Jia-Nan Wu Yuxin Wu Xinyu Zhou
references.title: Exploiting local structures with the Kronecker layer in convolutional networks
references.venue: arXiv preprint arXiv:1512.09194,
references.authors: Cijo Jose Moustapha Cissé Francois Fleuret
references.title: Kronecker recurrent units
references.venue: In International Conference on Machine Learning,
references.authors: Ivan V Oseledets
references.title: Tensor-train decomposition
references.venue: SIAM Journal on Scientific Computing,
references.authors: Emily Denton Wojciech Zaremba Joan Bruna Yann LeCun Rob Fergus
references.title: Exploiting linear structure within convolutional networks for efficient evaluation
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Alexander Novikov Dmitry Podoprikhin Anton Osokin Dmitry Vetrov
references.title: Tensorizing neural networks
references.venue: In Advances in Neural Information Processing Systems,
references.authors: Charles C Onu Jacob E Miller Doina Precup
references.title: A fully tensorized recurrent neural network
references.venue: arXiv preprint arXiv:2010.04196,
references.authors: Oleksii Hrinchuk Valentin Khrulkov Leyla Mirvakhabova Elena Orlova Ivan Oseledets
references.title: Tensorized embedding layers
references.venue: In Findings of the Association for Computational Linguistics: EMNLP 2020,
references.authors: Xindian Ma Peng Zhang Shuai Zhang Nan Duan Yuexian Hou Ming Zhou Dawei Song
references.title: A tensorized transformer for language modeling
references.venue: Advances in Neural Information Processing Systems,
references.authors: Thomas Frerix Joan Bruna
references.title: Approximating orthogonal matrices with effective Givens factorization
references.venue: In International Conference on Machine Learning,
references.authors: Harold V Henderson Shayle R Searle
references.title: The vec-permutation matrix, the vec operator and Kronecker products: A review
references.venue: Linear and multilinear algebra,
references.authors: Charles F Van Loan
references.title: The ubiquitous Kronecker product
references.venue: Journal of Computational and Applied Mathematics,
references.authors: Charles F Van Loan Nikos Pitsianis
references.title: Approximation with Kronecker products
references.authors: Julie Kamm James G Nagy
references.title: Kronecker product and SVD approximations in image restoration
references.venue: Linear Algebra and its Applications,
references.authors: Eugene Tyrtyshnikov
references.title: Kronecker-product approximations for some function-related matrices
references.venue: Linear Algebra and its Applications,
references.authors: Wallace Givens
references.title: Computation of plain unitary rotations transforming a general matrix to triangular form
references.venue: Journal of the Society for Industrial and Applied Mathematics,
references.authors: Colin Raffel Noam Shazeer Adam Roberts Katherine Lee Sharan Narang Michael Matena Yanqi Zhou Wei Li Peter J. Liu
references.title: Exploring the limits of transfer learning with a unified text-to-text transformer
references.venue: Journal of Machine Learning Research,
references.authors: Ond Onrej Bojar Christian Federmann Mark Fishel Yvette Graham Barry Haddow Matthias Huck Philipp Koehn Christof Monz
references.title: Findings of the 2018 Conference on Machine Translation (WMT18)
references.venue: In Proceedings of the Third Conference on Machine Translation,
references.authors: Mauro Cettolo J Niehues S Stüker Luisa Bentivogli Marcello Federico
references.title: Report on the 11th IWSLT evaluation campaign, IWSLT
references.venue: In IWSLT-International Workshop on Spoken Language Processing,
references.authors: Rico Sennrich Barry Haddow Alexandra Birch
references.title: Neural machine translation of rare words with subword units
references.venue: In Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers),
references.authors: Kishore Papineni Salim Roukos Todd Ward Wei-Jing Zhu
references.title: BLEU: a method for automatic evaluation of machine translation
references.venue: In Proceedings of the 40th annual meeting of the Association for Computational Linguistics,
references.authors: Thomas Wolf Lysandre Debut Victor Sanh Julien Chaumond Clement Delangue Anthony Moi Pierric Cistac Tim Rault Rémi Louf Morgan Funtowicz
references.title: HuggingFace’s Transformers: State-of-the-art natural language processing
references.authors: Adam Paszke Sam Gross Francisco Massa Adam Lerer James Bradbury Gregory Chanan Trevor Killeen Zeming Lin Natalia Gimelshein Luca Antiga
references.title: Pytorch: An imperative style, high-performance deep learning library
references.venue: Advances in Neural Information Processing Systems,
references.authors: Yang You Jing Li Sashank Reddi Jonathan Hseu Sanjiv Kumar Srinadh Bhojanapalli Xiaodan Song James Demmel Kurt Keutzer Cho-Jui Hsieh
references.title: Large batch optimization for deep learning: Training BERT in 76 minutes
references.venue: In International Conference on Learning Representations ICLR’2020,
sections.heading: 1 Introduction
sections.text: 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.
sections.heading: 2 Related Work
sections.text: 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.
sections.heading: 3 Limitations of Factorized Matrices in Deep Networks
sections.text: 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.
sections.heading: 4 Stacked Kronecker Product-based Representations
sections.text: 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.
sections.heading: 4.1 Expressiveness of stacked Kronecker-product Layers
sections.text: Certain types of matrices can easily be expressed as a sum of Kronecker products, but not in a r-rank factorized matrix form. For example, identity mapping on Rn does not admit low-rank factorization, but it can be decomposed as a single Kronecker product of two smaller identity matrices. More broadly, the output of a sum of Kronecker products layer is not limited to r-dimensional subspace; a Kronecker product of two orthogonal matrices is an orthogonal matrix [30], thus even a single A⊗B term can model isomorphisms Rn → Rn. The converse is not true, most matrices require Kronecker representation with a high rank, and the question of approximating various types of matrices with a sum of Kronecker products has received ample attention [31, 32, 30, 33].
  Here, we focus on a different question – how does the expressive power of a stack of layers, each involving a sum of Kronecker products instead of generic linear transformation, grow with the network depth. We analyze an idealized network formed by a stack of Kronecker layers with linear activations. Given high-enough rank, a single layer is enough to represent any matrix. We first analyze what the depth needs to be if rank is restricted to two, and then investigate how quickly the depth can be reduced as the rank increases.
  Below, we use the following notation. Support supp(x) of vector x is a set of indices at which x is not null. We use [n] = {1, ..., n}. I is identity matrix. 1jk is a matrix with unity at row j and column k and null elsewhere; 1k is a shorthand for 1kk. To avoid ambiguity, we use square brackets for individual entries of a matrix or a vector; thus A[j, :] represents j-th row of matrix A, while Aj represents j-th matrix out of a series, as used above. For index set Z of cardinality k and a square matrixA, byA[Z] we represent a k×k submatrix ofA obtained by taking the intersection of columns and rows from set Z.
  We first define k-variant matrices, which will be a tool in exploring stacked layers. 3If √ mn is not an integer, we round it up, and use a submatrix as representation of W .
  Definition 1 (k-variant matrix). For an n × n orthogonal matrix U and an index set Z ⊂ [n], by writing UZ we indicate that both U and its inverse UT act as identity on all the dimensions not in set Z. That is, for each x ∈ Rn, supp(Ux− x) ⊆ Z and supp(UTx− x) ⊆ Z. If Z has cardinality at most k, we call UZ a k-variant matrix.
  The notion of k-variant matrices generalizes Givens rotation matrices [34] used for example in QR decomposition. A Givens rotation in n dimensions is any rotation that acts on a plane spanned by two of the n coordinate axes – it is thus 2-variant. QR decomposition can be used to show that each matrix can be represented as a product of 2-variant matrices. We will then show that 2-variant matrices, or more broadly, k-variant matrices for k ≤ √ n, can be represented using sum of Kronecker products.
  Theorem 2 (Decomposition into layers involving Kronecker products). Let n = m2 for somem ∈ N. Any orthogonal n× n matrix U can be represented as
  U = L∏ i=1 2∑ j=1 Aij ⊗Bij ,
  where Aij and Bij are √ n× √ n matrices, and L = O ( n2 ) .
  Proof. QR decomposition using Givens matrices will result in Q being a product of up to n(n− 1)/2 Givens rotations, which are 2-variant, and R an upper triangular matrix which for orthogonal matrices is a diagonal matrix with +1 and −1 elements. The sign of any pair of −1 diagonal entries can be negated by a 2-variant matrix; at most n/2 such matrices are needed to convert the diagonal into identity, since there are at most n− 1 negative elements in R. In total, up to n2/2 2-variant matrices are needed to represent any given orthogonal matrix.
  Next, we set to show that any 2-variant matrix can be represented as a product of at most three matrices, each of the form A1 ⊗B1 +A2 ⊗B2, using m×m matrices.
  Let U{k,q} be an arbitrary 2-variant orthogonal m2 ×m2 matrix for arbitrary k, q ∈ [m2]. That is, y = U{k,q}x acts only on x[k] and x[q] to produce y[k] and y[q], for l 6= k, q we have y[l] = x[l]. Consider two vectors, xA, xB ∈ Rm. Kronecker product of x = xA ⊗ xB of these two vectors is a vector x ∈ Rn with entries defined through products
  x[k] = xA[κ]xB [λ] for k = (κ− 1)m+ λ, x[q] = xA[π]xB [ρ] for q = (π − 1)m+ ρ
  for κ, λ, π, ρ ∈ [m].
  A matrixA⊗B acts on vectors x = xA⊗xB as (A⊗B)x = (AxA)⊗(BxB). For example, letB{λ,ρ} be a 2-variant matrix action on λ, ρ, then 1ν ⊗B{λ,ρ} is a acting on dimensions k = (ν − 1)m+ λ and q = (ν − 1)m+ ρ in the same way as B{λ,ρ} acts on dimensions λ and ρ, and producing either identity or null elsewhere, as shown below:
  [ 0
  1 0
  ] ⊗ [ α −β
  1 β α
  ] =  0 . . . α −β 1 β α . . .
  0
   .
  First, consider k, q such that κ = π = ν for some ν ∈ [m]. Define a 2-variant B{λ,ρ} such that B[{λ, ρ}] = U [{k, q}], that is, B acts on λ, ρ in the same way as U acts on k, q. We then have
  U{k,q} = 1ν ⊗B{λ,ρ} + (I − 1ν)⊗ I.
  A symmetric construction holds if λ = ρ = ν. In both cases, we can represent a 2-variant orthogonal matrix by A1 ⊗B1 +A2 ⊗B2.
  The general case of U{k,q} for κ 6= π and λ 6= ρ can be addressed by converting it into the case κ = π via transforming κ into π, applying the special case, and then transforming π back into κ.
  We show that an orthogonal matrix V of the form V = C ⊗D + C ′ ⊗D′ exists such that U{k,q} = V T ( 1π ⊗B{λ,ρ} + (I − 1ν)⊗ I ) V,
  where B[{λ, ρ}] = U [{k, q}], the λ, ρ columns/rows of B are formed by taking submatrix of U defined by columns/rows k, q.
  To construct V , set C = Pκ→π, a permutation matrix that moves κ to π, and D = 1λ. Also set C ′ = I and D′ = I − 1λ. Columns of C ⊗D and C ′ ⊗D are linearly independent, and each of the two matrices is orthogonal, thus V is orthogonal.
  Consider x = xA ⊗ xB , we then have x′ = V x = Pκ→πxA ⊗ 1λxB + IxA ⊗ (I − 1λ)xB .
  The first term will use the permutation to move xA[κ] to position π, where it will be multiplied by xB [λ], and placed at x′[k′] for k′ = (π − 1)m+ λ; all other entries xA[·]xB [λ] will be rearranged as well by the permutation. For m = 3, λ = 2, κ = 1, π = 2, the transformation that moves x[k = 2] = x[(κ − 1)m + λ] = α1β2 into x′[k′ = 5] = x′[(π − 1)m + λ] and thus brings it into position ready for the special case, is illustrated below[
  0 1 0 1 0 0 0 0 1 ][ α1 α2 α3 ] ⊗ [ 0 0 0 0 1 0 0 0 0 ][ β1 β2 β3 ] = [ α2 α1 α3 ] ⊗ [ 0 β2 0 ] =
  = [0 α2β2 0 0 α1β2 0 0 α3β2 0] T .
  The second term will preserve xA[·]xB [·] for all entries of xB except λ; this includes ρ, since in the general case λ 6= ρ. The k, q entries of x will be now at k′, q entries of x′, and applying 1π ⊗ B{λ,ρ} + (I − 1ν)⊗ I on x′ will be equivalent to acting on positions k, q of x. The inverse orthogonal transformation V T will restore the transformed values at positions k′, q, as well as all other entries xA[·]xB [λ], back to their original positions. The transformation V T is of the same form as V , except using P−1κ→π instead of Pκ→π , since the other matrices involved in V are diagonal with unit-magnitude entries. In summary, in the general case, m×m matrices Aij and Bij exist such that any 2-variant orthogonal matrix can be represented as ∏3 i=1 ∑2 j=1Aij ⊗Bij .
  Together, any orthogonal n = m2 dimensional matrix U can be decomposed into U = ∏3n2/2 i=1 ∑2 j=1Aij ⊗Bij .
  Corollary 3. Any linear n ×m matrix W can be represented as a product of a finite sequence of sums of two Kronecker products of dne × dme matrices followed by selecting an n×m submatrix.
  Proof. The result immediately follows by considering an orthogonal matrix with dimension max(dne2 × dme2), observing that an arbitrary, possibly rank-deficient diagonal matrix D can be decomposed as a product of a series of matrices of the form 1k⊗(I−(D[j, k]−1)1j)+(I−1k)⊗I , one per diagonal entry D[j, k], which allows for producing arbitrary rank-deficient square matrices, from which non-square matrices can be selected.
  The above theorem shows that stacking layers represented compactly as ( √ nm× r) (r × √ nm) for r = 2 increases the expressive power of the network. In the idealized case with linear activation function, adding such layers allows the network to eventually produce arbitrary Rm → Rn linear transformations, unlike stacking standard factorized layers (n× r) (r ×m), which always maps into an r-dimensional subspace irrespective of the number of layers.
  With rank r = 2, we may need as many as 3n2/2 layers to represent an orthogonal n× n matrix. On the other extreme, only one layer is needed if the rank is n, since arbitrary n× n matrix can be seen as a concatenation of n tiles of size √ n× √ n. As long as we define n indicator √ n× √ n matrices 1jk and use them as tile position selectors, we can use a sum of Kronecker products of the selection with the corresponding tile to produce the matrix. In between these extremes, the number of layers needed to increase the expressiveness of the network to allow arbitrary transformations depends on the rank r in the Shapeshifter representation in the following way.
  Theorem 4 (Depth-rank tradeoff for deep models involving Kronecker products). Let n = m2 for some m ∈ N, and let r ≤ m. Any orthogonal n× n matrix U can be represented as a product of a sequence of sums of two Kronecker products of m×m matrices,
  U = L∏ i=1 r∑ j=1 Aij ⊗Bij ,
  where Aij and Bij are √ n× √ n matrices, and L = O ( n2/r ) .
  Proof. Instead of 2-variant orthogonal matrices, we use r-variant ones. Consider r = 3, with U{k,q,t}, and with k = (κ− 1)m+ λ, q = (π − 1)m+ ρ, and t = (τ − 1)m+ υ. We now have two groups of indices, κ, π, τ and λ, ρ, υ. If κ = π = τ = ν for some ν, we have a special case as in Theorem 2
  U{k,q,s} = 1ν ⊗B{λ,ρ,υ} + (I − 1ν)⊗ I;
  we have a similar special case if λ = ρ = υ.
  The general case involves no equality in either of the two groups. Then, we can proceed similarly as for 2-variant general case, but we need a permutation matrix Pκ→τ and separately a permutation matrix Pπ→τ , which allow us to define
  x′ = V x = Pκ→τxA ⊗ 1λxB + Pπ→τxA ⊗ 1ρxB + IxA ⊗ (I − 1λ − 1ρ)xB .
  Here, V is an orthogonal transformation that uses a sum of three Kronecker products as a transformation that maps the general case to the κ = π = τ case. The construction naturally extends into higher r as long as r ≤ m, the number of diagonal entries in matrices A: we need V to be a sum of r Kronecker products involving distinct dimensions of A. In summary, we need less than∏3 i=1 ∑r j=1Aij ⊗Bij to represent arbitrary r-variant orthogonal matrix. A product of two k-variant matrices is at most 2k-variant. Grouping r/2 Givens rotations or, more generally, 2-variant matrices can result in at most a r-variant matrix, thus we need at most L = O ( 3n2/r ) layers, each involving a Kronecker product of rank at most r.
  As is common in theoretical work on deep networks, the results are limited to idealized multilayer models with linear activations. However, they can provide insights into networks with nonlinearities; for example, activations such as ReLU that are linear on large part of their input will lead to piece-wise linear network, and the results above hold in a piece-wise fashion.
sections.heading: 5 Experimental Results
sections.text: We validated the approach on machine translation using sequence-to-sequence models. The experiments were performed on a single V100 (longest run: 6 days) or A100 GPU (longest run: 1 day). Time overhead experiments were performed on a dedicated workstation with a single NVIDIA RTX 3090 GPU. As the basis for Shapeshifter, we used Transformer [1], with embedding and encoder/decoder layers replaced with compact representations. In Shapeshifter, we used small rank for all matrices in the multi-head self-attention blocks. For embedding matrices, which are much larger and can be reduced more effectively, we used higher rank, while aiming to keep the total size of factorized embeddings below the total size of the rest of the factorized encoder/decoder. The time overhead
  (see Table 1) introduced by the representation is modest, up to 17% increase in running time during training, and up to 13% during predictions with a trained model.
  We compared our approach to two recently proposed approaches for reducing model size: DeLighT [14] and PHM Layers [18]. DeLighT uses separate vocabulary for the source language encoder embeddings, and a different vocabulary for the target language decoder embedding and output layer. PHM Layers, like its predecessor Quaternion Transformer [19], uses the same Byte-Pair Encoding [38] in both encoder and decoder. To facilitate comparisons, we followed both approaches and trained two Shapeshifter variants, one with separate embeddings and one with single embedding.
sections.heading: 5.1 Comparison with PHM Layers
sections.text: In comparisons with PHM layers [18], we used T5-small Transformer [35], which shares a single embedding matrix in both the encoder and the decoder. We use Huggingface Transformers [40] for PyTorch [41] implementation. We evaluate the approach on IWLSLT’14 German-to-English (De-En) [37] and WMT’18 English-to-Romanian (En-Ro) [36] datasets. For the first dataset we use learning
  rate 2e-3 with batch size of 128, while for the larger En-Ro dataset we use 3e-3 with batch size of 192. We used LAMB optimizer [42] with inverse square root scheduler, dropout 0.1, no weight decay, and 0.1 label-smoothed cross entropy loss.
  To facilitate comparison of our results with PHM Layers approach under the same conditions, we reimplemented the approach and applied it to T5-Small Transformer that we used as basis for experiments with our method. The small increase in size of the PHM Transformer observed in our re-implementation of PHM layers compared to the parameter size provided in the literature [18] (3.2M vs 2.9M parameters) results from the fact that the underlying T5 Transformer that we build on uses three separate dhid × dhid matrices for self-attention query, key, and value weights for all heads, while the Transformer used in [18] stores all three as a single, concatenated dhid × 3dhid matrix. The results presented in Table 2 and in the left panel of Figure ?? indicate that while our full models are four times smaller than full PHM models, they offer very similar on-task performance. The parameter reduction comes not only from factorized embeddings in Shapeshifter, but also from matrices inside the attention blocks in encoder/decoder: both Shapeshifter 16 / 256 and PHM n = 16 use a sum of 16 Kronecker products to represent a matrix, yet PHM results in a 1.5 times larger encoder/decoder when applied to the same exact base architecture.
sections.heading: 5.2 Comparison with DeLighT
sections.text: Following DeLighT [14], we used Transformer models as implemented in fairseq [45]. We benchmarked the Shapeshifter models on four datasets: IWSLT’14 German-to-English (De-En) [37], WMT’16 English-to-Romanian (En-Ro) [43], WMT’14 English-to-German (En-De) [44], and WMT’14 English-to-French (En-Fr) [44]. The specific Transformer architecture used by DeLighT, and also in our experiments, varies depending on the dataset. For the WMT’14 En-De we used 12 transformers blocks with 4 attention heads, 512 as the embedding dimension and 1024 for the fan out fully-connected layer. For the other three datasets we used 12 transformers blocks with 8 attention heads, 512 as the embedding dimension and 2048 for the fan out fully-connected layer. For the training setup we have used Adam optimizer, learning rate of 5e-4, inverse square root scheduler with 15K warmup steps, dropout 0.3, weight decay of 1e-4, and 0.1 label-smoothed cross entropy loss.
  Compared to the smallest of the DeLighT models available for each dataset, as summarized in Table 3 and in the right panel of Figure ??, the Shapeshifter models have between 1.5 and 4 times
  fewer parameters, yet offer similar quality; Shapeshifter achieves higher score than DeLighT on two datasets, and lower score on two other datasets.
sections.heading: 5.3 Comparison with Standard Low-rank Factorization
sections.text: To validate experimentally the theoretical results from Section 4.1 concerning the increased expressiveness of the factorization type used in Shapeshifter compared to standard n× r, r ×m low-rank factorization, we trained models factorized this way using the same encoder/decoder rank of 16, and the same embedding rank of 256 as in Shapeshifter. To test only the stack of Transformer layers, we also trained models that do not factorize the embedding matrix. We used one small dataset, IWSLT’14 German-to-English, and one large, WMT’18 English-to-Romanian. The results in Table 4 show that the size of the model is reduced substantially in the Kronecker-based approach compared to standard low-rank representation. The reduction comes predominantly from the embedding, which involve non-square matrices. Despite more than two-times larger number of parameters, standard low-rank representation shows much higher loss on the training set, with no significant difference in terms of generalization outside of the training set, confirming that the difference between the two representations results from increased expressiveness of a stack of low-rank Kronecker-product-based layers. Similar behavior can be observed when the comparison involves models where the embedding matrices are not factorized.
sections.heading: 6 Conclusion
sections.text: Our theoretical and experimental results show that deep models composed of stacked low-rank Kronecker-product-based representations are more expressive, yet smaller, than equivalent models that use standard low-rank factorized matrix representations. These observations help solve a puzzle of how recent methods such as word2ket [17] and PHM layers [18] achieve their impressive parameter reduction rates for embeddings and for encoder/decoder layers, respectively, despite training models de novo, without the help of an existing large model serving as a blueprint for compression or pruning, or as a teacher for knowledge distillation. Building on these results, we provide a comprehensive Shapeshifter4 approach that works both for the embeddings and for the encoder/decoder, using more effective way to exploit Kronecker products for factorized representations. Experimental comparisons with state-of-the-art model reduction approaches, not only PHM but also recently proposed DeLighT [14], on a range of machine translation problems, show that the approach offers similar translation quality with much fewer parameters.
  Funding Disclosure and Competing Interests
  This work was partially funded by NSF grant IIS-1453658 to T.A. The authors have no competing interests.
  4Code available at: https://github.com/tarodz/shapeshifter.
title: Shapeshifter: a Parameter-efficient Transformer using Factorized Reshaped Matrices