Year: 2021

Abstract:
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.

text: 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.

heading: 1 INTRODUCTION
text: 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.

heading: Summary of Contributions:
text: (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.

heading: 2 TENSOR DIAGRAM REPRESENTATION
text: 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.

heading: 2.2 TENSOR DIAGRAM REPRESENTATION OF MULTI-HEAD SELF-ATTENTION
text: The core of Transformer models is a multi-head self-attention (MHSA) module which allows the model to jointly attend to information from different representation sub-spaces at different positions (Vaswani et al., 2017). To distinguish inputs for query, key, value matrices, we use XQ ∈ RN×F , XK ,XV ∈ RM×F respectively. The MHSA module outputs a matrix M ∈ RN×F as
Q[h] = X QWQ[h]; K[h] = X KWK[h]; V[h] = X V W V[h], (1a) head[h] = softmax ( Q[h]K > [h] /√ D ) V[h], (1b)
M = [ head[1], head[2], · · · , head[H] ] WO. (1c)
In the above equations, we see how MHSA computes its output in three steps: (1a) Linear transformation on input embedding to get latent features. For each h, the query/key/value matrices Q[h], K[h], V[h] are linear transformations, parameterized by W Q [h], W K [h], and W V [h] ∈ RF×D respectively, of the input embedding. (1b) Scaled dot-product attention. Each head[h] ∈ RN×D computes a scaled dot-product between a latent feature matrix Q[h] ∈ RN×D (i.e., a query matrix) and another latent feature matrix K[h] ∈ RM×D (i.e., a key matrix) along D, and applies a softmax function to obtain the weights to be multiplied with a third latent feature matrix V[h] ∈ RM×D (i.e., a value matrix). (1c) Concatenation and contraction for MHSA. After concatenation of the H heads, a linear transformation of the concatenated heads, parameterized by WO ∈ RHD×F , is implemented to get the outcome M of the MHSA as a RN×F object.
Figure 3a shows how Single-Head Self-Attention (SHSA) is represented in tensor diagram according to Step (1b). Then we will see how to extend SHSA to MHSA as in Step (1c) (Figure 3b).
Multi-head self-attention in tensor diagram. The concatenation of H matrices {head[h]}Hh=1 can be viewed as constructing a three dimensional array with the third mode being the number of heads, i.e., three-legged nodes in tensor diagram notation. This way, the weight matrices {WQ[h]} H h=1 are stacked to be an order 3 tensor WQ. It is the same with other weight matrices. Since in MHSA, the attention calculation among key, query and value matrices are done separately within each head, batch multiplication along mode H is a proper way to denote these “head-wise” operations. Finally, we connect the resultant tensor with WO to obtain MHSA as in Figure 3b. See Appendix B for more details and the proof.

heading: 3 TUFORMERS: DATA-DRIVEN DESIGN WITH HIGHER EXPRESSIVE POWER

heading: 3.1 MOTIVATION: INTERPRETATION OF MHSA
text: In MHSA, Q[h] and K[h] can only contribute to head[h],∀h. In order to characterize the information flow in MHSA, we stack the weight matrices WQ[h],∀h together (similarly for W K [h] and W V [h]) and introduce a matrix C to interpret MHSA. Guided by its tensor diagram representation in Figure 3b, C is clearly characterized as in Proposition 1.
Proposition 1 (Interpretation of MHSA). With weight matrices stacked as WQ := [WQ[1], · · · ,W Q [H]],W K := [WK[1], · · · ,W K [H]],W V := [W V[1], · · · ,W V [H]] ∈ R
F×DH , and a core matrix C := IH ⊗ (1D1>D) ∈ RDH×DH ( i.e., a Kronecker product of an all-ones matrix 1D1 > D ∈ RD×D and an identity matrix IH ∈ RH×H ), MHSA is equivalent to
Q = XQW Q ; K = XKW K ; V = XV W V , (2a)
headr = softmax ( DH∑ s=1 Crs QsK > s /√ D ) V r, (2b)
M = [ head1,head2, · · · , headDH ] WO. (2c)
In tensor diagram notation, MHSA is equivalent to Figure 4a. Here, we use Qr,Kr,V r ∈ RF , headr ∈ RN to denote the rth column of the corresponding matrices, and Crs to denote the (r, s)th element of C.
Proposition 1 is proved in Appendix D.1. Note that {head1, · · · ,headD} correspond to head[1], {headD+1, · · · ,head2D} correspond to head[2] and so forth. The intuition of interpreting MHSA using Proposition 1 comes from the tensor diagram graphical representation of MHSA structure. From the tensor diagram of MHSA as shown in Figure 3b, we see two types of contractions: latent contractions along D and head contractions along H . The former corresponds to (1D1>D) part of C and the latter corresponds to IH part of C as shown in Figure 4a.
Motivation for an improved design: from MHSA (a preset C) to THSA (a learnable C). In MHSA, with the preset specific C = IH ⊗ (1D1>D), it restricts the possible weight sharing. In other words, Q[h] and K[h] can only contribute to head[h],∀h. However, there is no motivation why such a preset and fixed C is desirable or why we do not allow Q[h] or K[h] to contribute to other heads. Our motivation for an improved design is simple: rather than presetting the core matrix as C = IH ⊗ (1D1>D), we allow a data-driven learnable C, and we call such a design Tunable-head self-attention (THSA). Now with the learnable C in THSA, we allow Q[h] and K[h] to contribute to other heads head[h′], where h′ 6= h. In other words, THSA allows weight sharing across heads.

heading: 3.2 A NOVEL DESIGN: TUNABLE-HEAD SELF-ATTENTION (THSA)
text: As shown in Figure 4b, a Tunable-head self-attention (THSA) module has the same input and output domains as an MHSA, i.e., it takes the input embedding XQ ∈ RN×F and XK , XV ∈ RM×F as inputs and returns a matrix T ∈ RN×F as output. A THSA, with five weight matrices WQ,WK ,W V ∈ RF×R, WO ∈ RR×F , and C ∈ RR×R, is mathematically described as
Q = XQWQ; K = XKWK ; V = XV W V , (3a)
Ωr = softmax
( R∑
s=1
Crs QsK > s
/√ R ) Vr, (3b)
T = ΩWO = [Ω1,Ω2, · · · , ΩR]WO. (3c)
The intermediate result Ωr ∈ RN is the rth column of the matrix Ω. In MHSA, the softmax function is scaled, which normalizes its input by √ D, i.e., the square root of the latent dimension. The scaling ensures that the distribution of the attention matrix approximates a standard Gaussian distribution. To achieve a similar effect, in THSA, we change the constant to √ R accordingly.
THSA generalizes MHSA. THSA has learnable C. (1) If C = 1D1>D, i.e., an D × D all-ones matrix, the THSA module reduces to an MHSA with a single head and latent dimension D (i.e.,
a single-head self-attention with latent dimension D). (2) If C = IH , i.e., an H × H identity matrix, the THSA reduces to a MHSA with H heads and latent dimension 1 (i.e., a heads-only self-attention with H heads). (3) If C = IH ⊗ (1D1>D), i.e., a Kronecker product of an all-ones matrix 1D1>D ∈ RD×D and an identity matrix IH ∈ RH×H , the THSA reduces to an MHSA with H heads and latent dimension D, as illustrated in Figure 4a. (4) We call a Transformer with THSA module as Tunable Transformers (Tuformers). The tensor diagram representation of THSA is given in Figure 4b. It has an learnable core matrix C compared to MHSA in Figure 4a which has a preset core matrix.
Tuformers can be initialized with Pre-trained Transformers. Although different in the structure, Tuformers can always be initialized with pre-trained Transformer models through re-parameterization. MHSA is updated by weight matrices {WK[h]} H h=1, {W Q [h]} H h=1, {W V[h]} H h=1 and W
O while in THSA we have WQ, WK , W V , WO and C. To initialize THSA with MHSA, we follow the scheme proposed in Proposition 1. C is the kronecker product between a D × D all-one matrix and an H ×H identity matrix, where D is the latent dimension and H is the number of heads in MHSA. WO remains the same as that of MHSA. As for WQ, WK and W V ∈ RF×R, we concatenate {WK[h]} H h=1, {W Q [h]} H h=1, {W V[h]} H h=1 ∈ RF×D along mode D.
The notion of heads in THSA. Due to the learnable C and data-dependent weight sharing mechanism, the notion of heads in THSA is not as obvious as in MHSA. Although we have an upper bound R of the number of heads, which is achieved when the learned C decides not to share any weights across heads, the actual effective number of heads depends on C. As a result, we define the number of heads in THSA as the stable rank of the core matrix C, which is ∑ i σ 2 i / maxi σi, where σi is the singular value of C. We use the stable rank because it is largely unaffected by tiny singular values.

heading: 3.3 EXPRESSIVE POWER OF TUFORMERS
text: We first formally describe how to compare the expressive power between two models.
Definition 2 (Expressive Power). Suppose we have two function classes F, G with the same source domain X and target domain Y , i.e., each function f ∈ F (or g ∈ G) is a mapping from X to Y . We say G is more expressive than F if G ⊇ F : for any f ∈ F , there exists g ∈ G such that g(x) = f(x),∀x ∈ X . Furthermore, we say G is strictly more expressive than F if G ⊃ F: besides G ⊇ F , there exists g ∈ G such that given any f there exists x ∈ X and g(x) 6= f(x).
According to the definition, expressive power is a partial order set. Therefore, expressive power of different models is usually not comparable as two sets might have overlap and non-overlap regions. However, when one is a superset or subset of the other, we can unambiguously compare their expressive power. In the following theorem, we formally show that a Tunable-head self-attention (THSA) module has a higher expressive than multi-head self-attention (MHSA) if the rank R in THSA is equal to the product of the number of heads H and the latent dimension D in MHSA.
Theorem 3 (THSA is more expressive than MHSA). Given R = DH and other hyper-parameters being the same, a Tunable-head self-attention (THSA) module with rank R is more expressive than a multi-head self-attention (MHSA) module with H heads and latent dimension D.
Theorem 3, proved in Appendix D.2, shows that any function realizable by MHSA can be realized by THSA.
A Design Space. With the tensor diagram of MHSA in Section 2.2, it is clear that if we take a holistic view of the weight tensors WQ, WK , WV ,WO ∈ RH×D×F and consider them as components of an order 4 tensor of size F × F × F × F (the “super node” represented as the shaded ellipse in Figure 4a, although this order 4 tensor is not a tensor that is free to have any forms because of the limitation brought by the softmax box), then we can use the tensor representation theory to “re-parameterize” the 4th order tensor. This reveals a design space for Transformers. Our THSA serves as a specific design of the “super node”, with a proved higher expressive power than MHSA. Such a design space makes it possible for optimal structure search using tensor representation theories and can inspire better designs and deeper understandings of MHSA and Transformers.

heading: 4 EXPERIMENTS
text: In this section, we show that THSA (used in Tuformers) universally outperforms MHSA (used in Transformers) under a wide range of model scales in diverse tasks, including natural language processing, automatic speech recognition and image generation.
We also demonstrate that THSA is simple to use in practice. First, THSA is a plug-and-play module that can be used in existing Transformer architectures. Second, THSA can be initialized from pre-trained MHSA modules (to avoid expensive training from scratch). For instance, fine-tuning Tuformers initialized from pre-trained BERT and its variants further improves their performance. Finally, THSA can be combined with other kernel-based efficient Transformers to obtain linear computational and memory complexities regarding the sequence length.
Datasets, tasks and evaluation metrics. We evaluate Tuformers on 7 datasets for 5 tasks: word-level Language Modeling (LM) on Penn Treebank (PTB) (Marcus et al., 1993), Neural Machine Translation (NMT) on WMT16 ‘English-German’ (Sennrich et al., 2016), Automatic Speech Recognition (ASR) on LibriSpeech (Panayotov et al., 2015), Natural Language Inference (NLI) on MNLI (Williams et al., 2018) and QNLI (Wang et al., 2018) and Image Generation on CIFAR10 (Krizhevsky, 2009) and MNIST (Deng, 2012) datasets.
Baselines. We compare the performance of MHSA module with THSA module in the following 5 backbone architectures: (1) a vanilla Transformer model (Vaswani et al., 2017), (2) an TransformerXL model (Dai et al., 2019), (3) a BERT-large model (Devlin et al., 2019), (4) a RoBERTa-large model (Liu et al., 2019), and (5) an ALBERT model (Lan et al., 2019). In all baselines, we adopt all default settings used in the original paper (hyperparameters specified in Appendix E), and only replace/re-parameterize the MHSA to THSA.

heading: Performance Results of Tuformers.
text: (1) THSA consistently outperforms MHSA on diverse datasets for a variety of tasks under different model scales. In LM, NMT and ASR tasks, we compare THSA against the MHSA models under different model scales. MHSA is parameterized by four weight matrices WQ, WK , W V and WO whereas THSA adds an additional core matrix C of the same size. Adopting the size of the four original weight matrices in MHSA to THSA incurs an 25% increase in the number of parameters. For a fair comparison, we first compare MHSA against THSA with the same number of parameters (thus with smaller query/key/value/output weight matrices). To get a full picture, we also compare MHSA against THSA with the same sized four original weight matrices. As shown in Figure 5, under the same number of parameters, our THSA (blue curve) universally outperforms MHSA (red curve) under all model scales. By increasing the number of parameters of the attention units by 25% 1 (i.e., adding the core matrix C while keeping the size of the query/key/value/output weight matrices the same), our THSA (purple curve) further improves the performance universally under all model scales.
1As will be discussed later, the overall model size is not drastically increased.
(2) Tuformers initialized with pre-trained strong models further improve the performance. In NLI task, we initialize Tuformers with pre-trained BERT-large, RoBERTa-large and ALBERT models, and perform the same fine-tuning tasks. As shown in Table 1 (Left), THSA is compatible with existing pre-trained Transformer models and there is no need to train Tuformers from scratch: fine-tuning THSA initialized from pre-trained BERT and its variants further improves their performance.
(3) Combing THSA with kernels to obtain linear computation and memory complexities in sequence length. Although an effective design with guaranteed higher expressive power, Tuformers come with an extra memory overhead when calculating values within the softmax. In Tuformers, if done naively, the memory overhead is MNR where N,M denote the sequence length and R = DH is the rank of the core tensor, incurring a D times larger memory overhead than vanilla Transformers in calculating softmax. Combining Tuformers with some state-of-the-art kernel-based efficient Transformers such as Linear Transformers using ELU (Katharopoulos et al., 2020) and polynomial (Tsai et al., 2019) kernels, we reduce the computation and memory complexities of Tuformers to be linear in the sequence length. Kernel-based Transformers apply kernel methods to MHSA to remove the nonlinearity constrain in the softmax function, obtaining linear computation and memory complexities. For example, in Linear Transformer (N =M ), the space complexity is N ×D ×H . Applying the same idea to THSA leads to the same space complexity. We experiment on image generation tasks on CIFAR10 (Krizhevsky, 2009) and MNIST (Deng, 2012) datasets as in Table 1(Right) and show that Tuformers obtain strong performance and linear space complexity when combined with kernels.
Parameter efficiency of Tuformers. THSA introduces an additional trainable core matrix compared with MHSA. However, THSA does not increase the complexity in big O notation compared with MHSA. Moreover, in Transformer models, attention modules are not the only source of parameters and a 25% increase in MHSA often does not increase the total number of parameters significantly. For instance, if we replace MHSA with THSA by keeping the size of the query/key/value/output weight matrices to be the same, a Transformer encoder with 1.3M parameters will be transferred to a Tuformer with 1.6M parameters. Finally, we empirically verify that the core matrix design still maintains parameter efficiency. In Figure 6, we benchmark the training time of Tuformers in multiple tasks and show that THSA with 25% more number of parameters (blue curve) is slightly slower to train compared with MHSA (red curve) but significantly improves the performance (as shown in Figure 5). In addition, as shown previously in Figure 5(middle), using 10% parameters of a vanilla Transformer, a Tuformer obtains a 25.8 BLEU score whereas the original vanilla Transformer has a BLEU score of 27.1 (95% of the performance), demonstrating competitive parameter efficiency.
Ablation study: check-pointing to alleviate memory overhead without kernels. Although combining THSA with kernel-based efficient Transformers relieves the memory overhead as shown above, we investigate other solutions to alleviate the D times larger memory overhead compared with vanilla Transformers in calculating softmax when not using kernels. Specifically, we use checkpointing to eliminate the memory overhead in THSA. Since no intermediate results are saved in the forward pass using check-pointing, we recalculate the intermediate results of the softmax box in the back-propagation, which introduces some computational overhead. Figure 6 (purple curve) shows the per iteration training time overhead induced by check-pointing. From Figure 6 we see that applying check-pointing to Tuformers leads to only a slight increase in training time while relieving the memory overhead. We can reduce or even avoid the repetitive calculations by developing a proper
10% 20% 50% 100% 150% 200%
1.5
2
2.5
3
3.5
model scale
se co
nd pe
ri te
ra tio
n
Language Modeling on Penn Treebank
MHSA THSA THSA with checkpointing
10% 20% 50% 100% 150% 200%
0.2
0.4
0.6
0.8
1.0
1.2
1.4
model scale
se co
nd pe
ri te
ra tio
n
Language Modeling on Penn Treebank
MHSA THSA THSA with checkpointing
Figure 6: (Left) Training and (Right) inference run-time comparison between MHSA, THSA with the same number of parameters and THSA with the same number of parameters with checkpointing. The batch size is 256. The comparison of the NMT task, when THSA has the same sized query/key/value/outputmatrices as MHSA,the can be found in E.
backward module, although such development is beyond the scope of this paper and deferred for future works.

heading: 5 RELATED WORKS
text: Analysis of MHSA. Given the popularity of Transformers in a wide range of domains, a line of works focuses on understanding and improving MHSA. Voita et al. (2019) proposed a practical scheme to prune less informative heads. Michel et al. (2019) also showed that most attention heads are removable in testing without much performance compromise. Cordonnier et al. (2020) found empirical evidence of redundancy in key/query projections and proposed a re-parameterization scheme. These works bring valuable insights into MHSA, but a consistent theoretical analysis for these findings is missing. Our work differs from these works in that we focus on the theoretical understanding of MHSA in terms of expressive power, and we provide theoretical guarantees for a novel design, Tuformers, with higher expressive power and data-driven trainable heads.
Tensor theories and the expressive power of neural networks. There are a long line of works applying tensor methods in modern machine learning. Most of them focus on neural network compression using the tensor decomposition methods as (Ma et al., 2019; Su et al., 2018; Rabanser et al., 2017). There are also works analyzing the expressive power in neural networks using tensor representation theories (Cohen et al., 2016; Khrulkov et al., 2018). Cordonnier et al. (2019) used expressive power in comparing the relationship between the multi-head self-attention and convolutional layers. Yun et al. (2019) established that Transformer models are universal approximators of continuous permutation equivariant sequence-to-sequence functions with compact support. In our work, we represent MHSA in tensor diagram and form a design space where we compare the expressive power of different tensor structures among the weight matrices in MHSA.
Efficient Transformers. MHSA has a global receptive field, i.e., each token’s representation is updated by attending to all other tokens, therefore incurring a quadratic memory and computation complexities concerning the sequence length. An extensive line of works focusing on reliving such the dependencies on the sequence length, such as Performer (Choromanski et al., 2020), Reformer (Kitaev et al., 2020), Linformer (Wang et al., 2020) and Linear Transformer (Katharopoulos et al., 2020). Our proposed model, Tuformers, can be incorporated into some kernel-based efficient Transformers, which applies kernel-based methods to MHSA to remove the nonlinearity constraint brought by the softmax function, achieving the same state-of-the-art linear computation and memory complexities in the sequence length.

heading: 6 CONCLUSION
text: This paper introduces a mathematically rigorous yet intuitive tensor diagram representation of MHSA, formulates a design space where we can analyze the expressive power of MHSA and its variants and proposes Tuformers, a novel model design with a guaranteed higher expressive power. Furthermore, Tuformers have a data-driven structure where heads can be trained implicitly and initialized with pre-trained Transformer models. The introduction of the tensor diagram representations and the theory presented in this paper provide new tools and open new directions for future research searching for expressive and efficient architectures.

heading: ETHICS STATEMENT
text: Transformers are powerful neural network structures that achieve remarkable results in many domains. Understanding the expressive power of such neural networks is desirable but difficult. In this paper, we introduce tensor diagram representation of MHSA (the core component of Transformers) and formulate a design space where we can analyze the expressive power from a tensor perspective. Guided by the tensor diagram representation, we propose a plug-and-play THSA module to replace MHSA with THSA, obtaining Tuformers. Tuformers are theoretically guaranteed to have high expressive power than Transformers.
Our work benefits the Machine Learning community in following points:
(1) We introduce tensor diagram representation to the ML community as an analytical tool for understanding expressive power of neural networks. Equipping Transformer models with such a precise and yet intuitive tool, we are able to close the gap between the currently used simple but intuitive flowcharts and the precise but complex code/mathematical formulation.
(2) By tensorizing MHSA, we provide new possibilities for network analysis from a tensor representation perspective. We formulate a design space for MHSA, which can foster future studies for architecture design.
(3) We explore the theory behind MHSA and propose Tuformers with THSA. The theoretical framework we provide can lead to future explorations and developments of expressive and yet efficient models.
(4) Our proposed Tuformers are energy and resource efficient. (a) Tuformers can be initialized with pre-trained strong models such as BERT, RoBERTa and ALBERT models, drastically reducing the training cost for new downstream tasks. (b) THSA is a plug-and-play module that can be paired with ELU or polynomial kernels to obtain linear computation and memory efficiency in sequence length.

heading: REPRODUCIBILITY STATEMENT
text: For theoretical results, full details are given in Appendix D.
For experimental results, firstly, all the datasets used in this paper are public and contain no sensitive information. Then a detailed description of the datasets, tasks, the hyper-parameters and how the are chosen, and train/test data splitting are in Appendix E. The main code can be found in the supplementary materials.

heading: Appendix of Tuformer: Data-driven Design of Expressive Transformers with Learnable Weight Sharing Across Heads

heading: A SUPPLEMENTARY MATERIAL FOR TENSOR DIAGRAMS
text: In this section, we provide a comprehensive introduction to the tensor diagram representation.

heading: A.1 COMPONENTS OF TENSOR DIAGRAM: ARRAYS AS NODES WITH LEG(S)
text: 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. The order (the number of dimensions) of the array is denoted by the number of legs extending from the node. The legs do not need to be straight lines ,and the orientation does not matter. For example, matrices A and A> are equivalent in this notation. 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). If multiple modes of the same tensor have the same length, we use subscripts to differentiate them.

