Tuformer: Data-Driven Design of Expressive Transformer by Tucker Tensor Representation

Transformers are neural network architectures that achieve remarkable performance in many areas . However , the core component of Transformers , multi-head selfattention ( MHSA ) , is mainly derived from heuristics , and the interactions across its components are not well understood . To address the problem , we first introduce a mathematically rigorous and yet intuitive tensor diagram representation of MHSA . Guided by tensor diagram representations , we formulate a design space where we can analyze the expressive power of the network structure , providing new directions and possibilities for enhanced performance . We then propose a novel design , namely Tunable Transformer ( Tuformers ) , with a guaranteed higher expressive power than MHSA by allowing learnable weight sharing across heads . Unlike vanilla Transformer models , where the number of heads is a pre-defined fixed constant , Tuformer ’ s structure is data-driven , and the number of heads is trainable . Training of Tuformers could be made very efficient as it allows initialization from existing pre-trained Transformer models . We test Tuformers on various tasks across multiple domains and show competitive results under a wide range of model sizes . 1 INTRODUCTION . Transformer models are first introduced by Vaswani et al . ( 2017 ) in the context of sequence modelling . They have demonstrated impressive results on a wide variety of tasks in many fields , such as language model pre-training ( Sanh et al. , 2019 ) , speech recognition ( Moritz et al. , 2020 ) , image classification ( Dosovitskiy et al. , 2020 ) , and generation ( Parmar et al. , 2018 ) . The core component of Transformer models is the multi-head self-attention ( MHSA ) which is extended from the standard attention mechanism ( Bahdanau et al. , 2014 ) . Each attention head in MHSA has a global receptive field , i.e. , each token ’ s representation is updated by attending to all other tokens , and H attention heads are computed in parallel and concatenated together . The current MHSA design is mainly derived from empirical studies or heuristics , leaving some unresolved challenges . ( 1 ) Lack of solid theoretical understanding . The theory behind MHSA is only starting to catch up with practice . The role of the components and their interactions in MHSA are not well understood rigorously from a theoretical perspective , which may lead to inefficiencies in the design . For example , Michel et al . ( 2019 ) show that most attention heads can be removed in the testing phase without much performance compromise , while Cordonnier et al . ( 2020 ) find empirical evidence of redundancy in key/query projections and propose a re-parameterization scheme . These works focus more on practical solutions , leading to questions about whether theoretical patterns exist in these designs . ( 2 ) The number of heads is not trainable . Intuitively , the heads in multi-heads are expected to capture different context information through each head . However , the number of heads is fixed in training . Thus , although we could tune the hyper-parameter , an exhaustive search would be time-consuming or not practical for large-scale tasks . ( 3 ) Hard to analyze the expressive power . Analyzing the expressive power ( as will be defined in Definition 2 ) of a neural network , i.e. , proving that some architectures are more expressive than others , is a non-trivial and challenging task . Such analysis has been done in Convolutional Neural Networks ( CNNs ) ( LeCun et al. , 1995 ; Cohen et al. , 2016 ) , Recurrent Neural Networks ( RNNs ) ( Mikolov et al. , 2010 ; Khrulkov et al. , 2018 ) , but no such work exists for interpreting MHSA or guiding the structural design . In response to the above challenges , we first interpret MHSA from a tensor representation perspective using the intuitive graphical tool , the tensor diagram . Current prevalent descriptions of MHSA use flow charts to convey high-level intuitions , which could cause ambiguities . Therefore , it is inevitable to pair those flow charts with mathematical formulas to understand the mechanism precisely . However , these two separated descriptions create difficulties for interpretations and inspections of the operations implemented in MHSA . To address this issue , we propose a graphical representation of MHSA , which is both semantically intuitive and mathematically rigorous . Specifically , we modify and extend the vanilla tensor diagrams ( Penrose , 1971 ) , which conveniently allow for rigorous graphical representation of multi-linear operations between multi-dimensional arrays ( i.e. , higher-order tensors ) to represent nonlinear operations . We also propose a novel data-driven structure that is guaranteed to have higher expressive power . Inspired from the tensor diagram representation of MHSA , we obtain a novel design , namely TunableHead Self-Attention ( THSA ) , which is a re-parameterization of the weight matrices to allow learnable weight sharing across heads . Transformers with THSA , named Tunable Transformers ( Tuformers ) , have several advantages compared against vanilla Transformers : ( 1 ) A guaranteed higher expressive power . We prove that MHSA is a special case of THSA . ( 2 ) The number of heads is trainable . The concept of the number of heads in THSA generalizes to the stable rank of the core matrix , allowing data-driven implicit training . ( 3 ) Tuformers allow initialization from pre-trained Transformers such as BERT ( Devlin et al. , 2019 ) and its variants . We experiment Tuformers with several tasks across multiple domains , from language modeling , machine translation to image generation under a wide range of model sizes . We demonstrate competitive results not only on Tuformers but also in cases where Tuformers are initialized with pre-trained Transformers for other downstream tasks and when combined with Linear Transformer ( Katharopoulos et al. , 2020 ) on the image generation task . Using tensor diagram representation , we project the current design of MHSA into its tensor form , which renders a holistic view of the weights in MHSA for better interpretation of the information flow and exchange among the components of MHSA . More importantly , we formulate a design space for MHSA that allows improving the expressive power of the network by finding better tensor representations of the weights through tensor representation theory in followup works . Summary of Contributions : . ( 1 ) We propose a mathematically rigorous and semantically intuitive tensor diagram representation of the multi-head self-attention , introducing a new tool to the ML community for future studies on interpretation and improvements of Transformers . ( 2 ) We propose a novel design of the MHSA , Tunable-Head Self-Attention ( THSA ) , resulting in Tuformers , which has a theoretical guaranteed improvement in expressive power . ( 3 ) We provide a constructive re-parameterization scheme from vanilla Transformers to Tuformers so that Tuformers can be easily used in fine-tuning tasks with pre-trained models and efficient Transformers such as Linear Transformer ( Katharopoulos et al. , 2020 ) and Performer ( Choromanski et al. , 2020 ) with the state-of-the-art computation and memory complexities . ( 4 ) We formulate a design space for attention unit design and a framework for analyzing their expressive power using tensor representation theory under the intuitive guidance of tensor diagrams . 2 TENSOR DIAGRAM REPRESENTATION . Notations . We use lower case letters ( e.g. , v ) to denote vectors , upper case letters ( e.g. , M ) to denote matrices , and curly letters ( e.g. , T ) to denote general tensors . For a tensor T ∈ RI1×···IM , we refer to the number of indices as order , each individual index as mode , and the length of one mode as dimension . For instance , T is an M th order tensor that has dimension Im at its mth mode . We reserve superscripts to distinguish similar arrays ( e.g. , WQ , WK , W V are query/key/value weight matrices ) , and subscripts to index the elements in an array ( e.g. , Wij is the ( i , j ) th element of W ) . We use colon : to slice an array ( e.g. , Wi , : denotes the ith row of W ) . We propose to use tensor diagrams ( Penrose , 1971 ) , a commonly used rigorous/precise and intuitive graphical representation for multi-linear operations among higher-order arrays ( i.e. , tensors ) , to represent multi-head self-attention ( MHSA ) . Since MHSA consists of multilinear operations and a nonlinear softmax function , we will introduce tensor diagrams and our novel design extending tensor diagrams to denote the composition of nonlinear and multilinear operations in MHSA . 2.1 TENSOR DIAGRAM BASICS v M 𝓣 Vector order 1 tensor Matrix order 2 tensor 3 dimensional array order 3 tensor A BA A B C Figure 1 : Arrays in tensor diagram . A vector is denoted as a node with 1 leg , a matrix as a node with 2 legs and an N -dimensional array as a node with N legs . Arrays denoted as nodes with legs . An array is represented as a node with leg ( s ) in a tensor diagram as shown in Figure 1 . We denote the order ( the number of dimensions ) of the array by the number of legs extending from the node . Each labeled leg represents one mode of a tensor . Every mode of the tensor needs to be uniquely labeled . We usually use the dimension of the mode as the label ( i.e. , an associated positive integer written on top of each leg ) . The legs do not need to be straight lines , and their orientations do not matter . Matrices M ∈ RA×B and M > ∈ RB×A can be represented via the same tensor diagram as long as the M node has two legs ( to denote that it is a matrix , wherever the legs extend to ) , labeled as A and B ( to denote its size ) . Operations in tensor diagrams . There are three types of operations in the calculation of MHSA : , contraction , softmax , and batch multiplication , as shown in Figure 2 and explained in its caption . Evaluation of tensor diagrams : ( 1 ) Evaluation order . We can evaluate a tensor diagram in any pair-wise order except for the nodes in the softmax box . Since the softmax function is nonlinear , we must first evaluate the nodes in the softmax box with arbitrary order before those outside the box . ( 2 ) Reading a tensor diagram . We can easily identify the output shape by the dangling edges . For instant , in Section 2.1 , there are three dangling legs M , N , R. Thus the output is a 3rd order tensor with dimensions M , N , R. Note that the softmax function does not change the shape of the tensor . Advantages of tensor diagrams : ( 1 ) Tensor diagram is orientation invariant , meaning that we can represent X and X > using the same tensor diagram . Thus in multi-head , we obtain a universal graphical representation regardless of whether we represent an embedding of each token in the sequence as rows or columns of the input embedding matrix X . ( 2 ) Multi-linear operations are concisely represented and interpreted . ( 3 ) The representation is both precise in math and intuitive , and the information flow is clear , making analysis of network structure more accessible . In addition , with the labels of the legs , we can read the model complexity explicitly . We include a comprehensive introduction to tensor diagram representation in Appendix A .

Summary: The paper investigates the multi-head self-attention mechanism (MHSA) of transformer networks through the lens of tensor decompositions via tensor diagram notation. The authors propose an extension to MHSA inspired by the Tucker decomposition (termed THSA), analyze its expressive power, and demonstrate that it belongs to a class of more expressive functions than MHSA. Further, the authors show the positive effect of this drop-in replacement on the downstream tasks.