Original: T UFORMER : D ATA - DRIVEN D ESIGN OF T RANSFORMERS FOR I MPROVED G ENERALIZATION OR E FFICIENCY
Revised: T UFORMER : D ATA - DRIVEN D ESIGN OF T RANSFORMERS FOR I MPROVED G ENERALIZATION OR E FFICIENCY

Original: A BSTRACT
Revised: A BSTRACT

Original: Transformers are neural network architectures that achieve remarkable performance in many areas.
Revised: Transformers are neural network architectures that achieve remarkable performance in many areas.

Original: 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.
Revised: 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.

Original: To address the problem, we ﬁrst introduce a mathematically rigorous and yet intuitive tensor diagram representation of MHSA.
Revised: To address the problem, we ﬁrst introduce a mathematically rigorous and yet intuitive tensor diagram representation of MHSA.

Original: Guided by tensor diagram representations, we propose a novel design, namely
Revised: Guided by tensor diagram representations, we propose a novel design, namely

Original: Tunable Transformers (Tuformers), by allowing data-driven weights across heads, whereas MHSA adopts pre-deﬁned and ﬁxed weights across heads, as will be explained in our paper.
Revised: Tunable Transformers (Tuformers), by allowing data-driven weights across heads, whereas MHSA adopts pre-deﬁned and ﬁxed weights across heads, as will be explained in our paper.

Original: Tuformers naturally reveal a ﬂexible design space that a user, depending on the needs, can choose a structure that has either improved performance (generalization error) or higher model efﬁciency.
Revised: Tuformers naturally reveal a ﬂexible design space that a user, depending on the needs, can choose a structure that has either improved performance (generalization error) or higher model efﬁciency.

Original: Any pre-trained
Revised: Any pre-trained

Original: Transformer can be an initialization of the corresponding Tuformer with trainable number of heads for efﬁcient training and ﬁne-tuning.
Revised: Transformer can be an initialization of the corresponding Tuformer with trainable number of heads for efﬁcient training and ﬁne-tuning.

Original: Tuformers universally outperform Transformers on various tasks across multiple domains under a wide range of model sizes.
Revised: Tuformers universally outperform Transformers on various tasks across multiple domains under a wide range of model sizes.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Transformer models are ﬁrst introduced by Vaswani et al.
Revised: Transformer models are ﬁrst introduced by Vaswani et al.

Original: (2017) in the context of sequence modelling.
Revised: (2017) in the context of sequence modelling.

Original: They have demonstrated impressive results on a wide variety of tasks in many ﬁelds, such as language model pre-training (Sanh et al., 2019), speech recognition (Moritz et al., 2020), image classiﬁcation (Dosovitskiy et al., 2020), and generation (Parmar et al., 2018).
Revised: They have demonstrated impressive results on a wide variety of tasks in many ﬁelds, such as language model pre-training (Sanh et al., 2019), speech recognition (Moritz et al., 2020), image classiﬁcation (Dosovitskiy et al., 2020), and generation (Parmar et al., 2018).

Original: The core component of
Revised: The core component of

Original: Transformer models is the multi-head self-attention (MHSA) which is extended from the standard attention mechanism (Bahdanau et al., 2014).
Revised: Transformer models is the multi-head self-attention (MHSA) which is extended from the standard attention mechanism (Bahdanau et al., 2014).

Original: Each attention head in MHSA has a global receptive ﬁeld, 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.
Revised: Each attention head in MHSA has a global receptive ﬁeld, 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.

Original: The current MHSA design is mainly derived from empirical studies or heuristics, leaving some unresolved challenges.
Revised: The current MHSA design is mainly derived from empirical studies or heuristics, leaving some unresolved challenges.

Original: (1) Lack of solid theoretical understanding.
Revised: (1) Lack of solid theoretical understanding.

Original: The theory behind MHSA is only starting to catch up with practice.
Revised: The theory behind MHSA is only starting to catch up with practice.

Original: The role of the components and their interactions in MHSA are not well understood rigorously from a theoretical perspective, which may lead to inefﬁciencies in the design.
Revised: The role of the components and their interactions in MHSA are not well understood rigorously from a theoretical perspective, which may lead to inefﬁciencies in the design.

Original: For example, Michel et al.
Revised: For example, Michel et al.

Original: (2019) show that most attention heads can be removed in the testing phase without much performance compromise, while Cordonnier et al.
Revised: (2019) show that most attention heads can be removed in the testing phase without much performance compromise, while Cordonnier et al.

Original: ﬁnd empirical evidence of redundancy in key/query projections and propose a re-parameterization scheme.
Revised: ﬁnd empirical evidence of redundancy in key/query projections and propose a re-parameterization scheme.

Original: These works focus more on practical solutions, leading to questions about whether theoretical patterns exist in these designs.
Revised: These works focus more on practical solutions, leading to questions about whether theoretical patterns exist in these designs.

Original: The number of heads is not trainable.
Revised: The number of heads is not trainable.

Original: Intuitively, the heads in multi-heads are expected to capture different context information through each head.
Revised: Intuitively, the heads in multi-heads are expected to capture different context information through each head.

Original: However, the number of heads is ﬁxed in training.
Revised: However, the number of heads is ﬁxed in training.

Original: Thus, although we could tune the hyper-parameter, an exhaustive search would be time-consuming or not practical for large-scale tasks.
Revised: Thus, although we could tune the hyper-parameter, an exhaustive search would be time-consuming or not practical for large-scale tasks.

Original: (3) Hard to analyze the expressive power.
Revised: (3) Hard to analyze the expressive power.

Original: Analyzing the expressive power (as will be deﬁned in Deﬁnition 3) of a neural network, i.e., proving that some architectures are more expressive than others, is a non-trivial and challenging task.
Revised: Analyzing the expressive power (as will be deﬁned in Deﬁnition 3) of a neural network, i.e., proving that some architectures are more expressive than others, is a non-trivial and challenging task.

Original: 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.
Revised: 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.

Original: In response to the above challenges, we ﬁrst interpret MHSA from a tensor representation perspective using the intuitive graphical tool, the tensor diagram .
Revised: In response to the above challenges, we ﬁrst interpret MHSA from a tensor representation perspective using the intuitive graphical tool, the tensor diagram .

Original: Current prevalent descriptions of MHSA use ﬂow charts to convey high-level intuitions, which could cause ambiguities.
Revised: Current prevalent descriptions of MHSA use ﬂow charts to convey high-level intuitions, which could cause ambiguities.

Original: Therefore, it is inevitable to pair those ﬂow charts with mathematical formulas to understand the mechanism precisely.
Revised: Therefore, it is inevitable to pair those ﬂow charts with mathematical formulas to understand the mechanism precisely.

Original: However, these two separated descriptions create difﬁculties for interpretations and inspections of the operations implemented in MHSA.
Revised: However, these two separated descriptions create difﬁculties for interpretations and inspections of the operations implemented in MHSA.

Original: To address this issue, we propose a graphical representation of MHSA, which is both semantically intuitive and mathematically rigorous.
Revised: To address this issue, we propose a graphical representation of MHSA, which is both semantically intuitive and mathematically rigorous.

Original: Speciﬁcally, 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.
Revised: Speciﬁcally, 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.

Original: 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 ﬂow and exchange among the components of MHSA.
Revised: 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 ﬂow and exchange among the components of MHSA.

Original: We then propose a novel data-driven structure , namely Tunable-Head Self-Attention (THSA) , which is a re-parameterization of the weight matrices to allow learnable weight across heads.
Revised: We then propose a novel data-driven structure , namely Tunable-Head Self-Attention (THSA) , which is a re-parameterization of the weight matrices to allow learnable weight across heads.

Original: Transformers with THSA, named Tunable Transformers (Tuformers) , have several advantages compared against vanilla Transformers: (1) A guaranteed higher expressive power.
Revised: Transformers with THSA, named Tunable Transformers (Tuformers) , have several advantages compared against vanilla Transformers: (1) A guaranteed higher expressive power.

Original: We prove that MHSA is a special case of THSA.
Revised: We prove that MHSA is a special case of THSA.

Original: The number of heads is trainable.
Revised: The number of heads is trainable.

Original: The concept of the number of heads in THSA generalizes to the stable rank of the core matrix, allowing data-driven implicit training.
Revised: The concept of the number of heads in THSA generalizes to the stable rank of the core matrix, allowing data-driven implicit training.

Original: (3) Tuformers allow initialization from pre-trained
Revised: (3) Tuformers allow initialization from pre-trained

Original: Transformers such as BERT (Devlin et al., 2019) and its variants.
Revised: Transformers such as BERT (Devlin et al., 2019) and its variants.

Original: More importantly, we formulate a ﬂexible design space for MHSA where we explore the trade-off between generalization and efﬁciency, which can be adjusted by the expressive power if there is no over-ﬁtting.
Revised: More importantly, we formulate a ﬂexible design space for MHSA where we explore the trade-off between generalization and efﬁciency, which can be adjusted by the expressive power if there is no over-ﬁtting.

Original: Users can choose structures with either higher model efﬁciency or better generalization that satisfy their own needs.
Revised: Users can choose structures with either higher model efﬁciency or better generalization that satisfy their own needs.

Original: The proposed design space further allows improving the expressive power of the network by ﬁnding better tensor representations of the weights through tensor representation theory in future followup works.
Revised: The proposed design space further allows improving the expressive power of the network by ﬁnding better tensor representations of the weights through tensor representation theory in future followup works.

Original: We experiment Tuformers with several tasks across multiple domains, from language modeling, machine translation to image generation under a wide range of model sizes.
Revised: We experiment Tuformers with several tasks across multiple domains, from language modeling, machine translation to image generation under a wide range of model sizes.

Original: We demonstrate competitive results not only on Tuformers but also in cases where Tuformers are initialized with pre-trained
Revised: We demonstrate competitive results not only on Tuformers but also in cases where Tuformers are initialized with pre-trained

Original: Transformers for other downstream tasks, when combined with Linear Transformer (Katharopoulos et al., 2020) on the image generation task and their efﬁcient variants in the design space.
Revised: Transformers for other downstream tasks, when combined with Linear Transformer (Katharopoulos et al., 2020) on the image generation task and their efﬁcient variants in the design space.

Original: Summary of Contributions:
Revised: Summary of Contributions:

Original: (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.
Revised: (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.

Original: (2) We propose a novel design of the MHSA, Tunable-Head Self-Attention (THSA) , resulting in Tuformers .
Revised: (2) We propose a novel design of the MHSA, Tunable-Head Self-Attention (THSA) , resulting in Tuformers .

Original: Tuformers are structurally more expressive and are showing improved generalization error experimentally.
Revised: Tuformers are structurally more expressive and are showing improved generalization error experimentally.

Original: (3) We formulate a ﬂexible design space for attention unit design where users can choose to design structures, depending on the needs, that have better generalization or efﬁciency.
Revised: (3) We formulate a ﬂexible design space for attention unit design where users can choose to design structures, depending on the needs, that have better generalization or efﬁciency.

Original:  T ENSOR D IAGRAM R EPRESENTATION
Revised:  T ENSOR D IAGRAM R EPRESENTATION

Original: Notations.
Revised: Notations.

Original: 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.
Revised: 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.

Original: For a tensor T ∈ R I 1 ×··· I M , we refer to the number of indices as order , each individual index as mode , and the length of one mode as dimension .
Revised: For a tensor T ∈ R I 1 ×··· I M , we refer to the number of indices as order , each individual index as mode , and the length of one mode as dimension .

Original: For instance, T is an M th order tensor that has dimension I m at its m th mode.
Revised: For instance, T is an M th order tensor that has dimension I m at its m th mode.

Original: We reserve superscripts to distinguish similar arrays (e.g., W Q , W K , W V are query/key/value weight matrices), and subscripts to index the elements in an array (e.g., W ij is the ( i, j ) th element of W ).
Revised: We reserve superscripts to distinguish similar arrays (e.g., W Q , W K , W V are query/key/value weight matrices), and subscripts to index the elements in an array (e.g., W ij is the ( i, j ) th element of W ).

Original: We use colon : to slice an array (e.g., W i, : denotes the i th row of W ).
Revised: We use colon : to slice an array (e.g., W i, : denotes the i th row of W ).

Original: 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).
Revised: 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).

Original: 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.
Revised: 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.

Original: 2. T 
Revised: 2. T ENSOR D IAGRAM B ASICS

Original: [Figure]
Revised: [Figure]

Original: Arrays denoted as nodes with legs.
Revised: Arrays denoted as nodes with legs.