heading: A.2 OPERATIONS IN TENSOR DIAGRAM
text: There are three types of operations:contraction, softmax and batch multiplication. Contraction is an operation that does summation on element-wise products along a mode in object 1 and a corresponding mode (with the same dimension) in object 2 (i.e., in tensor diagram language, along a leg in node 1 and a corresponding leg in node 2), whereas softmax is an element-wise exponential function normalized along a certain mode. Tensor diagram has a convention to denote the contraction, but not the softmax. We will first introduce contraction in tensor diagram notation and then propose our design of softmax in tensor diagram language.
Contractions denoted as edges connecting the node legs. In tensor diagram, multiplying two matrices (or higher-order tensors with more than 2 legs) corresponds to “gluing” their corresponding legs (along a certain mode) as shown in Appendix A.2, and it is called tensor contraction. Since the representation of arrays in the tensor diagram is orientation invariant, contractions are also orientation invariant. This is especially useful in representing the self-attention unit. In some cases, the data object are represented as row vectors while in some other scenarios they are denoted as column vectors. If we use mathematical formula, we need to take care of the ordering of the matrix multiplication because A>B,AB>,BA>,B>A and more are all different. In contrast, tensor diagram provides a universal graphical representation as long as the corresponding legs are connected correctly.
Nonlinear activation denoted as dotted box with a labelled filled ball attached. We propose to denote the α-scaled softmax function softmax(αA) on A as a dotted box with a labeled filled ball (to distinguish itself from tensor objects, i.e., nodes which are blank circles) attached to one leg. The label right above the filled ball denote the scaling parameter α, and the leg that the ball is attached to indicate the mode where normalization is implemented along in the softmax operation. As shown in Appendix A.2, let A denote the resultant of some operations (for example Q[h]K>[h]), the 1√ D
-scaled softmax on A which normalize along leg N is depicted. Softmax does not change the shape of the input. That is why we introduce a dotted box, which does not change the shape either, to denote it.
Batch Multiplication Given A and B, the batch multiplication is an element-wise product along the connected legs. The resultant tensor C is obtained by merging connected nodes, maintaining the dangling legs M , N and R. Since in MHSA, the attention calculation among key, query and value matrices are done separately within each head, batch multiplication is a proper way to denote these “head-wise” operations.

