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: Focusing on the multi-head self-attention (MHSA) structure, this paper proposes an extension of the tensor diagram to denote self-attention (SA) structures more intuitively. Then inspired by the Tucker format, this paper also proposes a new form of SA named Tucker-Head Self-Attention (THSA), which can also be illustrated with the tensor diagram extension. Although THSA derives some success, THSA may lack theoretical guarantee and experiments seems insufficient.