Original: An array is represented as a node with leg(s) in a tensor diagram as shown in Figure 1.
Revised: An array is represented as a node with leg(s) in a tensor diagram as shown in Figure 1.

Original: We denote the order (the number of dimensions) of the array by the number of legs extending from the node.
Revised: We denote the order (the number of dimensions) of the array by the number of legs extending from the node.

Original: Each labeled leg represents one mode of a tensor.
Revised: Each labeled leg represents one mode of a tensor.

Original: Every mode of the tensor needs to be uniquely labeled.
Revised: Every mode of the tensor needs to be uniquely labeled.

Original: We usually use the dimension of the mode as the label (i.e., an associated positive integer written on top of each leg).
Revised: We usually use the dimension of the mode as the label (i.e., an associated positive integer written on top of each leg).

Original: The legs do not need to be straight lines, and their orientations do not matter.
Revised: The legs do not need to be straight lines, and their orientations do not matter.

Original: Matrices M ∈ R A × B and
Revised: Matrices M ∈ R A × B and

Original: M (cid:62) ∈
Revised: M (cid:62) ∈

Original: R B × 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). 
Revised: R B × 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). [Figure] [Equation] M

Revised: N

Original: [Figure]
Revised: [Figure]

Original: [Equation] 
Revised: [Equation] [Figure] [Equation] [Figure]

Revised: C M

Original: [Figure]
Revised: N = R A M B

Revised: N

Revised: R

Revised: A mr B nr = C nmr

Original: Operations in tensor diagrams.
Revised: Operations in tensor diagrams.

Original: 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.
Revised: 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.

Original: Evaluation of tensor diagrams: (1) Evaluation order.
Revised: Evaluation of tensor diagrams: (1) Evaluation order.

Original: We can evaluate a tensor diagram in any pair-wise order except for the nodes in the softmax box.
Revised: We can evaluate a tensor diagram in any pair-wise order except for the nodes in the softmax box.

Original: Since the softmax function is nonlinear, we must ﬁrst evaluate the nodes in the softmax box with arbitrary order before those outside the box.
Revised: Since the softmax function is nonlinear, we must ﬁrst evaluate the nodes in the softmax box with arbitrary order before those outside the box.

Original: (2) Reading a tensor diagram.
Revised: (2) Reading a tensor diagram.

Original: We can easily identify the output shape by the dangling edges.
Revised: We can easily identify the output shape by the dangling edges.

Original: For instant, in Section 2.1, there are three dangling legs M , N , R .
Revised: For instant, in Section 2.1, there are three dangling legs M , N , R .

Original: Thus the output is a 3 rd order tensor with dimensions M , N , R .
Revised: Thus the output is a 3 rd order tensor with dimensions M , N , R .

Original: Note that the softmax function does not change the shape of the tensor.
Revised: Note that the softmax function does not change the shape of the tensor.

Original: Advantages of tensor diagrams: (1) Tensor diagram is orientation invariant , meaning that we can represent X and X (cid:62) using the same tensor diagram.
Revised: Advantages of tensor diagrams: (1) Tensor diagram is orientation invariant , meaning that we can represent X and X (cid:62) using the same tensor diagram.

Original: 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 .
Revised: 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 .

Original: (2) Multi-linear operations are concisely represented and interpreted.
Revised: (2) Multi-linear operations are concisely represented and interpreted.

Original: The representation is both precise in math and intuitive, and the information ﬂow is clear, making analysis of network structure more accessible.
Revised: The representation is both precise in math and intuitive, and the information ﬂow is clear, making analysis of network structure more accessible.

Original: In addition, with the labels of the legs, we can read the model complexity explicitly.
Revised: In addition, with the labels of the legs, we can read the model complexity explicitly.

Original: We include a comprehensive introduction to tensor diagram representation in Appendix A.
Revised: We include a comprehensive introduction to tensor diagram representation in Appendix A.

Original: 2. T ENSOR D IAGRAM R EPRESENTATION OF MULTI - HEAD SELF - ATTENTION
Revised: 2. T ENSOR D IAGRAM R EPRESENTATION OF MULTI - HEAD SELF - ATTENTION

Original: 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).
Revised: 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).

Original: To distinguish inputs for query, key, value matrices, we use X Q ∈ R N × F ,
Revised: To distinguish inputs for query, key, value matrices, we use X Q ∈ R N × F ,

Original: X K , X V ∈ R M × F respectively.
Revised: X K , X V ∈ R M × F respectively.

Original: The MHSA module outputs a matrix M ∈ R N × F as
Revised: The MHSA module outputs a matrix M ∈ R N × F as

Original: [Equation]
Revised: [Equation]

Original: In the above equations, we see how MHSA computes its output in three steps :
Revised: In the above equations, we see how MHSA computes its output in three steps :

Original: (1a) Linear transformation on input embedding to get latent features.
Revised: (1a) Linear transformation on input embedding to get latent features.

Original: For each h
Revised: For each h

Original: the query/
Revised: the query/

Original: key
Revised: key

Original: /value matrices Q [ h ] , K
Revised: /value matrices Q [ h ] , K

Original: [ h ] ,
Revised: [ h ] ,