heading: A.3 ADVANTAGES OF TENSOR DIAGRAM
text: Tensor diagram representations enjoy a few advantages compared with using flow charts together with math formulas.
(1) The notation is orientation invariant. The input embedding matrix X can be represented as either RN×F or RF×N with rows or columns being the embedding of each token in the sequence respectively. Although we take the former convention in our mathematical notations in Equation B.1, there are papers that use the latter, creating discrepancies between notations. However, in tensor diagram, since the node legs are orientation invariant, we obtain a universal graphical representation of multi-head self-attention, irrespective of how the input embedding matrix is represented, as shown in Figure 8c.
(2) The representation is both precise in math and intuitive, making analysis easier. The tensor diagram itself is mathematically rigorous.
With the labels of the legs, the model complexity (number of parameters) is explicitly displayed, requiring no supplementary information such as the size of the parameters WQ, WK and W V as in mathematical formulas. In addition,the resultant from an operation or a sequence of operations (i.e., a connected sub-graph of the entire tensor diagram graph) can be treated as “merging” any connected nodes in the sub-graph. This is particularly convenient to obtain the size of the resultant: any the dangling leg becomes a leg of the resultant.
(3) Multi-linear operations are concisely represented and interpreted. Using tensor diagram, we can concisely represent the multi-head self-attention rigorously in one diagram as shown in Figure 8c. More important, the multi-linear interaction between the latent features Q[h], K[h] and V[h], which essentially attribute to the multi-linear interaction between the weight tensors WQ, WK and WB , is clearly illustrated in the tensor diagram. This illustration, which is not achieved by any existing representations, is crucial for our principled understanding of the multi-head self-attention.
(4) The information flow is clear. For example, in Figure 8a there is no direct information pass from WK[h], W Q [h] to W V [h] while in Figure 8c, such exchange is done through edge Ds and Hs. It is such observation (there are multiple paths that connect the weight parameter matrices) that leads to our idea of comparing expressive power between different structures.

heading: B SUPPLEMENTARY MATERIAL FOR MULTI-HEAD SELF-ATTENTION
text: In the section, we prove the equivalence between the tensor diagram and the equations for multi-head self-attention (MHSA). To establish such an equivalence, we will need to show that both tensor diagram and the matrix equations lead to the same result in element-wise notation.
For convenience of reference, we recap both representations in Equation (B.1) and Figure 8.
Q[h] = X QWQ[h]; K[h] = X KWK[h]; V[h] = X V W V[h]. (B.1a) head[h] = softmax ( Q[h]K > [h] /√ D ) V[h], (B.1b)
M = [ head[1],head[2], · · · , head[H] ] WO, (B.1c)
Note that the query weight tensor WQ comes from a concatenation of H query weight matrices {WQ[h]} H h=1 such that W Q h,:,: = W Q [h]. Similarly, W
K , WV are concatenated from {WK[h]}Hh=1, {WV [h]}Hh=1 such that W K h,:,: = W K [h],W V h,:,: = W V [h]. Furthermore, we denote Q, K, V as the reshaped tensors from {Q[h]}Hh=1, {K[h]}Hh=1, {V[h]}Hh=1. Finally, W O is reshaped from WO.
To differentiate between equations and tensor diagrams in the proof, we add an over-line to every object in the tensor diagram throughout this section.