Original: V [
Revised: V [

Original: h ] are linear transformations, parameterized by W Q [ h
Revised: h ] are linear transformations, parameterized by W Q [ h

Original: , W K [ h
Revised: , W K [ h

Original: and W V [ h ] ∈ R F × D respectively, of the input embedding.
Revised: and W V [ h ] ∈ R F × D respectively, of the input embedding.

Original: (1b)
Revised: (1b)

Original: Scaled dot-product attention.
Revised: Scaled dot-product attention.

Original: Each head [ h ] ∈ R N × D computes a scaled dot-product between a latent feature matrix Q [ h ] ∈ R N × D (i
Revised: Each head [ h ] ∈ R N × D computes a scaled dot-product between a latent feature matrix Q [ h ] ∈ R N × D (i

Original: e., a query matrix) and another latent feature matrix K [ h ] ∈ R M × 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 ] ∈ R M × D (i.e., a value matrix).
Revised: e., a query matrix) and another latent feature matrix K [ h ] ∈ R M × 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 ] ∈ R M × D (i.e., a value matrix).

Original: (1c) Concatenation and contraction for MHSA
Revised: (1c) Concatenation and contraction for MHSA

Original: After concatenation of the H heads, a linear transformation of the concatenated heads, parameterized by W O ∈ R HD × F , is implemented to get the outcome M of the MHSA as a R N × F object.
Revised: After concatenation of the H heads, a linear transformation of the concatenated heads, parameterized by W O ∈ R HD × F , is implemented to get the outcome M of the MHSA as a R N × F object.

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F 1 D D
Revised: F 1 D D

Original: D
Revised: D

Original: W Q [ h ] W K [ h ]
Revised: W Q [ h ] W K [ h ]

Original: W V [ h ]
Revised: W V [ h ]

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K (a) Single-head Self-Attention
Revised: X K (a) Single-head Self-Attention

Original: M
Revised: M

Original: N
Revised: N

Original: F
Revised: F

Original: F 1 D
Revised: F 1 D

Original: D
Revised: D

Original: H
Revised: H

Original: Head Contraction
Revised: Head Contraction

Original: Latent Contraction
Revised: Latent Contraction

Original: F
Revised: F

Original: F D Latent Contraction
Revised: F D Latent Contraction

Original: X Q 𝒲 Q 𝒲 K
Revised: X Q 𝒲 Q 𝒲 K

Original: X K 𝒲 O 𝒲 V
Revised: X K 𝒲 O 𝒲 V

Original: X V (b) Multi-Head Self-Attention
Revised: X V (b) Multi-Head Self-Attention

Original: Multi-head self-attention in tensor diagram.
Revised: Multi-head self-attention in tensor diagram.

Original: 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.
Revised: 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.

Original: This way, the weight matrices { W Q [ h ] } Hh =1 are stacked to be an order 3 tensor W Q .
Revised: This way, the weight matrices { W Q [ h ] } Hh =1 are stacked to be an order 3 tensor W Q .

Original: It is the same with other weight matrices.
Revised: It is the same with other weight matrices.

Original: 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.
Revised: 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.

Original: Finally, we connect the resultant tensor with W O to obtain MHSA as in Figure 3b.
Revised: Finally, we connect the resultant tensor with W O to obtain MHSA as in Figure 3b.

Original: See Appendix B for more details and the proof.
Revised: See Appendix B for more details and the proof.

Original:  T UFORMERS : D ATA -D RIVEN D ESIGN WITH H IGHER E XPRESSIVE P OWER 3. M OTIVATION : I NTERPRETATION OF MHSA
Revised:  T UFORMERS : D ATA -D RIVEN D ESIGN WITH H IGHER E XPRESSIVE P OWER 3. M OTIVATION : I NTERPRETATION OF MHSA

Original: In MHSA, Q [ h ] and K [ h ] can only contribute to head [ h ] , ∀ h .
Revised: In MHSA, Q [ h ] and K [ h ] can only contribute to head [ h ] , ∀ h .

Original: In order to characterize the information ﬂow in MHSA, we stack the weight matrices W Q [ h ] , ∀ h together (similarly for W K [ h ] and W V [ h ] ) and introduce a matrix C to interpret MHSA.
Revised: In order to characterize the information ﬂow in MHSA, we stack the weight matrices W Q [ h ] , ∀ h together (similarly for W K [ h ] and W V [ h ] ) and introduce a matrix C to interpret MHSA.

Original: Guided by its tensor diagram representation in Figure 3b,
Revised: Guided by its tensor diagram representation in Figure 3b,

Original: C is clearly characterized as in Proposition 1.
Revised: C is clearly characterized as in Proposition 1.

Original: Proposition 1 (Interpretation of MHSA) .
Revised: Proposition 1 (Interpretation of MHSA) .

Original: With weight matrices stacked as W Q :=
Revised: With weight matrices stacked as W Q :=

Original: [ W Q [1] , · · · , W Q [ H ] ] , W K :=
Revised: [ W Q [1] , · · · , W Q [ H ] ] , W K :=

Original: [ W K [1] , · · · , W K [ H ] ] , W V :=
Revised: [ W K [1] , · · · , W K [ H ] ] , W V :=

Original: [ W V [1] , · · · , W V [ H ] ] ∈ R F × DH , and a core matrix C :
Revised: [ W V [1] , · · · , W V [ H ] ] ∈ R F × DH , and a core matrix C :

Original: = I H ⊗ ( 1 D 1 (cid:62) D ) ∈ R DH × DH ( i.e., a Kronecker product of an all-ones matrix 1 D 1 (cid:62) D ∈ R D
Revised: = I H ⊗ ( 1 D 1 (cid:62) D ) ∈ R DH × DH ( i.e., a Kronecker product of an all-ones matrix 1 D 1 (cid:62) D ∈ R D

Original: × D and an identity matrix I H ∈ R H × H ), MHSA is equivalent to
Revised: × D and an identity matrix I H ∈ R H × H ), MHSA is equivalent to

Original: [Equation]
Revised: [Equation]

Original: In tensor diagram notation, MHSA is equivalent to Figure 4a.
Revised: In tensor diagram notation, MHSA is equivalent to Figure 4a.

Original: Here, we use Q r , K r , V r ∈ R F , head r ∈ R N to denote the r th column of the corresponding matrices, and C rs to denote the ( r, s ) th element of C .
Revised: Here, we use Q r , K r , V r ∈ R F , head r ∈ R N to denote the r th column of the corresponding matrices, and C rs to denote the ( r, s ) th element of C .

Original: Proposition 1 is proved in Appendix D.1.
Revised: Proposition 1 is proved in Appendix D.1.

Original: Note that { head 1 , · · · , head D } correspond to head [1] , { head D +1 , · · · , head 2 D } correspond to head [2] and so forth.
Revised: Note that { head 1 , · · · , head D } correspond to head [1] , { head D +1 , · · · , head 2 D } correspond to head [2] and so forth.

Original: The intuition of interpreting MHSA using Proposition 1 comes from the tensor diagram graphical representation of MHSA structure.
Revised: The intuition of interpreting MHSA using Proposition 1 comes from the tensor diagram graphical representation of MHSA structure.

Original: 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 .
Revised: 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 .

Original: The former corresponds to ( 1 D 1 (cid:62) D ) part of C and the latter corresponds to I H part of C as shown in Figure 4a.
Revised: The former corresponds to ( 1 D 1 (cid:62) D ) part of C and the latter corresponds to I H part of C as shown in Figure 4a.

Original: Motivation for an improved design: from MHSA (a preset C ) to THSA (a learnable C ).
Revised: Motivation for an improved design: from MHSA (a preset C ) to THSA (a learnable C ).

Original: In
Revised: In

Original: MHSA, with the preset speciﬁc C
Revised: MHSA, with the preset speciﬁc C

Original: =
Revised: =

Original: I H ⊗ ( 1 D 1 (cid:62) D ) , it restricts the possible weight sharing.
Revised: I H ⊗ ( 1 D 1 (cid:62) D ) , it restricts the possible weight sharing.

Original: In other words, Q [ h ] and K [ h ] can only contribute to head [ h ] , ∀ h .
Revised: In other words, Q [ h ] and K [ h ] can only contribute to head [ h ] , ∀ h .

Original: However, there is no motivation why such a preset and ﬁxed C is desirable or why we do not allow Q [ h ] or K
Revised: However, there is no motivation why such a preset and ﬁxed C is desirable or why we do not allow Q [ h ] or K

Original: [ h ] to contribute to other heads.
Revised: [ h ] to contribute to other heads.

Original: Our motivation for an improved design is simple: rather than presetting the core matrix as
Revised: Our motivation for an improved design is simple: rather than presetting the core matrix as

Original: [Equation]
Revised: [Equation]

Original: , we allow a data-driven learnable C , and we call such a design Tunable-head self-attention (THSA) .
Revised: , we allow a data-driven learnable C , and we call such a design Tunable-head self-attention (THSA) .

Original: Now with the learnable C in THSA, we allow Q [ h ] and K [ h ] to contribute to other heads head [ h (cid:48) ] , where h (cid:48) (cid:54)
Revised: Now with the learnable C in THSA, we allow Q [ h ] and K [ h ] to contribute to other heads head [ h (cid:48) ] , where h (cid:48) (cid:54)

Original: = h .
Revised: = h .

Original: In other words, THSA allows learnable weights across heads.
Revised: In other words, THSA allows learnable weights across heads.

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F F 1 R 𝒲 Q 𝒲 K 𝒲 O 𝒲 V
Revised: F F 1 R 𝒲 Q 𝒲 K 𝒲 O 𝒲 V

Original: I H
Revised: I H

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K 1 D 1 D
Revised: X K 1 D 1 D

Original: D
Revised: D

Original: D
Revised: D

Original: H H
Revised: H H

Original: C =
Revised: C =

Original: I h ⊗( 1 ⊤ D 1 D )
Revised: I h ⊗( 1 ⊤ D 1 D )

Original: [Figure]
Revised: [Figure]

Original: M
Revised: M

Original: N
Revised: N

Original: F
Revised: F

Original: F 1 R R
Revised: F 1 R R

Original: F
Revised: F

Original: F
Revised: F

Original: [Figure]
Revised: [Figure]

Original: Q
Revised: Q

Original: [Figure]
Revised: [Figure]

Original: K
Revised: K

Original: W O
Revised: W O

Original: W V
Revised: W V

Original: [Figure]
Revised: [Figure]

Original: C fully
Revised: C fully

Original: [Figure]
Revised: [Figure]

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F F 1 R 𝒲 Q 𝒲 K 𝒲 O 𝒲 V
Revised: F F 1 R 𝒲 Q 𝒲 K 𝒲 O 𝒲 V

Original: C
Revised: C

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K
Revised: X K

Original: C 2 C 

Revised: C 2 C

Original: D D
Revised: D D

Original: H H
Revised: H H

Original: C = C 1 ⊗( C ⊤ 2 C 3 )
Revised: C = C 1 ⊗( C ⊤ 2 C 3 )

Original: [Figure] [Figure]
Revised: [Figure] [Figure]

Original: [Equation]
Revised: [Equation]

Original: EAD S ELF -A TTENTION (THSA)
Revised: EAD S ELF -A TTENTION (THSA)

Original: As shown in Figure 4b, a Tunable-head self-attention (THSA) 1 module has the same input and output domains as an MHSA, i.e., it takes the input embedding X Q ∈ R N × F and X K , X V ∈ R M × F as inputs and returns a matrix T ∈ R N × F as output.
Revised: As shown in Figure 4b, a Tunable-head self-attention (THSA) 2 module has the same input and output domains as an MHSA, i.e., it takes the input embedding X Q ∈ R N × F and X K , X V ∈ R M × F as inputs and returns a matrix T ∈ R N × F as output.

Original: A THSA, with ﬁve weight matrices
Revised: A THSA, with ﬁve weight matrices

Original: W Q , W K , W V ∈ R F × R , W O ∈ R R × F , and C ∈ R R × R , is mathematically described as
Revised: W Q , W K , W V ∈ R F × R , W O ∈ R R × F , and C ∈ R R × R , is mathematically described as

Original: [Equation]
Revised: [Equation]

Original: The intermediate result Ω r ∈ R N is the r th column of the matrix Ω .
Revised: The intermediate result Ω r ∈ R N is the r th column of the matrix Ω .

Original: In MHSA, the softmax function is scaled, which normalizes its input by √ D , i.e., the square root of the latent dimension.
Revised: In MHSA, the softmax function is scaled, which normalizes its input by √ D , i.e., the square root of the latent dimension.

Original: The scaling ensures that the distribution of the attention matrix approximates a standard Gaussian distribution.
Revised: The scaling ensures that the distribution of the attention matrix approximates a standard Gaussian distribution.

Original: To achieve a similar effect, in THSA, we change the constant to √ R accordingly.
Revised: To achieve a similar effect, in THSA, we change the constant to √ R accordingly.

Original: THSA generalizes MHSA.
Revised: THSA generalizes MHSA.

Original: THSA has learnable C .
Revised: THSA has learnable C .

Original: If C = I
Revised: If C = I

Original: H ⊗ ( 1 D 1 (cid:62) D ) , i.e., a Kronecker product of an all-ones matrix 1 D 1 (cid:62) D ∈ R D
Revised: H ⊗ ( 1 D 1 (cid:62) D ) , i.e., a Kronecker product of an all-ones matrix 1 D 1 (cid:62) D ∈ R D

Original: × D and an identity matrix
Revised: × D and an identity matrix

Original: I H ∈
Revised: I H ∈

Original: R H
Revised: R H

Original: × H , the THSA reduces to an
Revised: × H , the THSA reduces to an

Original: MHSA with H heads and latent dimension D , as illustrated in Figure 4a.
Revised: MHSA with H heads and latent dimension D , as illustrated in Figure 4a.

Original: We call a Transformer with
Revised: We call a Transformer with

Original: THSA module as Tunable Transformers (Tuformers).
Revised: THSA module as Tunable Transformers (Tuformers).

Original: The tensor diagram representation of THSA is given in Figure 4b.
Revised: The tensor diagram representation of THSA is given in Figure 4b.

Original: It has an learnable core matrix C compared to MHSA in Figure 4a which has a preset core matrix.
Revised: It has an learnable core matrix C compared to MHSA in Figure 4a which has a preset core matrix.

Original: Tuformers can be initialized with Pre-trained Transformers.
Revised: Tuformers can be initialized with Pre-trained Transformers.

Original: Although different in the structure,
Revised: Although different in the structure,

Original: Tuformers can always be initialized with pre-trained Transformer models through re-parameterization.
Revised: Tuformers can always be initialized with pre-trained Transformer models through re-parameterization.

Original: MHSA is updated by weight matrices { W K [ h ] } Hh =1 , { W Q [ h ] } Hh =1 , { W V [ h ] } Hh =1 and W O while in THSA we have W Q , W K , W V , W O and C .
Revised: MHSA is updated by weight matrices { W K [ h ] } Hh =1 , { W Q [ h ] } Hh =1 , { W V [ h ] } Hh =1 and W O while in THSA we have W Q , W K , W V , W O and C .

Original: To initialize THSA with MHSA, we follow the scheme proposed in Proposition 1.
Revised: To initialize THSA with MHSA, we follow the scheme proposed in Proposition 1.

Original: C is the kronecker product between a D × D all-one matrix and an
Revised: C is the kronecker product between a D × D all-one matrix and an

Original: H ×
Revised: H ×

Original: H identity matrix, where D is the latent dimension and H is the number of heads in MHSA.
Revised: H identity matrix, where D is the latent dimension and H is the number of heads in MHSA.

Original: W O remains the same as that of MHSA.
Revised: W O remains the same as that of MHSA.

Original: As for W Q , W K and W V ∈ R F × R , we concatenate { W K
Revised: As for W Q , W K and W V ∈ R F × R , we concatenate { W K

Original: [ h ] } Hh =1 , { W Q [ h ] } Hh
Revised: [ h ] } Hh =1 , { W Q [ h ] } Hh

Original: =1 , { W V [ h ] } Hh =1 ∈ R F × D along mode D .
Revised: =1 , { W V [ h ] } Hh =1 ∈ R F × D along mode D .

Original: The notion of heads in THSA.
Revised: The notion of heads in THSA.

Original: Due to the learnable C and data-dependent weight sharing mechanism, the notion of heads in THSA is not as obvious as in MHSA.
Revised: Due to the learnable C and data-dependent weight sharing mechanism, the notion of heads in THSA is not as obvious as in MHSA.

Original: Although we have an upper bound
Revised: Although we have an upper bound

Original: 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 .
Revised: 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 .

Original: As a result, we deﬁne the number of heads in THSA as the stable rank of the core matrix C , which is (cid:80) i σ 2 i (cid:14) max i σ i , where σ i is the singular value of C .
Revised: As a result, we deﬁne the number of heads in THSA as the stable rank of the core matrix C , which is (cid:80) i σ 2 i (cid:14) max i σ i , where σ i is the singular value of C .

Original: We use the stable rank because it is largely unaffected by tiny singular values.
Revised: We use the stable rank because it is largely unaffected by tiny singular values.

Original: THSA as a ﬂexible design space: C in THSA ( Figure 4b ) naturally reveals a design space.
Revised: THSA as a ﬂexible design space: C in THSA ( Figure 4b ) naturally reveals a design space.

Original: In
Revised: In

Original: MHSA, C is completely ﬁxed and in THSA C is fully trainable.
Revised: MHSA, C is completely ﬁxed and in THSA C is fully trainable.

Original: We can also have ﬂexible C designs for user-speciﬁc needs of either increasing generalization or improving efﬁciency.
Revised: We can also have ﬂexible C designs for user-speciﬁc needs of either increasing generalization or improving efﬁciency.

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F F 1 R 𝒲 Q 𝒲 K 𝒲 O 𝒲 V
Revised: F F 1 R 𝒲 Q 𝒲 K 𝒲 O 𝒲 V

Original: I H
Revised: I H

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K
Revised: X K

Original: C 2 C 

Revised: C 2 C

Original: D D
Revised: D D

Original: H
Revised: H

Original: H C
Revised: H C

Original: =
Revised: =

Original: I h ⊗( C ⊤ 2 C 3 )
Revised: I h ⊗( C ⊤ 2 C 3 )

Original: [Figure]
Revised: [Figure]

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F F 1 R
Revised: F F 1 R

Original: C
Revised: C

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K 1 D 1 D
Revised: X K 1 D 1 D

Original: D
Revised: D

Original: D
Revised: D

Original: H
Revised: H

Original: H 𝒲 Q 𝒲 K 𝒲 O 𝒲 V
Revised: H 𝒲 Q 𝒲 K 𝒲 O 𝒲 V

Original: C = C 1 ⊗( 1 ⊤ D 1 D )
Revised: C = C 1 ⊗( 1 ⊤ D 1 D )

Original: [Figure]
Revised: [Figure]

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F F 1 R W Q
Revised: F F 1 R W Q

Original: W K X Q
Revised: W K X Q

Original: X V
Revised: X V

Original: X K H
Revised: X K H

Original: W V W O
Revised: W V W O

Original: C 1 H C
Revised: C 1 H C

Original: [Figure]
Revised: [Figure]

Original: C
Revised: C

Original: [Figure]
Revised: [Figure]

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F F 1 R W Q
Revised: F F 1 R W Q

Original: W K X Q
Revised: W K X Q

Original: [Figure]
Revised: [Figure]

Original: V
Revised: V

Original: [Figure]
Revised: [Figure]

Original: H
Revised: H

Original: [Figure]
Revised: [Figure]

Original: We demonstrate some of them extending from the structure of MHSA where C takes the form of a Kronecker product as shown in Figure 4c.
Revised: We demonstrate some of them extending from the structure of MHSA where C takes the form of a Kronecker product as shown in Figure 4c.

Original: The relationship between the generalization and efﬁciency of these designs is summarized in Figure 6.
Revised: The relationship between the generalization and efﬁciency of these designs is summarized in Figure 6.

Original: [Figure]
Revised: [Figure]

Original: Better generalization.
Revised: Better generalization.

Original: Better generalization can be obtained by increasing the expressive power if there is no over-ﬁtting.
Revised: Better generalization can be obtained by increasing the expressive power if there is no over-ﬁtting.

Original: A naive way to increase the expressive power is to add more parameters in the weight matrices in MHSA, such as increasing the latent dimension D or the number of heads H .
Revised: A naive way to increase the expressive power is to add more parameters in the weight matrices in MHSA, such as increasing the latent dimension D or the number of heads H .

Original: We ﬁnd that a more effective way to increase the expressive power for better generalization is to ﬁnd more expressive structures.
Revised: We ﬁnd that a more effective way to increase the expressive power for better generalization is to ﬁnd more expressive structures.

Original: We demonstrate a few of them in Figure 5 and Figure 4b.
Revised: We demonstrate a few of them in Figure 5 and Figure 4b.

Original: In MHSA, C is pre-set ﬁxed where C =
Revised: In MHSA, C is pre-set ﬁxed where C =

Original: I H ⊗ (1 (cid:62) D 1 D ) .
Revised: I H ⊗ (1 (cid:62) D 1 D ) .

Original: In Figure 5a (skip-left C
Revised: In Figure 5a (skip-left C

Original: =
Revised: =

Original: I H ⊗ ( C (cid:62) 2 C 3 ) ) and Figure 5b (skip-right C
Revised: I H ⊗ ( C (cid:62) 2 C 3 ) ) and Figure 5b (skip-right C

Original: =
Revised: =

Original: I H ⊗ ( C (cid:62) 2 C 3 ) ), we maintain the form of a Kronecker product but allow either C 1 or
Revised: I H ⊗ ( C (cid:62) 2 C 3 ) ), we maintain the form of a Kronecker product but allow either C 1 or

Original: C 2 , C 3 trainable.
Revised: C 2 , C 3 trainable.

Original: We can think of skip-left and skipright as skip-connection designs in C .
Revised: We can think of skip-left and skipright as skip-connection designs in C .

Original: In THSA, C is fully trainable.
Revised: In THSA, C is fully trainable.

Original: It is easy to observe an increasing in expressive power from MHSA to skip-left, skip-right, and ﬁnally THSA, together with an increasing in generalization in the same order.
Revised: It is easy to observe an increasing in expressive power from MHSA to skip-left, skip-right, and ﬁnally THSA, together with an increasing in generalization in the same order.

Original: Our empirical study also show that these structurally more expressive models more effectively improve the generalization than naively increase the number of parameters.
Revised: Our empirical study also show that these structurally more expressive models more effectively improve the generalization than naively increase the number of parameters.

Original: More Efﬁciency.
Revised: More Efﬁciency.

Original: We can obtain different levels of efﬁcient structures in the design space by adjusting the expressive power.
Revised: We can obtain different levels of efﬁcient structures in the design space by adjusting the expressive power.

Original: Compared to MHSA, skip-left and skip-right have more generalization in various tasks but almost the same number of parameters.
Revised: Compared to MHSA, skip-left and skip-right have more generalization in various tasks but almost the same number of parameters.

Original: Being special cases of skip-right, Figure 5c ( trainable heads-only C = C 1 ) and Figure 5d ( heads-only C = I H ) are less expressive.
Revised: Being special cases of skip-right, Figure 5c ( trainable heads-only C = C 1 ) and Figure 5d ( heads-only C = I H ) are less expressive.

Original: However, trainable heads-only only has around 6% of the parameters compared to MHSA while heads-only only has 1 . 5% with most performance maintained.
Revised: However, trainable heads-only only has around 6% of the parameters compared to MHSA while heads-only only has 1 . 5% with most performance maintained.

Original: There is a trade-off between generalization and efﬁciency in the design space C as demonstrated in
Revised: There is a trade-off between generalization and efﬁciency in the design space C as demonstrated in

Original: Figure 6.
Revised: Figure 6.

Original: We empirically verify that without over-ﬁtting, generalization and efﬁciency have similar hierarchy as expressive power.
Revised: We empirically verify that without over-ﬁtting, generalization and efﬁciency have similar hierarchy as expressive power.

Original: Thus we can design structures in different level of expressive power, tailoring needs for either better generalization or more efﬁciency.
Revised: Thus we can design structures in different level of expressive power, tailoring needs for either better generalization or more efﬁciency.

Original: THSA has the best generalization, but 125% parameters compared to MHSA.
Revised: THSA has the best generalization, but 125% parameters compared to MHSA.

Original: Since skip-left and skip-right are special cases of THSA, it is natural that their expressive power and generalization are lower.
Revised: Since skip-left and skip-right are special cases of THSA, it is natural that their expressive power and generalization are lower.

Original: However, they are more efﬁcient.
Revised: However, they are more efﬁcient.

Original: With almost the same number of parameters, they have better generalization than MHSA.
Revised: With almost the same number of parameters, they have better generalization than MHSA.

Original: Heads-only and trainable heads-only are very efﬁciency sacriﬁcing an acceptable amount generalization.
Revised: Heads-only and trainable heads-only are very efﬁciency sacriﬁcing an acceptable amount generalization.

Original: More details can be found in Section 4.
Revised: More details can be found in Section 4.

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Datasets, tasks and evaluation metrics.
Revised: Datasets, tasks and evaluation metrics.

Original: We evaluate Tuformers on 7 datasets for 5 tasks: word-level
Revised: We evaluate Tuformers on 7 datasets for 5 tasks: word-level

Original: 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.
Revised: 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.

Original: Baselines.
Revised: Baselines.

Original: We compare the performance of MHSA module with THSA module in the followingbackbone 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).
Revised: We compare the performance of MHSA module with THSA module in the followingbackbone 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).

Original: In all baselines, we adopt all default settings used in the original paper (hyperparameters speciﬁed in Appendix E), and only replace/re-parameterize the MHSA to THSA.
Revised: In all baselines, we adopt all default settings used in the original paper (hyperparameters speciﬁed in Appendix E), and only replace/re-parameterize the MHSA to THSA.

Original: 4. E MPIRICAL R ESULTS OF T UFORMERS .
Revised: 4. E MPIRICAL R ESULTS OF T UFORMERS .

Original: In this part, we ﬁrst show that THSA (used in Tuformers) universally outperforms MHSA (used in
Revised: In this part, we ﬁrst show that THSA (used in Tuformers) universally outperforms MHSA (used in

Original: Transformers) diverse tasks under different model scales and THSA is simple to use in practice.
Revised: Transformers) diverse tasks under different model scales and THSA is simple to use in practice.

Original: (1) THSA consistently outperforms MHSA on diverse datasets for a variety of tasks under different model scales.
Revised: (1) THSA consistently outperforms MHSA on diverse datasets for a variety of tasks under different model scales.

Original: In LM, NMT and ASR tasks, we compare THSA against the MHSA models under different model scales.
Revised: In LM, NMT and ASR tasks, we compare THSA against the MHSA models under different model scales.

Original: MHSA is parameterized by four weight matrices W Q , W K , W V and W O whereas THSA adds an additional core matrix C of the same size.
Revised: MHSA is parameterized by four weight matrices W Q , W K , W V and W O whereas THSA adds an additional core matrix C of the same size.

Original: Adopting the size of the four original weight matrices in MHSA to THSA incurs an 25% increase in the number of parameters.
Revised: Adopting the size of the four original weight matrices in MHSA to THSA incurs an 25% increase in the number of parameters.

Original: For a fair comparison, we ﬁrst compare MHSA against THSA with the same number of parameters (thus with smaller query/key/value/output weight matrices).
Revised: For a fair comparison, we ﬁrst compare MHSA against THSA with the same number of parameters (thus with smaller query/key/value/output weight matrices).

Original: To get a full picture, we also compare
Revised: To get a full picture, we also compare

Original: MHSA against THSA with the same sized four original weight matrices.
Revised: MHSA against THSA with the same sized four original weight matrices.

Original: As shown in Figure 7, under the same number of parameters, our THSA ( blue curve) universally outperforms MHSA ( red curve) under all model scales.
Revised: As shown in Figure 7, under the same number of parameters, our THSA ( blue curve) universally outperforms MHSA ( red curve) under all model scales.