heading: B.1 SINGLE-HEAD SELF-ATTENTION (SHSA)
text: We will prove that the tensor diagram in Figure 8a is equivalent to Equations (B.1a) and (B.1b).
Proof. For simplicity, we omit the subscript [h] in this subsection.
We first show that a contraction of the nodes within the softmax box is equivalent to QK>. Let A = QK> = XQWQ(WKXK)>, where XQ ∈ RN×F , XK ∈ RM×F , WQ, WK ∈ RF×D. We use subscripts in the element-wise notation to differentiate edges when multiple edges with same dimension appear in the same equation.
Qnd = F∑ f1=1 XQnf1W Q f1d , (B.2a)
Kmd = F∑ f2=1 XKmf2W K f2d, (B.2b)
Anm = D∑ d=1 QndKdm = D∑ d=1 F∑ f1=1 F∑ f2=1 XQnf1W Q f1d WKf2dX K mf2 . (B.2c)
Denote be the result of the subset of nodes in the softmax box as A.
Anm = D∑ d=1 F∑ f1=1 F∑ f2=1 X Q nf1W Q f1dW K f2dX K mf2 . (B.3)
Comparing Equations (B.2c) and (B.3), we have A = A. Let B = softmax(A) and B = softmax(A), we further have B = softmax(A) = softmax(A) = B.
Let M be the result of the tensor diagram and M be the matrix representation of the single-head attention. We have
Mnd = N∑ m=1 F∑ f=1 BnmX V mfW V fd. (B.4)
Mnd = N∑ m=1 F∑ f=1 BnmX V mfW V fd. (B.5)
Comparing Equations (B.4) and (B.5), we have C is equivalent to C, which completes the proof.