Original: By increasing the number of parameters of the attention units by 25%(i.e.
Revised: By increasing the number of parameters of the attention units by 25%(i.e.

Original: , 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.
Revised: , 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.

Original: (2) Tuformers initialized with pre-trained strong models further improve the performance.
Revised: (2) Tuformers initialized with pre-trained strong models further improve the performance.

Original: In NLI task, we initialize Tuformers with pre-trained BERT-large, RoBERTa-large and ALBERT models, and perform the same ﬁne-tuning tasks.
Revised: In NLI task, we initialize Tuformers with pre-trained BERT-large, RoBERTa-large and ALBERT models, and perform the same ﬁne-tuning tasks.

Original: As shown in Table 1, THSA is compatible with
Revised: As shown in Table 1, THSA is compatible with

Original: [Figure]
Revised: [Figure]

Original: CIFAR-10 datasets, the lower the better.
Revised: CIFAR-10 datasets, the lower the better.

Original: Results show that Tuformers can be extended to image generation tasks and improve the performance of other efﬁcient designs with linear computational and memory complexities.
Revised: Results show that Tuformers can be extended to image generation tasks and improve the performance of other efﬁcient designs with linear computational and memory complexities.

Original: [Table]
Revised: [Table]

Original: existing pre-trained Transformer models and there is no need to train Tuformers from scratch:
Revised: existing pre-trained Transformer models and there is no need to train Tuformers from scratch:

Original: ﬁnetuning THSA initialized from pre-trained BERT and its variants further improves their performance.
Revised: ﬁnetuning THSA initialized from pre-trained BERT and its variants further improves their performance.

Original: (3) Combing THSA with kernels to obtain linear computation and memory complexities in sequence length.
Revised: (3) Combing THSA with kernels to obtain linear computation and memory complexities in sequence length.

Original: Tuformers can also be incorporated into efﬁcient variants of Transformers for lower computation and memory complexities.
Revised: Tuformers can also be incorporated into efﬁcient variants of Transformers for lower computation and memory complexities.

Original: We experiment on image generation tasks on CIFAR10 (Krizhevsky, 2009) and MNIST (Deng, 2012) datasets as in Table 1 and show that Tuformers obtain strong performance and linear space complexity when combined with kernels.
Revised: We experiment on image generation tasks on CIFAR10 (Krizhevsky, 2009) and MNIST (Deng, 2012) datasets as in Table 1 and show that Tuformers obtain strong performance and linear space complexity when combined with kernels.

Original: 4. F LEXIBLE DESIGN SPACE
Revised: 4. F LEXIBLE DESIGN SPACE

Original: (1) THSA reveals a ﬂexible design space C for Transformer.
Revised: (1) THSA reveals a ﬂexible design space C for Transformer.

Original: We empirically demonstrate that there exists a trade-off between generalization and efﬁciency in this design space, which could be adjusted by the expressive power.
Revised: We empirically demonstrate that there exists a trade-off between generalization and efﬁciency in this design space, which could be adjusted by the expressive power.

Original: By making C fully-trainable, partially trainable or pre-set ﬁxed as shown in
Revised: By making C fully-trainable, partially trainable or pre-set ﬁxed as shown in

Original: Figure 5 and Figure 4, we demonstrate that some structures in C improves the generalization while some other designs improve the efﬁciency of MHSA.
Revised: Figure 5 and Figure 4, we demonstrate that some structures in C improves the generalization while some other designs improve the efﬁciency of MHSA.

Original: [Table]
Revised: [Table]

Original: (2) Models with higher expressive power in the design space have better generalization.
Revised: (2) Models with higher expressive power in the design space have better generalization.

Original: As shown in
Revised: As shown in

Original: Table 2 and Figure 7, we compare two ways that increase the expressive power for better generalization: naively increasing the number of parameters and designing more expressive structures.
Revised: Table 2 and Figure 7, we compare two ways that increase the expressive power for better generalization: naively increasing the number of parameters and designing more expressive structures.

Original: Both approaches increase generalization.
Revised: Both approaches increase generalization.

Original: However, naively increasing the number of parameters is not as effective as adopting more expressive structures.
Revised: However, naively increasing the number of parameters is not as effective as adopting more expressive structures.

Original: With the same number of parameters, more expressive structures, such as THSA, skip-left and skip-right, universally outperform MHSA.
Revised: With the same number of parameters, more expressive structures, such as THSA, skip-left and skip-right, universally outperform MHSA.

Original: (3) Models can also be made very efﬁcient.
Revised: (3) Models can also be made very efﬁcient.

Original: Moreover, in Figure 5c ( trainable heads-only C = C 1 ), we trade generalization for efﬁciency.
Revised: Moreover, in Figure 5c ( trainable heads-only C = C 1 ), we trade generalization for efﬁciency.

Original: With 6% parameters, we maintain 95% performance in the
Revised: With 6% parameters, we maintain 95% performance in the

Original: NMT task.
Revised: NMT task.

Original: In a more extreme case, if we only maintain the head contraction as shown in Figure 5d, the number of parameters goes down to 1 .
Revised: In a more extreme case, if we only maintain the head contraction as shown in Figure 5d, the number of parameters goes down to 1 .

Original: 5% with a 42% faster training and 30% faster testing.
Revised: 5% with a 42% faster training and 30% faster testing.

Original: Surprisingly, almost 94% performance is obtained with this light-weighted model.
Revised: Surprisingly, almost 94% performance is obtained with this light-weighted model.

Original: 4. A BLATION STUDY : CHECK - POINTING TO ALLEVIATE MEMORY OVERHEAD WITHOUT KERNELS .
Revised: 4. A BLATION STUDY : CHECK - POINTING TO ALLEVIATE MEMORY OVERHEAD WITHOUT KERNELS .

Original: Although an effective design with guaranteed higher expressive power, Tuformers come with an extra memory overhead when calculating values within the softmax.
Revised: Although an effective design with guaranteed higher expressive power, Tuformers come with an extra memory overhead when calculating values within the softmax.

Original: 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.
Revised: 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.

Original: Combining Tuformers with some state-of-the-art kernel-based efﬁcient Transformers such as Linear Transformers using ELU (Katharopoulos et al., 2020) and polynomial
Revised: Combining Tuformers with some state-of-the-art kernel-based efﬁcient Transformers such as Linear Transformers using ELU (Katharopoulos et al., 2020) and polynomial

Original: (Tsai et al., 2019) kernels, we reduce the computation and memory complexities of Tuformers to be linear in the sequence length by removing the nonlinearity constrain in the softmax.
Revised: (Tsai et al., 2019) kernels, we reduce the computation and memory complexities of Tuformers to be linear in the sequence length by removing the nonlinearity constrain in the softmax.

Original: We can also alleviate the D times larger memory overhead compared with vanilla Transformers in calculating softmax without using kernels.
Revised: We can also alleviate the D times larger memory overhead compared with vanilla Transformers in calculating softmax without using kernels.

Original: Speciﬁcally, we use check-pointing to eliminate the memory overhead in THSA.
Revised: Speciﬁcally, we use check-pointing to eliminate the memory overhead in THSA.

Original: Since no intermediate results are saved in the forward pass using checkpointing, we recalculate the intermediate results of the softmax box in the back-propagation, which introduces some computational overhead.
Revised: Since no intermediate results are saved in the forward pass using checkpointing, we recalculate the intermediate results of the softmax box in the back-propagation, which introduces some computational overhead.

Original: We observe that applying check-pointing to Tuformers leads to only a slight increase in training time while relieving the memory overhead.
Revised: We observe that applying check-pointing to Tuformers leads to only a slight increase in training time while relieving the memory overhead.

Original: We can reduce or even avoid the repetitive calculations by developing a proper backward module, although such development is beyond the scope of this paper and deferred for future works.
Revised: We can reduce or even avoid the repetitive calculations by developing a proper backward module, although such development is beyond the scope of this paper and deferred for future works.

Original: See Figure 12 for more details.
Revised: See Figure 12 for more details.

Original:  R ELATED W ORKS
Revised:  R ELATED W ORKS

Original: Analysis of MHSA.
Revised: Analysis of MHSA.

Original: Given the popularity of Transformers in a wide range of domains, a line of works focuses on understanding and improving MHSA.
Revised: Given the popularity of Transformers in a wide range of domains, a line of works focuses on understanding and improving MHSA.

Original: Voita et al.
Revised: Voita et al.

Original: (2019) proposed a practical scheme to prune less informative heads.
Revised: (2019) proposed a practical scheme to prune less informative heads.

Original: Michel et al.
Revised: Michel et al.

Original: (2019) also showed that most attention heads are removable in testing without much performance compromise.
Revised: (2019) also showed that most attention heads are removable in testing without much performance compromise.

Original: Cordonnier et al.
Revised: Cordonnier et al.

Original: (2020) found empirical evidence of redundancy in key/query projections and proposed a re-parameterization scheme.
Revised: (2020) found empirical evidence of redundancy in key/query projections and proposed a re-parameterization scheme.

Original: These works bring valuable insights into MHSA, but a consistent theoretical analysis for these ﬁndings is missing.
Revised: These works bring valuable insights into MHSA, but a consistent theoretical analysis for these ﬁndings is missing.

Original: 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.
Revised: 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.

Original: Efﬁcient Transformers.
Revised: Efﬁcient Transformers.

Original: MHSA has a global receptive ﬁeld, 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.
Revised: MHSA has a global receptive ﬁeld, 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.

Original: 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).
Revised: 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).

Original: Our proposed model, Tuformers, can be incorporated into some kernel-based efﬁcient 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.
Revised: Our proposed model, Tuformers, can be incorporated into some kernel-based efﬁcient 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.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: 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.
Revised: 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.

Original: Furthermore,
Revised: Furthermore,

Original: Tuformers have a data-driven structure where heads can be trained implicitly and initialized with pre-trained Transformer models.
Revised: Tuformers have a data-driven structure where heads can be trained implicitly and initialized with pre-trained Transformer models.

Original: 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 efﬁcient architectures.
Revised: 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 efﬁcient architectures.

Original: A S UPPLEMENTARY M ATERIAL FOR T ENSOR D IAGRAMS
Revised: A S UPPLEMENTARY M ATERIAL FOR T ENSOR D IAGRAMS

Original: In this section, we provide a comprehensive introduction to the tensor diagram representation.
Revised: In this section, we provide a comprehensive introduction to the tensor diagram representation.

Original: A. C OMPONENTS OF T ENSOR D IAGRAM : A RRAYS AS N ODES WITH L EG ( S )
Revised: A. C OMPONENTS OF T ENSOR D IAGRAM : A RRAYS AS N ODES WITH L EG ( S )

Original: Arrays denoted as Nodes with Legs.
Revised: Arrays denoted as Nodes with Legs.

Original: An array is represented as a node with leg(s) in a tensor diagram as shown in Figure 1.
Revised: An array is represented as a node with leg(s) in a tensor diagram as shown in Figure 1.

Original: The order (the number of dimensions) of the array is denoted by the number of legs extending from the node.
Revised: The order (the number of dimensions) of the array is denoted by the number of legs extending from the node.

Original: The legs do not need to be straight lines ,and the orientation does not matter.
Revised: The legs do not need to be straight lines ,and the orientation does not matter.

Original: For example, matrices A and A (cid:62) are equivalent in this notation.
Revised: For example, matrices A and A (cid:62) are equivalent in this notation.

Original: Each labeled leg represents one mode of a tensor.
Revised: Each labeled leg represents one mode of a tensor.

Original: Every mode of the tensor needs to be uniquely labeled.
Revised: Every mode of the tensor needs to be uniquely labeled.

Original: We usually use the dimension of the mode as the label (i.e., an associated positive integer written on top of each leg).
Revised: We usually use the dimension of the mode as the label (i.e., an associated positive integer written on top of each leg).

Original: If multiple modes of the same tensor have the same length, we use subscripts to differentiate them.
Revised: If multiple modes of the same tensor have the same length, we use subscripts to differentiate them.

Original: A. O PERATIONS IN T ENSOR D IAGRAM
Revised: A. O PERATIONS IN T ENSOR D IAGRAM

Original: [Equation]
Revised: [Equation]

Original: There are three types of operations: contraction , softmax and batch multiplication .
Revised: There are three types of operations: contraction , softmax and batch multiplication .

Original: 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.
Revised: 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.

Original: Tensor diagram has a convention to denote the contraction, but not the softmax.
Revised: Tensor diagram has a convention to denote the contraction, but not the softmax.

Original: We will ﬁrst introduce contraction in tensor diagram notation and then propose our design of softmax in tensor diagram language.
Revised: We will ﬁrst introduce contraction in tensor diagram notation and then propose our design of softmax in tensor diagram language.

Original: Contractions denoted as edges connecting the node legs.
Revised: Contractions denoted as edges connecting the node legs.

Original: 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 .
Revised: 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 .

Original: Since the representation of arrays in the tensor diagram is orientation invariant, contractions are also orientation invariant.
Revised: Since the representation of arrays in the tensor diagram is orientation invariant, contractions are also orientation invariant.

Original: This is especially useful in representing the self-attention unit.
Revised: This is especially useful in representing the self-attention unit.

Original: In some cases, the data object are represented as row vectors while in some other scenarios they are denoted as column vectors.
Revised: In some cases, the data object are represented as row vectors while in some other scenarios they are denoted as column vectors.

Original: If we use mathematical formula, we need to take care of the ordering of the matrix multiplication because A (cid:62) B , AB (cid:62) , BA (cid:62) , B (cid:62) A and more are all different.
Revised: If we use mathematical formula, we need to take care of the ordering of the matrix multiplication because A (cid:62) B , AB (cid:62) , BA (cid:62) , B (cid:62) A and more are all different.

Original: In contrast, tensor diagram provides a universal graphical representation as long as the corresponding legs are connected correctly.
Revised: In contrast, tensor diagram provides a universal graphical representation as long as the corresponding legs are connected correctly.

Original: Nonlinear activation denoted as dotted box with a labelled ﬁlled ball attached.
Revised: Nonlinear activation denoted as dotted box with a labelled ﬁlled ball attached.

Original: We propose to denote the α -scaled softmax function softmax ( α A ) on A as a dotted box with a labeled ﬁlled ball (to distinguish itself from tensor objects, i.e., nodes which are blank circles) attached to one leg.
Revised: We propose to denote the α -scaled softmax function softmax ( α A ) on A as a dotted box with a labeled ﬁlled ball (to distinguish itself from tensor objects, i.e., nodes which are blank circles) attached to one leg.

Original: The label right above the ﬁlled 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.
Revised: The label right above the ﬁlled 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.

Original: As shown in
Revised: As shown in

Original: Appendix A.2, let A denote the resultant of some operations (for example Q [ h ] K (cid:62)
Revised: Appendix A.2, let A denote the resultant of some operations (for example Q [ h ] K (cid:62)

Original: [ h ] ), the 1 √ D -scaled softmax on A which normalize along leg N is depicted.
Revised: [ h ] ), the 1 √ D -scaled softmax on A which normalize along leg N is depicted.

Original: Softmax does not change the shape of the input.
Revised: Softmax does not change the shape of the input.

Original: That is why we introduce a dotted box, which does not change the shape either, to denote it.
Revised: That is why we introduce a dotted box, which does not change the shape either, to denote it.

Original: Batch Multiplication Given A and B , the batch multiplication is an element-wise product along the connected legs.
Revised: Batch Multiplication Given A and B , the batch multiplication is an element-wise product along the connected legs.

Original: The resultant tensor C is obtained by merging connected nodes, maintaining the dangling legs M , N and R .
Revised: The resultant tensor C is obtained by merging connected nodes, maintaining the dangling legs M , N and R .

Original: 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.
Revised: 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.

Original: A. A DVANTAGES OF T ENSOR D IAGRAM
Revised: A. A DVANTAGES OF T ENSOR D IAGRAM

Original: Tensor diagram representations enjoy a few advantages compared with using ﬂow charts together with math formulas.
Revised: Tensor diagram representations enjoy a few advantages compared with using ﬂow charts together with math formulas.

Original: (1) The notation is orientation invariant.
Revised: (1) The notation is orientation invariant.

Original: The input embedding matrix X can be represented as either R N × F or R F × N with rows or columns being the embedding of each token in the sequence respectively.
Revised: The input embedding matrix X can be represented as either R N × F or R F × N with rows or columns being the embedding of each token in the sequence respectively.

Original: 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.
Revised: 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.

Original: 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 9c.
Revised: 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 9c.

Original: The representation is both precise in math and intuitive, making analysis easier.
Revised: The representation is both precise in math and intuitive, making analysis easier.

Original: The tensor diagram itself is mathematically rigorous.
Revised: The tensor diagram itself is mathematically rigorous.

Original: 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 W Q , W K and W V as in mathematical formulas.
Revised: 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 W Q , W K and W V as in mathematical formulas.

Original: 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.
Revised: 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.

Original: This is particularly convenient to obtain the size of the resultant: any
Revised: This is particularly convenient to obtain the size of the resultant: any

Original: the
Revised: the

Original: dangling leg
Revised: dangling leg

Original: becomes a leg of the resultant.
Revised: becomes a leg of the resultant.

Original: (3) Multi-linear operations are concisely represented and interpreted.
Revised: (3) Multi-linear operations are concisely represented and interpreted.

Original: Using tensor diagram, we can concisely represent the multi-head self-attention rigorously in one diagram as shown in Figure 9c.
Revised: Using tensor diagram, we can concisely represent the multi-head self-attention rigorously in one diagram as shown in Figure 9c.

Original: 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 W Q , W K and W B , is clearly illustrated in the tensor diagram.
Revised: 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 W Q , W K and W B , is clearly illustrated in the tensor diagram.

Original: This illustration, which is not achieved by any existing representations, is crucial for our principled understanding of the multi-head self-attention.
Revised: This illustration, which is not achieved by any existing representations, is crucial for our principled understanding of the multi-head self-attention.

Original: (4) The information ﬂow is clear.
Revised: (4) The information ﬂow is clear.

Original: For example, in Figure 9a there is no direct information pass from W K [ h ] , W Q [ h ] to W V [ h ] while in Figure 9c, such exchange is done through edge D s and H s. 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
Revised: For example, in Figure 9a there is no direct information pass from W K [ h ] , W Q [ h ] to W V [ h ] while in Figure 9c, such exchange is done through edge D s and H s. 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

Original: B S UPPLEMENTARY M ATERIAL FOR M ULTI -H EAD S ELF -A TTENTION
Revised: B S UPPLEMENTARY M ATERIAL FOR M ULTI -H EAD S ELF -A TTENTION

Original: In the section, we prove the equivalence between the tensor diagram and the equations for multi-head self-attention (MHSA).
Revised: In the section, we prove the equivalence between the tensor diagram and the equations for multi-head self-attention (MHSA).

Original: 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.
Revised: 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.

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F 1 D D
Revised: F 1 D D

Original: D
Revised: D

Original: W Q [ h ] W K [ h ]
Revised: W Q [ h ] W K [ h ]

Original: W V [ h ]
Revised: W V [ h ]

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K
Revised: X K

Original: [Figure]
Revised: [Figure]

Original: N
Revised: N

Original: M
Revised: M

Original: F
Revised: F

Original: F
Revised: F

Original: F F 1 D D
Revised: F F 1 D D

Original: H H
Revised: H H

Original: H
Revised: H

Original: H 𝒲 Q
Revised: H 𝒲 Q

Original: [Figure]
Revised: [Figure]

Original: K
Revised: K

Original: [Figure]
Revised: [Figure]

Original: F
Revised: F

Original: [Figure]
Revised: [Figure]

Original: K
Revised: K

Original: [Figure]
Revised: [Figure]

Original: V
Revised: V

Original: [Figure]
Revised: [Figure]

Original: For convenience of reference, we recap both representations in Equation (B.1) and Figure 9.
Revised: For convenience of reference, we recap both representations in Equation (B.1) and Figure 9.

Original: [Equation]
Revised: [Equation]

Original: B.1a)
Revised: B.1a)

Original: [Equation]
Revised: [Equation]

Original: B.1b)
Revised: B.1b)

Original: [Equation]
Revised: [Equation]

Original: B.1c)
Revised: B.1c)

Original: Note that the query weight tensor W Q comes from a concatenation of H query weight matrices { W
Revised: Note that the query weight tensor W Q comes from a concatenation of H query weight matrices { W

Original: [Equation]
Revised: [Equation]

Original: =1 such that W Qh, : , : = W Q
Revised: =1 such that W Qh, : , : = W Q

Original: [ h ] .
Revised: [ h ] .

Original: Similarly, W K , W V are concatenated from { W K [ h ] } Hh =1 , { W V [ h ] } Hh =1 such that W Kh, : , : =
Revised: Similarly, W K , W V are concatenated from { W K [ h ] } Hh =1 , { W V [ h ] } Hh =1 such that W Kh, : , : =

Original: W K [ h ] , W Vh, : , : =
Revised: W K [ h ] , W Vh, : , : =

Original: W V [ h ] .
Revised: W V [ h ] .

Original: Furthermore, we denote Q , K , V as the reshaped tensors from { Q [ h ] } Hh =1 , { K [ h ] } Hh =1 , { V [ h ] } Hh =1 .
Revised: Furthermore, we denote Q , K , V as the reshaped tensors from { Q [ h ] } Hh =1 , { K [ h ] } Hh =1 , { V [ h ] } Hh =1 .

Original: Finally, W O is reshaped from W O .
Revised: Finally, W O is reshaped from W O .

Original: 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.
Revised: 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.

Original: B. S INGLE -H EAD S ELF -A TTENTION (SHSA)
Revised: B. S INGLE -H EAD S ELF -A TTENTION (SHSA)

Original: We will prove that the tensor diagram in Figure 9a is equivalent to Equations (B.1a) and (B.1b).
Revised: We will prove that the tensor diagram in Figure 9a is equivalent to Equations (B.1a) and (B.1b).

Original: Proof.
Revised: Proof.

Original: For simplicity, we omit the subscript [ h ] in this subsection.
Revised: For simplicity, we omit the subscript [ h ] in this subsection.

Original: We ﬁrst show that a contraction of the nodes within the softmax box is equivalent to QK (cid:62) .
Revised: We ﬁrst show that a contraction of the nodes within the softmax box is equivalent to QK (cid:62) .

Original: Let
Revised: Let

Original: A = QK (cid:62) =
Revised: A = QK (cid:62) =

Original: X Q W Q ( W K X K ) (cid:62) , where X Q ∈ R N × F , X K ∈ R M × F , W Q , W K ∈ R F × D .
Revised: X Q W Q ( W K X K ) (cid:62) , where X Q ∈ R N × F , X K ∈ R M × F , W Q , W K ∈ R F × D .

Original: We use subscripts in the element-wise notation to differentiate edges when multiple edges with same dimension appear in the same equation.
Revised: We use subscripts in the element-wise notation to differentiate edges when multiple edges with same dimension appear in the same equation.

Original: [Equation]
Revised: [Equation]

Original: B.2a)
Revised: B.2a)

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Denote be the result of
Revised: Denote be the result of

Original: [Equation]
Revised: [Equation]

Original: Comparing Equations (B.2c) and (B.3), we have A = A .
Revised: Comparing Equations (B.2c) and (B.3), we have A = A .

Original: Let B = softmax( A ) and B = softmax( A ) , we further have B = softmax( A ) =
Revised: Let B = softmax( A ) and B = softmax( A ) , we further have B = softmax( A ) =

Original: softmax( A ) =
Revised: softmax( A ) =

Original: B .
Revised: B .

Original: Let M be the result of the tensor diagram and M be the matrix representation of the single-head attention.
Revised: Let M be the result of the tensor diagram and M be the matrix representation of the single-head attention.

Original: We have
Revised: We have

Original: [Equation]
Revised: [Equation]

Original: Comparing Equations (B.4) and (B.5), we have C is equivalent to C , which completes the proof.
Revised: Comparing Equations (B.4) and (B.5), we have C is equivalent to C , which completes the proof.

Original: B. M ULTI - HEAD S ELF -A TTENTION (MHSA)
Revised: B. M ULTI - HEAD S ELF -A TTENTION (MHSA)

Original: We now prove that tensor diagram in Figure 9c is equivalent to Equation (1c).
Revised: We now prove that tensor diagram in Figure 9c is equivalent to Equation (1c).

Original: Proof.
Revised: Proof.

Original: Let H be the number of heads.
Revised: Let H be the number of heads.

Original: Let A ∈ R H
Revised: Let A ∈ R H

Original: × N × M be the result of a contraction of the nodes within the softmax box.
Revised: × N × M be the result of a contraction of the nodes within the softmax box.

Original: From the last section, we know that A h, : , : =
Revised: From the last section, we know that A h, : , : =

Original: A [ h ] = Q [ h ] K (cid:62) [ h ] .
Revised: A [ h ] = Q [ h ] K (cid:62) [ h ] .

Original: Applying the softmax function B = softmax( A ) , we have B h, : , : = softmax( Q [ h ] K (cid:62)
Revised: Applying the softmax function B = softmax( A ) , we have B h, : , : = softmax( Q [ h ] K (cid:62)

Original: [ h ] ) .
Revised: [ h ] ) .

Original: Let T be result of the tensor diagram and T be that of Equation (B.1c).
Revised: Let T be result of the tensor diagram and T be that of Equation (B.1c).

Original: [Equation]
Revised: [Equation]

Original: Comparing T and T , we complete the proof.
Revised: Comparing T and T , we complete the proof.

Original: C C LOSER I NVESTIGATION OF MHSA AND ITS C OMPARISON TO THSA
Revised: C C LOSER I NVESTIGATION OF MHSA AND ITS C OMPARISON TO THSA

Original: Rethinking MHSA.
Revised: Rethinking MHSA.

Original: There are two types contractions in MHSA as shown in Figure 4a: HeadContraction , a global contraction among all weights W Q , W K , W V and W O along mode H , and
Revised: There are two types contractions in MHSA as shown in Figure 4a: HeadContraction , a global contraction among all weights W Q , W K , W V and W O along mode H , and

Original: Latent-Contraction , a local contraction between W Q and W K as well as between W V and W O both along mode D .
Revised: Latent-Contraction , a local contraction between W Q and W K as well as between W V and W O both along mode D .

Original: 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 W O .
Revised: 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 W O .

Original: 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.
Revised: 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.

Original: It turns out that in both situation, MHSA falls into a special case of THSA.
Revised: It turns out that in both situation, MHSA falls into a special case of THSA.

Original: Repeated Proposition 1 with two special cases.
Revised: Repeated Proposition 1 with two special cases.

Original: Given R = DH and other hyperparameters being the same, THSA module includes MHSA as a special case.
Revised: Given R = DH and other hyperparameters being the same, THSA module includes MHSA as a special case.

Original: Speciﬁcally,
Revised: Speciﬁcally,

Original: (1) If C = 1 D 1 (cid:62) 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 ).
Revised: (1) If C = 1 D 1 (cid:62) 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 ).

Original: (2) If C = I H , 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).
Revised: (2) If C = I H , 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).

Original: (3) If C = I
Revised: (3) If C = I

Original: H ⊗ ( 1 D 1 (cid:62) D ) , i.e., a Kronecker product of an all-ones matrix 1 D 1 (cid:62) D ∈ R D
Revised: H ⊗ ( 1 D 1 (cid:62) D ) , i.e., a Kronecker product of an all-ones matrix 1 D 1 (cid:62) D ∈ R D

Original: × D and an identity matrix
Revised: × D and an identity matrix

Original: I H ∈
Revised: I H ∈

Original: R H
Revised: R H

Original: × H , the THSA reduces to an MHSA with H heads and latent dimension D .
Revised: × H , the THSA reduces to an MHSA with H heads and latent dimension D .

Original: The proof of the above propositions is given in Appendix D.
Revised: The proof of the above propositions is given in Appendix D.

Original: Deﬁnition 2 (Stable Rank) .
Revised: Deﬁnition 2 (Stable Rank) .

Original: Let A ∈ R N × M , the stable rank of A of is deﬁned as (cid:80) i σ
Revised: Let A ∈ R N × M , the stable rank of A of is deﬁned as (cid:80) i σ

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: where σ i is the singular value of A .
Revised: where σ i is the singular value of A .

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F F 1 D
Revised: F F 1 D

Original: D
Revised: D

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K
Revised: X K

Original: D stable rank =
Revised: D stable rank =

Original: W Q
Revised: W Q

Original: W K
Revised: W K

Original: W O
Revised: W O

Original: W V 1 D 1 D
Revised: W V 1 D 1 D

Original: N
Revised: N

Original: M
Revised: M

Original: F
Revised: F

Original: F F F
Revised: F F F

Original: H
Revised: H

Original: W Q
Revised: W Q

Original: W K
Revised: W K

Original: W O
Revised: W O

Original: W V stable rank
Revised: W V stable rank

Original: =
Revised: =