heading: B.2 MULTI-HEAD SELF-ATTENTION (MHSA)
text: We now prove that tensor diagram in Figure 8c is equivalent to Equation (1c).
Proof. Let H be the number of heads. Let A ∈ RH×N×M be the result of a contraction of the nodes within the softmax box. From the last section, we know that Ah,:,: = A[h] = Q[h]K>[h]. Applying the softmax function B = softmax(A), we have Bh,:,: = softmax(Q[h]K>[h]).
Let T be result of the tensor diagram and T be that of Equation (B.1c). Tnf4 = ∑
m,h,d1,d2
softmax(Q[h],nd1Q[h],md1)V[h],md2W O [h],d2f4 . (B.6)
T n1f4 = ∑
m,h,d1,d2
softmax(Qhnd1Khmd1)Vhmd2W O hmd2 . (B.7)
Comparing T and T , we complete the proof.

heading: C CLOSER INVESTIGATION OF MHSA AND ITS COMPARISON TO THSA
text: Rethinking MHSA. There are two types contractions in MHSA as shown in Figure 4a: HeadContraction, a global contraction among all weights WQ, WK , WV and WO along mode H , and Latent-Contraction, a local contraction between WQ and WK as well as between WV and WO both along mode D. If we remove the Head-Contraction and maintain the Latent-Contraction only, multi-head is reduced to single-head with an additional contraction with the weight matrix WO. To understand the role of the two types of contractions, we propose ablation studies by removing one of them and study the changes in the structure and the expressive power. It turns out that in both situation, MHSA falls into a special case of THSA.
Repeated Proposition 1 with two special cases. Given R = DH and other hyperparameters being the same, THSA module includes MHSA as a special case. Specifically, (1) If C = 1D1>D, i.e., an D ×D all-ones matrix, the THSA module reduces to an MHSA with a single head and latent dimension D (i.e., a single-head self-attention with latent dimension D). (2) If C = IH , i.e., an H ×H identity matrix, the THSA reduces to a MHSA with H heads and latent dimension 1 (i.e., a heads-only self-attention with H heads). (3) If C = IH ⊗ (1D1>D), i.e., a Kronecker product of an all-ones matrix 1D1>D ∈ RD×D and an identity matrix IH ∈ RH×H , the THSA reduces to an MHSA with H heads and latent dimension D. The proof of the above propositions is given in Appendix D.
Definition 4 (Stable Rank). Let A ∈ RN×M , the stable rank of A of is defined as ∑
i σ 2 i / maxi σi
where σi is the singular value of A.
The notion of heads in THSA. The concept of the number of heads in THSA generalizes to the stable rank of the core matrix, allowing a data-driven implicit training. The number of heads H in MHSA corresponds to the stable rank of the core matrix C, defined as ∑ i σ 2 i / maxi σi where σi is the singular value. It is a useful surrogate for the rank because it is largely unaffected by tiny singular values.
As shown in Appendix C when C is constructed by the outer product of two vectors, the stable rank of C is 1, while in this case, the whole structure is equivalent to a single-head self-attention. When C is an identical matrix as shown in appendix C, the stable rank of C equals to the rank of the identical matrix, which is H , corresponding to the number of heads of MHSA. Also, the structure on the right can be obtained by removing the latent contraction edges D from Figure 4a. When C is the Kronecker product of the all-one matrix and the identical matrix, which corresponds to MHSA, the stable rank of C is H , and this number does not change in the training because that C is a fixed matrix in MHSA.