Original: H
Revised: H

Original: I H
Revised: I H

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K
Revised: X K

Original: N
Revised: N

Original: M
Revised: M

Original: F F
Revised: F F

Original: F F 1 R 𝒲 Q 𝒲 K 𝒲 O 𝒲 V
Revised: F F 1 R 𝒲 Q 𝒲 K 𝒲 O 𝒲 V

Original: I H
Revised: I H

Original: X Q
Revised: X Q

Original: X V
Revised: X V

Original: X K 1 D 1 D
Revised: X K 1 D 1 D

Original: D
Revised: D

Original: D
Revised: D

Original: H
Revised: H

Original: [Figure]
Revised: [Figure]

Original: The notion of heads in THSA.
Revised: The notion of heads in THSA.

Original: The concept of the number of heads in THSA generalizes to the stable rank of the core matrix, allowing a data-driven implicit training.
Revised: The concept of the number of heads in THSA generalizes to the stable rank of the core matrix, allowing a data-driven implicit training.

Original: The number of heads H in
Revised: The number of heads H in

Original: MHSA corresponds to the stable rank of the core matrix C , deﬁned as (cid:80) i σ 2 i (cid:14) max i σ i where σ i is the singular value.
Revised: MHSA corresponds to the stable rank of the core matrix C , deﬁned as (cid:80) i σ 2 i (cid:14) max i σ i where σ i is the singular value.

Original: It is a useful surrogate for the rank because it is largely unaffected by tiny singular values.
Revised: It is a useful surrogate for the rank because it is largely unaffected by tiny singular values.

Original: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: Also, the structure on the right can be obtained by removing the latent contraction edges D from Figure 4a.
Revised: Also, the structure on the right can be obtained by removing the latent contraction edges D from Figure 4a.

Original: 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 ﬁxed matrix in MHSA.
Revised: 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 ﬁxed matrix in MHSA.

Original: D S UPPLEMENTARY M ATERIAL FOR T UFORMER AND P ROOFS
Revised: D S UPPLEMENTARY M ATERIAL FOR T UFORMER AND P ROOFS

Original: We ﬁrst formally describe how to compare the expressive power between two models.
Revised: We ﬁrst formally describe how to compare the expressive power between two models.

Original: Deﬁnition 3 (Expressive Power) .
Revised: Deﬁnition 3 (Expressive Power) .

Original: 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
Revised: 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

Original: Y .
Revised: Y .

Original: We say G is more expressive than F if G ⊇ F : for any f ∈ F , there exists g ∈ G such that g ( x ) =
Revised: We say G is more expressive than F if G ⊇ F : for any f ∈ F , there exists g ∈ G such that g ( x ) =

Original: f
Revised: f

Original: ( x ) , ∀ x ∈ X .
Revised: ( x ) , ∀ x ∈ X .

Original: Furthermore, we say G is strictly more expressive than F if G ⊃ F : besides
Revised: Furthermore, we say G is strictly more expressive than F if G ⊃ F : besides

Original: G ⊇ F , there exists g ∈ G such that given any f there exists x ∈ X and g ( x ) (cid:54) =
Revised: G ⊇ F , there exists g ∈ G such that given any f there exists x ∈ X and g ( x ) (cid:54) =

Original: f ( x ) .
Revised: f ( x ) .

Original: According to the deﬁnition, expressive power is a partial order set.
Revised: According to the deﬁnition, expressive power is a partial order set.

Original: Therefore, expressive power of different models is usually not comparable as two sets might have overlap and non-overlap regions.
Revised: Therefore, expressive power of different models is usually not comparable as two sets might have overlap and non-overlap regions.

Original: However, when one is a superset or subset of the other, we can unambiguously compare their expressive power.
Revised: However, when one is a superset or subset of the other, we can unambiguously compare their expressive power.

Original: 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
Revised: 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

Original: THSA is equal to the product of the number of heads H and the latent dimension D in MHSA.
Revised: THSA is equal to the product of the number of heads H and the latent dimension D in MHSA.

Original: Theorem 4 (THSA is more expressive than MHSA) .
Revised: Theorem 4 (THSA is more expressive than MHSA) .

Original: 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 .
Revised: 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 .

Original: Theorem 4, proved in Appendix D.2, shows that any function realizable by MHSA can be realized by
Revised: Theorem 4, proved in Appendix D.2, shows that any function realizable by MHSA can be realized by

Original: THSA.
Revised: THSA.

Original: In this section, we will prove Proposition 1 and Theorem 4 in Section 3.
Revised: In this section, we will prove Proposition 1 and Theorem 4 in Section 3.

Original: 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 .
Revised: 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 .

Original: 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.
Revised: 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.

Original: Multi-head Self-Attention (MHSA).
Revised: Multi-head Self-Attention (MHSA).

Original: An MHSA module has its learnable parameters as H sets of (query, key, value) weight matrices { W Q [ h ] , W K [ h ] , W V [ h ] } Hh =1 (with W Q [ h ] , W K [ h ] , W V [ h ] ∈ R F × D for each h ) and an output weight matrix W O ∈ R HD × F .
Revised: An MHSA module has its learnable parameters as H sets of (query, key, value) weight matrices { W Q [ h ] , W K [ h ] , W V [ h ] } Hh =1 (with W Q [ h ] , W K [ h ] , W V [ h ] ∈ R F × D for each h ) and an output weight matrix W O ∈ R HD × F .

Original: The module maps three input matrices
Revised: The module maps three input matrices

Original: X Q ∈ R N × F , X K , X V ∈ R M × F into an output matrix M ∈ R N × F .
Revised: X Q ∈ R N × F , X K , X V ∈ R M × F into an output matrix M ∈ R N × F .

Original: [Equation]
Revised: [Equation]

Original: In Equation (D.1a), Q [ h ] ∈ R N × D , K
Revised: In Equation (D.1a), Q [ h ] ∈ R N × D , K

Original: [ h ] , V [ h ] ∈ R M × D are query, key, value matrices respectively.
Revised: [ h ] , V [ h ] ∈ R M × D are query, key, value matrices respectively.

Original: 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 ] ∈ R N × D is the resulted matrix for head h .
Revised: 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 ] ∈ R N × D is the resulted matrix for head h .

Original: Tunable-Head Self-Attention.
Revised: Tunable-Head Self-Attention.

Original: A THSA module is parameterized by ﬁve matrices: a query weight matrix W Q ∈ R F × R , a key weight matrix W K ∈ R F × R , a value weight matrix W V ∈ R F × R , an output weight matrix W O ∈ R R × F , and an additional core matrix C ∈ R R × R .
Revised: A THSA module is parameterized by ﬁve matrices: a query weight matrix W Q ∈ R F × R , a key weight matrix W K ∈ R F × R , a value weight matrix W V ∈ R F × R , an output weight matrix W O ∈ R R × F , and an additional core matrix C ∈ R R × R .

Original: The module has the same input and output domains as in MHSA, i.e., it takes three matrices X Q ∈ R N × F ,
Revised: The module has the same input and output domains as in MHSA, i.e., it takes three matrices X Q ∈ R N × F ,

Original: X K , X V ∈ R M × F as inputs and returns a matrix T ∈ R N × F as output.
Revised: X K , X V ∈ R M × F as inputs and returns a matrix T ∈ R N × F as output.

Original: [Equation]
Revised: [Equation]

Original: To distinguish THSA from MHSA, we use over-scored symbols whenever needed.
Revised: To distinguish THSA from MHSA, we use over-scored symbols whenever needed.

Original: In Equation (D.2a),
Revised: In Equation (D.2a),

Original: Q ∈ R N × R , K , V ∈ R M × R are query, key, value matrices respectively.
Revised: Q ∈ R N × R , K , V ∈ R M × R are query, key, value matrices respectively.

Original: In Equation (D.2b), we again omit the scaling factor 1 / √ R in the softmax function.
Revised: In Equation (D.2b), we again omit the scaling factor 1 / √ R in the softmax function.

Original: Equation (D.2b) leads to a head matrix head ∈ R F × N , where head r ∈ R F denotes the r th column of the matrix head .
Revised: Equation (D.2b) leads to a head matrix head ∈ R F × N , where head r ∈ R F denotes the r th column of the matrix head .

Original: D. P ROOF OF P ROPOSITION
Revised: D. P ROOF OF P ROPOSITION

Original: Proof of Proposition 1.
Revised: Proof of Proposition 1.

Original: 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:
Revised: 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:

Original: [Equation]
Revised: [Equation]

Original: In addition, we set W O = W O .
Revised: In addition, we set W O = W O .

Original: We have W Q , W K , W V ∈ R F × R , and W O ∈ R R × F .
Revised: We have W Q , W K , W V ∈ R F × R , and W O ∈ R R × F .

Original: (1) For C = 1 R 1 (cid:62) 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 .
Revised: (1) For C = 1 R 1 (cid:62) 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 .

Original: Since there is only one head in the module, we have W Q = W Q [1] , W K = W K [1] , and W V = W V [1] .
Revised: Since there is only one head in the module, we have W Q = W Q [1] , W K = W K [1] , and W V = W V [1] .

Original: As an immediate result,
Revised: As an immediate result,

Original: [Equation]
Revised: [Equation]

Original: Since C is an all-ones matrix, i.e., C rs = 1 , ∀ r, s , we rewrite Equation (D.2b) as:
Revised: Since C is an all-ones matrix, i.e., C rs = 1 , ∀ r, s , we rewrite Equation (D.2b) as:

Original: [Equation]
Revised: [Equation]

Original: Notice that the equation holds for each column r , we further write all R equations jointly as:
Revised: Notice that the equation holds for each column r , we further write all R equations jointly as:

Original: [Equation]
Revised: [Equation]

Original: Finally, we express the output matrix as:
Revised: Finally, we express the output matrix as:

Original: [Equation]
Revised: [Equation]

Original: which concludes the reduction for the special case (1) .
Revised: which concludes the reduction for the special case (1) .

Original: (2) For C
Revised: (2) For C

Original: =
Revised: =

Original: I R , an R × R identity matrix, we aim to prove that the THSA module reduces to a heads-only self-attention with H
Revised: I R , an R × R identity matrix, we aim to prove that the THSA module reduces to a heads-only self-attention with H

Original: = R and D = 1 .
Revised: = R and D = 1 .

Original: Since the latent dimension D = 1 , i.e., each W Q [ h ] (or W K [ h ] , W V [ h ] ) is a vector, we have W Qr = W Q
Revised: Since the latent dimension D = 1 , i.e., each W Q [ h ] (or W K [ h ] , W V [ h ] ) is a vector, we have W Qr = W Q

Original: [ r ] , W Kr
Revised: [ r ] , W Kr

Original: = W K
Revised: = W K

Original: [ r ] , and W Vr = W V [ r ] .
Revised: [ r ] , and W Vr = W V [ r ] .

Original: Therefore,
Revised: Therefore,

Original: [Equation]
Revised: [Equation]

Original: Since C is an identity matrix, i.e., C rr = 1 and C rs = 0 , ∀ r (cid:54)
Revised: Since C is an identity matrix, i.e., C rr = 1 and C rs = 0 , ∀ r (cid:54)

Original: =
Revised: =

Original: s , we rewrite Equation (D.2b) as:
Revised: s , we rewrite Equation (D.2b) as:

Original: [Equation]
Revised: [Equation]

Original: Finally, we express the output matrix as:
Revised: Finally, we express the output matrix as:

Original: [Equation]
Revised: [Equation]

Original: W
Revised: W

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: 10) which concludes the reduction for the special case (2) .
Revised: 10) which concludes the reduction for the special case (2) .

Original: (3) For C
Revised: (3) For C

Original: =
Revised: =

Original: I H ⊗ ( 1 D 1 (cid:62) D ) , i.e., a Kronecker product between an identity matrix
Revised: I H ⊗ ( 1 D 1 (cid:62) D ) , i.e., a Kronecker product between an identity matrix

Original: I H ∈
Revised: I H ∈

Original: R H
Revised: R H

Original: × H and an all-ones matrix 1 D 1 (cid:62) D ∈ R D × 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 .
Revised: × H and an all-ones matrix 1 D 1 (cid:62) D ∈ R D × 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 .

Original: In this general case, we have
Revised: In this general case, we have

Original: [Equation]
Revised: [Equation]

Original: Notice that C is a block diagonal matrix blkdiag ( 1 D 1 (cid:62) D , · · · , 1 D 1 (cid:62) D ) , where each block is an all-ones matrix.
Revised: Notice that C is a block diagonal matrix blkdiag ( 1 D 1 (cid:62) D , · · · , 1 D 1 (cid:62) D ) , where each block is an all-ones matrix.