heading: D SUPPLEMENTARY MATERIAL FOR TUFORMER AND PROOFS
text: In this section, we will prove Proposition 1 and Theorem 3 in Section 3. In both theorems, we assume the rank in Tunable-head self-attention (THSA) equals to the product of the number of heads H and the latent dimension D in multi-head self-attention (MHSA), i.e., R = DH .
For convenience of reference, we recap the mathematical expressions and tensor diagrams for both multi-head self-attention (MHSA) and Tunable-head self-attention (THSA) in the following.
Multi-head Self-Attention (MHSA). An MHSA module has its learnable parameters as H sets of (query, key, value) weight matrices {WQ[h],W K [h],W V [h]} H h=1 (with W Q [h],W K [h],W V [h] ∈ R F×D for each h) and an output weight matrix WO ∈ RHD×F . The module maps three input matrices XQ ∈ RN×F , XK ,XV ∈ RM×F into an output matrix M ∈ RN×F .
Q[h] = X QWQ[h]; K[h] = X KWK[h]; V[h] = X V W V[h], (D.1a) head[h] = softmax ( Q[h]K > [h] ) V[h], (D.1b)
M = [ head[1], head[2], · · · , head[H] ] WO. (D.1c)
In Equation (D.1a), Q[h] ∈ RN×D,K[h],V[h] ∈ RM×D are query, key, value matrices respectively. In Equation (D.1b), we omit the scaling factor 1/ √ D for the softmax function (the scalar can be merged into the learnable parameters), and head[h] ∈ RN×D is the resulted matrix for head h. Tunable-Head Self-Attention. A THSA module is parameterized by five matrices: a query weight matrix W
Q ∈ RF×R, a key weight matrix WK ∈ RF×R, a value weight matrix W V ∈ RF×R, an output weight matrix W
O ∈ RR×F , and an additional core matrix C ∈ RR×R. The module has the same input and output domains as in MHSA, i.e., it takes three matrices XQ ∈ RN×F , XK ,XV ∈ RM×F as inputs and returns a matrix T ∈ RN×F as output.
Q = XQW Q ; K = XKW K ; V = XV W V , (D.2a)
headr = softmax
( R∑
s=1
Crs QsK > s ) V r, (D.2b)
T = head W O = [ head1, head2, · · · , headR ] . (D.2c)
To distinguish THSA from MHSA, we use over-scored symbols whenever needed. In Equation (D.2a), Q ∈ RN×R,K,V ∈ RM×R are query, key, value matrices respectively. In Equation (D.2b), we again omit the scaling factor 1/ √ R in the softmax function. Equation (D.2b) leads to a head matrix head ∈ RF×N , where headr ∈ RF denotes the rth column of the matrix head.

heading: D.1 PROOF OF PROPOSITION 1
text: Proof of Proposition 1. We constructively prove that THSA reduces to three cases of MHSA if the query/key/value weight matrices in THSA concatenate the corresponding matrices in MHSA:
W Q = [ WQ[1], W Q [2], · · · , W Q [H] ] , (D.3a)
W K = [ WK[1], W K [2], · · · , W K [H] ] , (D.3b)
W V = [ W V[1], W V [2], · · · , W V [H] ] . (D.3c)
In addition, we set W O = WO. We have W Q ,W K ,W V ∈ RF×R, and WO ∈ RR×F . (1) For C = 1R1>R, i.e. an R×R all-ones matrix, we aim to prove that the THSA module reduces to a single-head self-attention with H = 1 and D = R. Since there is only one head in the module, we have W
Q = WQ[1], W K = WK[1] , and W V = W V[1]. As an immediate result,
Q = XQW Q = XQWQ[1] = Q[1], (D.4a)
K = XKW K = XKWK[1] = K[1], (D.4b)
V = XV W V = XV W V[1] = V[1]. (D.4c)
Since C is an all-ones matrix, i.e., Crs = 1,∀r, s, we rewrite Equation (D.2b) as:
headr = softmax
( R∑
s=1
QsK > s ) V r = softmax ( QK >) V r. (D.5)
Notice that the equation holds for each column r, we further write all R equations jointly as: head = softmax ( QK >) V = softmax ( Q[1]K > [1] ) V[1] = head[1]. (D.6)
Finally, we express the output matrix as:
T = head W O = head[1]W O = M , (D.7)
which concludes the reduction for the special case (1).
(2) For C = IR, an R × R identity matrix, we aim to prove that the THSA module reduces to a heads-only self-attention with H = R and D = 1. Since the latent dimension D = 1, i.e., each WQ[h] (or WK[h],W V [h]) is a vector, we have W Q r = W Q [r], W K r = W K [r] , and W V r = W V [r]. Therefore,
Qr = X QW
Q r = X QWQ[r] = Q[r], (D.8a)
Kr = X KW
K r = X KWK[r] = K[r], (D.8b)
V r = X V W
V r = X V W V[r] = V[r]. (D.8c)
Since C is an identity matrix, i.e., Crr = 1 and Crs = 0,∀r 6= s, we rewrite Equation (D.2b) as: headr = softmax ( QrK > r ) V r = softmax ( Q[r]K > [r] ) V[r] = head[r]. (D.9)
Finally, we express the output matrix as: T = head W O = [ head[1], head[2], · · · ; head[R] ] WO = M , (D.10)
which concludes the reduction for the special case (2).
(3) For C = IH ⊗ (1D1>D), i.e., a Kronecker product between an identity matrix IH ∈ RH×H and an all-ones matrix 1D1>D ∈ RD×D, we aim to prove that the THSA module reduces to a multi-head self-attention (MHSA) module with H heads and latent dimension D. In this general case, we have
Q = XQW Q = XQ [ WQ[1], W Q [2], · · · , W Q [H] ] = [ Q[1], Q[2], · · · , Q[H] ] (D.11)
K = XKW K = XK [ WK[1], W K [2], · · · , W K [H] ] = [ K[1], K[2], · · · , K[H] ] (D.12)
V = XV W V = XV [ W V[1], W V [2], · · · , W V [H] ] = [ V[1], V[2], · · · , V[H] ] (D.13)
Notice that C is a block diagonal matrix blkdiag(1D1>D, · · · ,1D1>D), where each block is an all-ones matrix. Therefore we can rewrite Equation (D.2b) for r = (h− 1)D + d as:
head(h−1)D+d = softmax
( R∑
s=1
1(h−1)D<s≤hD ( QsK > s )) V r (D.14)
= softmax
( D∑
d=1
Q[h],dK > [h],d ) V[h],d (D.15)
= softmax ( Q[h]K > [h] ) V[h],d. (D.16)
Since Equation (D.16) holds for each column d, we write all D equations jointly as:[ head(h−1)D+d, · · · , headhD ] = softmax ( Q[h]K > [h] ) V[h] = head[h]. (D.17)
Again, since Equation (D.17) holds for each block h, we combine all H equations as: head = [ head[1], head[2], · · · , head[H] ] . (D.18)
Finally, we express the output matrix as: T = head W O = [ head[1], head[2], · · · , head[H] ] WO = M , (D.19)
which concludes the reduction for the case (3).

heading: D.2 PROOF OF THEOREM 3
text: Proof of Theorem 3. We have shown that THSA reduces to MHSA when C takes specific forms. This shows that THSA is more expressive than MHSA — any mapping by an MHSA module can also be realized by a THSA module.
To prove that THSA is strictly more expressive than MHSA, we need find a mapping by a THSA module that can not be realized by an MHSA module. It suffices to show that the matrix M is rank-deficient, while the matrix T can be full-rank. We assume MH < min(N,F ).
For convenience, we divide the matrix WO for MHSA into H blocks.
WO = [ WO[1] > , WO[2] > , · · · , WO[H] ]> , (D.20)
where WO[h] ∈ R D×F is the matrix for the hth head. We now rewrite the output matrix M as:
M = H∑ h=1 softmax ( Q[h]K > [h] ) ︸ ︷︷ ︸
RN×M
( V[h]W O [h] ) ︸ ︷︷ ︸
RM×F
. (D.21)
Since the matrix inside summation is a product of two matrices of size RN×M and RM×F , its rank is at most M . Use the property rank(A +B) ≤ rank(A) + rank(B), we have rank(M) ≤ MH < min(N,F ). Similarly, we divide the matrix W O for THSA into R rows and rewrite the output matrix T as:
T = R∑ r=1 softmax
( R∑
s=1
Crs Q:,rK > :,r ) ︸ ︷︷ ︸
RN×M
( V :,rW O
r,: ) ︸ ︷︷ ︸
RM×F
. (D.22)
Using the same argument, we have rank(T ) ≤ MR. Since H < R (in fact, R = DH), we can always find an example such that rank(T ) > rank(M). This completes the proof.

heading: E SUPPLEMENTARY MATERIAL FOR EXPERIMENTS
text: All the experiments are run on computing nodes with 4 NVIDIA GeForce GPUs.
(1) Language Modeling on Penn Treebank. Language modeling is the task of computing the probability of a sentence or a sequence of words. The model performance is measured by per-word Perplexity (PPL) which is the lower the better.
Penn Treebank (Marcus et al., 1993) is under the LDC User Agreement for Non-Members. We adopt the default settings as in the vanilla Transformer models(Vaswani et al., 2017). Specifically, we have N (length of the sequence) to be 2048, F (length of the initial embedding) to be 512, 8 number of heads, 6 layers encoder-decoder structure, D (latent feature dimension) to be 64. We use Adam optimizer and have 4000 warm-up steps. We also use label smoothing(0.1). The model is trained for 50 epochs.
(2) Neural Machine Translation on WMT16. The goal of NMT is to generate a corresponding sequence in one language given a sequence in another. The performance is measured by BLEU scores. (the higher, the better).
WMT16 English-German dataset (Sennrich et al., 2016) is under the MIT license. In this experiment, we use a Transformer model with 8 layers and we have F = 512 and R = 1024. The learning rate is 1e−5 and 4000 warmups. We also set label-smooting to be 0.1. To speed up the experiment, we use the mix-precision trick with an optimization level being O1.
(3) Image Generation on MNIST and CIFAR-10. The image generation task is to predict an image pixel by pixel. The performance is evaluated by bits per dimension (BPD). The image generation task is chosen to demonstrate, firstly, that Tuformers work beyond the language domain, and secondly, Tuformers can improve other efficient designs.
CIFAR10 (Krizhevsky, 2009) dataset is under the MIT license (MIT), MNIST(LeCun et al., 2010) is under the Creative Commons Attribution-Share Alike 3.0 license. The reason why we choose to evaluate Tuformer in image generation tasks on these two datasets is that we want to show the compatibility of our model to other efficient designs. Katharopoulos et al. (2020) introduces a linear transformer that has a state-of-the-art complexity. As a result, we evaluate our design on the same experiments as they do in their paper. We adopt basically the same settings except for GPU resources. Concretely, we use a 8 layer transformer model with 8 heads. The embedding size F is 256 anD is 32. N is set to be 1024. We also use “RAdam” (Ma & Yarats, 2021) optimizer for a state-of-the-art result. The model is trained for 250 epochs for MNIST dataset and 50 epochs for CIFAR10 dataset. The batch size is 10 in MNIST training and 4 in CIFAR training. We additionally use he mix-precision trick with an optimization level set to be O1 to speed up the training.
(4) Automatic Speech Recognition on LibriSpeech Dataset ASR consists of transcribing audio speech segments into text. The performance is measured by Word Error Rate (WER). (the lower, the better). We adopt the default settings as in Fairseq (Ott et al., 2019), except that we replace MHSA with THSA.
(5) Natural Language Inference on MNLI and QNLI. NLI is a task of determining whether the given “hypothesis” and “premise” logically follow (entailment) or unfollow (contradiction) or are undetermined (neutral) to each other. The results are evaluated by accuracy. Using MNLI dataset, the model tries to predict whether sentence A entails or contradicts B while in QNLI dataset, the model is trained to answer whether sentence B contains answers to the question in sentence A. We use a batch size of 32 for 10 epochs, with a learning rate 1e−5.

heading: Additional Results

Review summaries:
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.

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.