Original: Therefore we can rewrite Equation (D.2b) for r =
Revised: Therefore we can rewrite Equation (D.2b) for r =

Original: ( h − 1) D + d as:
Revised: ( h − 1) D + d as:

Original: [Equation]
Revised: [Equation]

Original: =
Revised: =

Original: [Equation]
Revised: [Equation]

Original: Since Equation (D.16) holds for each column d , we write all D equations jointly as:
Revised: Since Equation (D.16) holds for each column d , we write all D equations jointly as:

Original: [Equation]
Revised: [Equation]

Original: Again, since Equation (D.17) holds for each block h , we combine all H equations as:
Revised: Again, since Equation (D.17) holds for each block h , we combine all H equations as:

Original: [Equation]
Revised: [Equation]

Original: Finally, we express the output matrix as:
Revised: Finally, we express the output matrix as:

Original: [Equation]
Revised: [Equation]

Original: which concludes the reduction for the case (3) .
Revised: which concludes the reduction for the case (3) .

Original: D. P ROOF OF T HEOREM
Revised: D. P ROOF OF T HEOREM

Original: Proof of Theorem 4.
Revised: Proof of Theorem 4.

Original: We have shown that THSA reduces to MHSA when C takes speciﬁc forms.
Revised: We have shown that THSA reduces to MHSA when C takes speciﬁc forms.

Original: This shows that THSA is more expressive than MHSA — any mapping by an MHSA module can also be realized by a THSA module.
Revised: This shows that THSA is more expressive than MHSA — any mapping by an MHSA module can also be realized by a THSA module.

Original: To prove that THSA is strictly more expressive than MHSA, we need ﬁnd a mapping by a THSA module that can not be realized by an MHSA module.
Revised: To prove that THSA is strictly more expressive than MHSA, we need ﬁnd a mapping by a THSA module that can not be realized by an MHSA module.

Original: It sufﬁces to show that the matrix M is rank-deﬁcient, while the matrix T can be full-rank.
Revised: It sufﬁces to show that the matrix M is rank-deﬁcient, while the matrix T can be full-rank.

Original: We assume MH < min( N, F ) .
Revised: We assume MH < min( N, F ) .

Original: For convenience, we divide the matrix W O for MHSA into H blocks.
Revised: For convenience, we divide the matrix W O for MHSA into H blocks.

Original: [Equation]
Revised: [Equation]

Original: where W O [ h ] ∈ R D
Revised: where W O [ h ] ∈ R D

Original: × F is the matrix for the h th head.
Revised: × F is the matrix for the h th head.

Original: We now rewrite the output matrix M as:
Revised: We now rewrite the output matrix M as:

Original: [Equation]
Revised: [Equation]

Original: Since the matrix inside summation is a product of two matrices of size R N × M and R M × F , its rank is at most M .
Revised: Since the matrix inside summation is a product of two matrices of size R N × M and R M × F , its rank is at most M .

Original: Use the property rank ( A + B ) ≤ rank ( A ) + rank ( B ) , we have rank ( M ) ≤ MH <
Revised: Use the property rank ( A + B ) ≤ rank ( A ) + rank ( B ) , we have rank ( M ) ≤ MH <

Original: min
Revised: min

Original: N, F ) .
Revised: N, F ) .

Original: Similarly, we divide the matrix W O for THSA into R rows and rewrite the output matrix T as:
Revised: Similarly, we divide the matrix W O for THSA into R rows and rewrite the output matrix T as:

Original: [Equation]
Revised: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Using the same argument, we have rank ( T ) ≤ MR .
Revised: Using the same argument, we have rank ( T ) ≤ MR .

Original: Since H < R (in fact, R = DH ), we can always ﬁnd
Revised: Since H < R (in fact, R = DH ), we can always ﬁnd

Original: an example such that rank ( T ) > rank ( M ) .
Revised: an example such that rank ( T ) > rank ( M ) .

Original: This completes the proof.
Revised: This completes the proof.

Original: E S UPPLEMENTARY M ATERIAL FOR E XPERIMENTS
Revised: E S UPPLEMENTARY M ATERIAL FOR E XPERIMENTS

Original: All the experiments are run on computing nodes with 4 NVIDIA GeForce GPUs.
Revised: All the experiments are run on computing nodes with 4 NVIDIA GeForce GPUs.

Original: (1) Language Modeling on Penn Treebank.
Revised: (1) Language Modeling on Penn Treebank.

Original: Language modeling is the task of computing the probability of a sentence or a sequence of words.
Revised: Language modeling is the task of computing the probability of a sentence or a sequence of words.

Original: The model performance is measured by per-word Perplexity (PPL) which is the lower the better
Revised: The model performance is measured by per-word Perplexity (PPL) which is the lower the better

Original: Penn Treebank (Marcus et al., 1993) is under the LDC User Agreement for Non-Members.
Revised: Penn Treebank (Marcus et al., 1993) is under the LDC User Agreement for Non-Members.

Original: We adopt the default settings as in the vanilla Transformer models(Vaswani et al., 2017).
Revised: We adopt the default settings as in the vanilla Transformer models(Vaswani et al., 2017).

Original: Speciﬁcally, we have
Revised: Speciﬁcally, we have

Original: 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.
Revised: 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.

Original: We use Adam optimizer and have 4000 warm-up steps
Revised: We use Adam optimizer and have 4000 warm-up steps

Original: We
Revised: We

Original: also use label smoothing(0.
Revised: also use label smoothing(0.

Original: The model is trained
Revised: The model is trained

Original: for 50 epochs.
Revised: for 50 epochs.

Original: (2) Neural Machine Translation on WMT16.
Revised: (2) Neural Machine Translation on WMT16.

Original: The goal of NMT is to generate a corresponding sequence in one language given a sequence in another.
Revised: The goal of NMT is to generate a corresponding sequence in one language given a sequence in another.

Original: The performance is measured by BLEU scores.
Revised: The performance is measured by BLEU scores.

Original: (the higher, the better).
Revised: (the higher, the better).

Original: WMT
Revised: WMT

Original: English-German dataset (Sennrich et al., 2016)
Revised: English-German dataset (Sennrich et al., 2016)

Original: is
Revised: is

Original: under the MIT license
Revised: under the MIT license

Original: In this experiment, we
Revised: In this experiment, we

Original: use
Revised: use

Original: a Transformer model with 8 layers and we have F = 512 and R =
Revised: a Transformer model with 8 layers and we have F = 512 and R =

Original: The learning
Revised: The learning

Original: rate is 1 e
Revised: rate is 1 e

Original: − 5 and
Revised: − 5 and

Original: warmups
Revised: warmups

Original: We also set label-smooting to be 0 . 1 .
Revised: We also set label-smooting to be 0 . 1 .

Original: To speed up the experiment,
Revised: To speed up the experiment,

Original: we
Revised: we

Original: use the mix-precision trick with an optimization
Revised: use the mix-precision trick with an optimization

Original: level being O1.
Revised: level being O1.

Original: (3) Image Generation on MNIST and CIFAR-10.
Revised: (3) Image Generation on MNIST and CIFAR-10.

Original: The image generation task is to predict an image pixel by pixel.
Revised: The image generation task is to predict an image pixel by pixel.

Original: The performance is evaluated by bits per dimension (BPD) .
Revised: The performance is evaluated by bits per dimension (BPD) .

Original: The image generation task is chosen to demonstrate, ﬁrstly, that Tuformers work beyond the language domain
Revised: The image generation task is chosen to demonstrate, ﬁrstly, that Tuformers work beyond the language domain

Original: and secondly, Tuformers can improve other efﬁcient designs.
Revised: and secondly, Tuformers can improve other efﬁcient designs.

Original: CIFAR
Revised: CIFAR

Original: (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.
Revised: (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.

Original: 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 efﬁcient designs.
Revised: 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 efﬁcient designs.

Original: Katharopoulos et al.
Revised: Katharopoulos et al.

Original: (2020) introduces a linear transformer that has a state-of-the-art complexity.
Revised: (2020) introduces a linear transformer that has a state-of-the-art complexity.

Original: As a result, we evaluate our design on the same experiments as they do in their paper.
Revised: As a result, we evaluate our design on the same experiments as they do in their paper.

Original: We adopt basically the same settings except for GPU resources.
Revised: We adopt basically the same settings except for GPU resources.

Original: Concretely, we use a 8 layer transformer model with 8 heads.
Revised: Concretely, we use a 8 layer transformer model with 8 heads.

Original: The embedding size F is 256 an D is 32.
Revised: The embedding size F is 256 an D is 32.

Original: N is set to be 1024.
Revised: N is set to be 1024.

Original: We also use “RAdam” (Ma & Yarats, 2021) optimizer for a state-of-the-art result.
Revised: We also use “RAdam” (Ma & Yarats, 2021) optimizer for a state-of-the-art result.

Original: The model is trained for 250 epochs for MNIST dataset and 50 epochs for CIFAR10 dataset.
Revised: The model is trained for 250 epochs for MNIST dataset and 50 epochs for CIFAR10 dataset.

Original: The batch size is 10 in MNIST training and 4 in CIFAR training.
Revised: The batch size is 10 in MNIST training and 4 in CIFAR training.

Original: We additionally use he mix-precision trick with an optimization level set to be O1 to speed up the training.
Revised: We additionally use he mix-precision trick with an optimization level set to be O1 to speed up the training.

Original: (4) Automatic Speech Recognition on LibriSpeech Dataset ASR consists of transcribing audio speech segments into text.
Revised: (4) Automatic Speech Recognition on LibriSpeech Dataset ASR consists of transcribing audio speech segments into text.

Original: The performance is measured by Word Error Rate (WER) .
Revised: The performance is measured by Word Error Rate (WER) .

Original: (the lower, the better).
Revised: (the lower, the better).

Original: We adopt the default settings as in Fairseq (Ott et al., 2019), except that we replace MHSA with THSA.
Revised: We adopt the default settings as in Fairseq (Ott et al., 2019), except that we replace MHSA with THSA.

Original: (5) Natural Language Inference on MNLI and QNLI.
Revised: (5) Natural Language Inference on MNLI and QNLI.

Original: 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.
Revised: 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.

Original: The results are evaluated by accuracy.
Revised: The results are evaluated by accuracy.

Original: 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.
Revised: 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.

Original: We use a batch size of 32 for 10 epochs, with a learning rate 1 e − 5 .
Revised: We use a batch size of 32 for 10 epochs, with a learning rate 1 e − 5 .

Original: Memory Overhead in THSA.
Revised: Memory Overhead in THSA.

Original: As mentioned in Section 4.3, naive Tuformers come with an extra memory overhead when calculating values within the softmax.
Revised: As mentioned in Section 4.3, naive Tuformers come with an extra memory overhead when calculating values within the softmax.

Original: 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.
Revised: 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.

Original: We argue that combining Tuformers with kernel-based efﬁcient variants such as Linear
Revised: We argue that combining Tuformers with kernel-based efﬁcient variants such as Linear

Original: Transformers (Katharopoulos et al., 2020) and polynomial kernels (Tsai et al., 2019) kernels could not only cancel the memory overhead but also achieve the same linear memory complexity by removing the nonlinear constrain in the softmax.
Revised: Transformers (Katharopoulos et al., 2020) and polynomial kernels (Tsai et al., 2019) kernels could not only cancel the memory overhead but also achieve the same linear memory complexity by removing the nonlinear constrain in the softmax.

Original: Furthermore, the extra memory overhead can be eliminated from the engineering side using the check-pointing technique: intermediate values are not saved in the memory but recalculated when needed.
Revised: Furthermore, the extra memory overhead can be eliminated from the engineering side using the check-pointing technique: intermediate values are not saved in the memory but recalculated when needed.

Original: We compare the runtime in both training and inference phase with check-pointing.
Revised: We compare the runtime in both training and inference phase with check-pointing.

Original: As shown in ﬁg. 12, under the same number of parameters, there is only slight increase in the training and basically the same runtime in the inference phase.
Revised: As shown in ﬁg. 12, under the same number of parameters, there is only slight increase in the training and basically the same runtime in the inference phase.

Original: With a direct transformation from Transformer to Tuformer (for example, when initialized with pre-trained models), the runtime is still under a reasonable range.
Revised: With a direct transformation from Transformer to Tuformer (for example, when initialized with pre-trained models), the runtime is still under a reasonable range.

Original: [Figure]
Revised: [Figure]