Tuformer: Data-driven Design of Transformers for Improved Generalization or Efficiency
Authors: Xiaoyu Liu
Jiahao Su
Furong Huang
Venue: ICLR 2022 Poster (2022)
Keywords: Attention Modules
Transformers
Data-driven Model Design
Trainable Heads
Expressive Power
Tensor Methods.

Abstract: Transformers are neural network architectures that achieve remarkable performance in many areas. However, the core component of Transformers, multi-head self-attention (MHSA), is mainly derived from heuristics, and the interactions across its components are not well understood. To address the problem, we first introduce a mathematically rigorous and yet intuitive tensor diagram representation of MHSA. Guided by tensor diagram representations, we propose a novel design, namely Tunable Transformers (Tuformers), by allowing data-driven weights across heads, whereas MHSA adopts pre-defined and fixed weights across heads, as will be explained in our paper. Tuformers naturally reveal a flexible design space that a user, depending on the needs, can choose a structure that has either improved performance (generalization error) or higher model efficiency. Any pre-trained Transformer can be an initialization of the corresponding Tuformer with trainable number of heads for efficient training and fine-tuning. Tuformers universally outperform Transformers on various tasks across multiple domains under a wide range of model sizes.

Published as a conference paper at ICLR 2022
TUFORMER: DATA-DRIVEN DESIGN OF TRANSFORM- ERS FOR IMPROVED GENERALIZATION OR EFFICIENCY
Xiaoyu Liu
Jiahao Su
Department of Computer Science
Department of Computer Science
University of Maryland, College Park
University of Maryland, College Park
MD 20740, USA
MD 20740, USA
x1iu1231e umd. edu
j jiahaosu @udd edu
Furong Huang
Department of Computer Science University of Maryland, College Park
MD 20740, USA Furongh@umd.est
ABSTRACT
Transformers are neural network architectures that achieve remarkable performance in many areas. However, the core component of Transformers, multi-head self- attention (MHSA), is mainly derived from heuristics, and the interactions across its components not well understood. To address the problem, we first introduce a mathematically rigorous and yet intuitive tensor diagram representation of MHSA. Guided by tensor diagram representations, Tunable Transformers (Tuformers), by allowing data-driven weights across heads, whereas MHSA adopts pre-defined and fixed weights across heads, as will be explained in our paper. Tuformers naturally reveal flexible design space that a user, depending on the needs, can choose a structure that has either improved performance (generalization error) or higher model efficiency. Any pre-trained Transformer can an initialization of the corresponding Tuformer with trainable number of heads for efficient training and fine-tuning. Tuformers universally outperform Transformers on various tasks across multiple domains under a wide range of model sizes
propose novel design, namely
INTRODUCTION
Transformer models are first introduced by Vaswani et al. (017)) in the context of sequence modelling. They have demonstrated impressive results language model pre- pre-traing training (Sanh et al.) 2019), speech recognition (Moritz et al.) 2220)), image classification (Dosovitskiy et al. 2000), and generation (Parmar et al.) 2018). The core component of Transformer models is the mulli-head self-attention (MHSA) which is extended from the standard attention mechanism (Bahdanau et al. 2014). Each attention head in MHSA has a global receptive field, i.c.. each token'ss representation updated by attending to all other tokens, and H attention heads are computed in parallel and concatenated together.
wide variety of tasks in many fields, such as
The current MHSA design is mainly derived from empirical studies or heuristics, leaving some unresolved challenges.
(1) Lack of solid theoretical understanding. The theory behind MHSA is only starting to catch up with practice. The role of the components and their interactions in MHSA are not well understood rigorously from theoretical perspective, which may lead to inefficiencies in the design. For example, Miicell et al. (2019) show that most attention heads can be removed in the testing phase without much performance compromise, while Cordonnier et al. (2020) find empirical evidence of redundancy in key/query projections and propose parammtterization scheme. These works focus more on practical solutions, leading to questions about whether theoretical patterns exist in these designs. (2) The number of heads is not trainable. Intuitively, the heads in multi-heads
code available at https//ggihuu.com/mm--uuu lab/ufoomrr

---

Published as a conference paper at ICLR 2022
expected to capture different context information through each head. However, the number of heads is fixed in training. Thus, although we could tune the hyper-panee parameter, an exhaustive search would be time-consuming or not practical for large-scale tasks. (3) Hard to analyze the expressive power, Analyzing the expressive power (as will be defined in Definition 33 of a neural network, i.e., proving that some architectures are more expressive than others, is non-trivial and challenging task. Such analysis has been done in Convolutional Neural Networks (CNNs) (LeCun et al., 1995 Cohen
ct al. 2016), Recurrent Neural Networks (RNNs) (Mikolov et al., 2010 Khrulkov al. 2018). but
no such work exists for interpreting MHSA or guiding the structural design.
In response to the above challenges, we first interpret MHSA from a tensor representation perspective using the intuitive graphical tool. the tensor diagram. Current prevalent descriptions of MHSA use flow charts to convey high-level intuitions, which could cause ambiguities. Therefore, is inevitable to pair those flow charts with mathematical formulas to understand the mechanism precisely. However, these two separated descriptions create difficulties for interpretations and inspections of the operations implemented in MHSA. To address this issue, we propose a graphical representation of MHSA, which is both semantically intuitive and mathematically rigorous. Specifically, we modify and extend the vanilla tensor diagrams (Penrose 1971), which conveniently allow for rigorous graphical representation of multi-linear operations between multi-ddimensonnl arrays (i.e., higher-order tensors) to represent nonlinear operations.
Using tensor diagram representation, we project the current design of MHSA into its tensor form, which renders holistic view of the weights in MHSA for better interpretation of the information flow and exchange among the components of MHSA. We then propose a novel data-driven structure, namely Tunable-Heed Self-Artention (THSA), which a re-parameterization of the weight matrices to allow learnable weight across heads. Transformers with THSA, named Twnable Transformers (Tuformers). have several advantages compared against vanilla Transformers: expressive power. We prove that MHSA is special case of THSA. (2) The number of heads is trainable. The concept of the number of heads THSA generalizes to the stable rank of the core matrix, allowing data-driven implicit training.
(1) A guaranteed higher
(3) Tuformers allow initialization from pre-trained Transformers such as BERT Devlin et al. 2019) and its variants.
More importantly, formulate flexible design space for MHSA where explore the trade-off between generalization and efficiency, which can be adjusted by the expressive power if there is no over- ever-fitting. Users can choose structures with either higher model efficiency or better generalization that satisfy their own needs. The proposed design space further allows improving the expressive power of the network by finding better tensor representations of the weights through tensor representation theory in future followup works.
We experiment Tuformers with several tasks across multiple domains, from language modeling, chine translation to image generation under a wide range of model sizes. We demonstrate competitive results not only on Tuformers but also in cases where Tuformers are initialized with pre-trained Transformers for other downstream tasks, when combined with Linear Transformer (Katharopoulos al. 2020) on the image generation task and their efficient variants in the design space.
Summary of Contributions:
(1) We propose a mathematically rigorous and semantically intuitive tensor diagram representation of the multi-head self-attention introducing a new tool to the ML community for future studies on interpretation and improvements of Transformers.
(2) We propose novel design of the MHSA, Tunable-Heea Self-Attention (THSA), resulting in Tuformers. Tuformers are structurally more expressive and are showing improved generalization error experimentally.
(3) We formulate a flexible design space for attention unit design where users can choose to design structures, depending the needs, that have better generalization efficiency.
2 TENSOR DIAGRAM REPRESENTATION REPRESENTATION
We use lower case letters (e.g., v) to denote vectors, upper case letters (e.g., M) to denote matrices, and curly letters (e.g.. T) to denote general tensors. For a tensor T E Rh x Ta we refer to the number of indices as order, each individual index mode, and the length of one mode as dimension. For instance, T is an Mh order tensor that has dimension Im at its th mode. We reserve superscripts to distinguish similar arrays (e.g., Wº, WK V are query/key/value weight
Notations.
2

---

Published as a conference paper at ICLR 2022
matrices), and subseripts to index the elements in an array (e.g., Wij is the (i,5)th element of W). We use colon to slice array (e.g.. W.: denotes the 7th row of W).
We propose to use tensor diagrams (Penrose 1971]. a commonly used rigorous/precise and intuitive graphical representation for multi-linear operations among higher-order arrays (i.e., tensors), to represent multi-head self-attention (MHSA). Since MHSA consists of multilinear operations and a nonlinear softmax function, will introduce tensor diagrams and our novel design extending tensor diagrams to denote the composition of nonlinear and multilinear operations in MHSA.
2.1 TENSOR DIAGRAM BASICS
Arrays denoted nodes with legs. An array is represented as a node with leg(s) in tensor diagram shown Figure [I We denote the order (the number of dimensions) of the array by the number of legs extending from the node. Each labeled leg represents Every mode the tensor needs to be uniquely labeled. We usually use the dimension of the mode as the label (i.e., an associated positive integer written
A
B
M
Vector order tensor order tensor
Matrix
3 dimensional array order tensor
mode of a tensor.
Figure 1: Arrays in tensor diagram. A vector is de-
noted as a node with I leg, matrix as a node with legs and an N-dimensional array as node with N legs.
top of each leg). The legs
do not need to be straight lines, and their orientations do not matter. Matrices M E RAx and M BxA can be represented via the same tensor diagram long 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).
B N
M C
A
N
A
B
K Ami Binn Cmn
cxpla Amn)
A Amr Bur Camr
Cmn
En'= exp(a Amn')
1=1
Figure 2: Tensor diagrams for atomic operations. (1) (left) Contraction is operation that generalizes the matrix multiplication. It sums element-wise products on mode object A and corresponding mode (with the same dimension) in object B (i.e., along a leg in node A and a corresponding leg in node B). In the tensor diagram, multiplying two matrices (or higher-order tensors with more than 2 legs) corresponds to "gluing" their corresponding legs (on certain mode). (2) (middle) Softmax is element- wise exponential function normalized along certain mode. We propose denote ax-caaled softmax function softmax((a on A as dotted box with labeled filled ball (to distinguish itself from tensor objects, i.c., nodes which are blank circles) attached to one leg. (3) (right) Batch multiplication is an element-wise product along the connected legs.
Operations in tensor diagrams. There three types of operations in the calculation of MHSA:, contraction, softmax, softtma,,
and batch multiplication, as shown in Figure 2 and explained in its caption.
Evaluation of tensor diagrams: (1) Evaluation order. We can evaluate a tensor diagram in any pair-wise order except for the nodes in the softmax box. Since the softmax function is nonlinear, we must first evaluate the nodes in the softmax box with arbitrary order before those outside the box. (2) Reading a tensor diagram. We can easily identify the output shape by the dangling edges. For instant, in Section 2.1 there are three dangling legs M, N, R. Thus the output is 3rd order tensor with dimensions M, N. R. Note that the softmax function does not change the shape of the tensor.
Advantages of tensor diagrams: (1) Tensor diagram orientation invariant, meaning that we can represent X and X using the same tensor diagram. Thus in multi-head, we obtain a universal graphical representation regardless of whether we represent an embedding of each token in the sequence as rows or columns of the input embedding matrix X. (2) Multi-linear operations are concisely represented and interpreted. (3) The representation is both precise in math and intuitive, and the information flow is clear, making analysis of network structure more accessible. In addition, with the labels of the legs, we can read the model complexity explicitly.
We include comprehensive introduction to tensor diagram representation in Appendix A
3

---

Published as a conference paper at ICLR 2022
2.2 TENSOR DIAGRAM REPRESENTATION OF MULTI-HEAD
The core of Transformer models is a multi-head self-altention (MHSA) module which allows the model to jointly attend information from different representation sub-spaces at different posi- tions Vaswani et al. 2017). To distinguish inputs for query, key, value matrices, we use X E Rx X* X
_XXx respectively. The MHSA module outputs matrix M RNXF as
Q
X
K
[M] = X W
VIN = X W [h]
=X
(la)
[A]
[A]
head [M] softmax (M [N]
IN/
D//VDD Vio
(Ib)
M [head []}, head [2]: ... head H]] W
(1c)
In the above equations, we see how MHSA computes its output in three steps:
la Linear transformation on input embedding to get latent features. For each h. the query/key/value
matrices [A)). K(M-- VIN] are linear transformations, parameterized by w W. and W [h] [h]
[4]
R FxD respectively, of the input embedding.
Ib
Scaled dot-product attention. Each head [k] E R NxD computes a scaled dot- dot-product be-
Tween a latent feature matrix Qix: E NRD (i.e., query matrix) and another latent feature matrix
K [N] E RRMDD (i.e., a key matrix) along D, and applies softmax function to obtain the weights to be multiplied with third latent feature matrix ViN E R RMxP (i.e., value matrix).
(Tc) Concatenation and contraction for MHSA. After concatenation of the H heads, a linear trans- formation of the concatenated heads, parameterized by W E RDD is implemented get the outcome M of the MHSA as a RR X object.
Latent Contractio
X
WE
X
IN
X
7F
X
V
Head Contraction
M
T
W
X
Wx
X
D
(a) Single- head Self-Attention Multi-head self-attention in tensor diagram. The concatenation of H matrices [head /= can be viewed as constructing a three dimensional array with the third mode being the number of heads, i.e. three-legged nodes in tensor diagram notation. This way, the weight matrices {wo) Q !
(b) Multi-
]
[N] Jh=1 are
stacked to be an order 3 tensor W It is the same with other weight matrices. Since in MHSA, the attention calculation among key, query and value matrices are done separately within each head, batch multiplication along mode H is proper way to denote these "head-wise" operations. Finally, we connect the resultant tensor with y to obtain MHSA as in Figure 3b See Appendix details and the proof.
B for more
TUFORMERS: DATA-DRIVEN DESIGN WITH HIGHER EXPRESSIVE POWER
3.1 MOTIVATION: INTERPRETATION MHSA
In MHSA, ( and K [n] can only contribute to head Yh. In order to characterize the information flow in MHSA, we stack the weight matrices W-Q Vh together (similarly for W
[A] and WV
(M] and
[k] introduce a matrix C' to interpret MHSA. Guided by its tensor diagram representation in Figure 3b
C is clearly characterized as in Proposition
Proposition (Interpretation of MHSA). With weight matrices stacked as W @
(H) E BFDDH and care matrix C :- In @ (1p15) E DHDD (iee, Kronecker product of all-ones matrix 1 RDDx and identity matrix I HR ), MHSA is equivalent to
[ pp [ W = [W
uip W. W = -W WW
[
4

---

Published as a conference paper at ICLR 2022
Q= x K ==x= V= x'w",
(2a)
DH
head, = sofimax E Q,K. /VD V,,
(2b)
M = head,, head2,
head DM WO.
(2c)
In tensor diagram notation, MHSA is equivalent to Figure 4a Here, we use K, R head, E TR.N to denote the p column of the corresponding matrices, and Crs to denote the (r, 8)th element of
Proposition is proved in Appendix D.1 Note that (head], .... he correspond to head [I]-
(head D+1 D+1. head2D} correspond to head and so forth. The intuition of interpreting MHSA using Proposition comes from the tensor diagram graphical representation of MHSA structure. From the tensor diagram of MHSA shown in Figure we two types of contractions: latent contractions along D and head contractions along H. The former corresponds to (1n1p) part of C and the latter corresponds to In part of C' as shown in Figure 4a
Motivation for an improved design: from MHSA (a preset C) to THSA (a learnable C). In MHSA, with the preset specific C = In (Inlb), restricts the possible weight sharing. In other words, Qi, and K(N] can only contribute to head [A], Vh. However, there is no motivation why such a preset and fixed C' is desirable or why we do not allow Qı, or K to contribute other heads. Our motivation for an improved design is simple: rather than presetting the core matrix as C C= @ (1n1p), we allow data-driven learnable C, and we call such a design Tunable- head self-atiention (THSA). Now with the learnable C in THSA, we allow Q [A] and K to contribute to other heads head wr], where h f h. In other words, THSA allows learnable weights across heads.
X
WP
G
W
X
(a) Multi- head Self-Attention (MHSA)
(b) Tunable- Head Self Belf-Attention
(c) C as design space: example.
(THSA)
C= & (I]ID)
C' fully trainable
C
C₁ @ (C. Cs).
Figure MHSA as special case THSA. Figure (left) represents MHSA. Figure 4b (middle) represents
THSA. As shown in Proposition [
weight matrices in MHSA collaborate in a very specific way. When
C= In @ (1)lD) in Proposition THSA degrades to to MHSA. Therefore, it is natural to make C' fully trainable for higher expressive power. THSA also reveals design space. For example, can extend MHSA by allowing trainable components C1, C2 and Cs as in Figure 4e
3.2 A NOVEL DESIGN: TUNABLE- TUNABLE-HEAD SELF-ATTENTION (THSA)
As shown in Figure 4b 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 E Nx F and X K X X as inputs and returns matrix T E NxF output. A THSA, with five weight matrices Wº, WK WV FxR wo R and E RRXR is mathematically described as
Q = X
X"W", XVW, R Ω = softmax C Q.K. / V,,
(3a)
(3b)
T = ΩW S2
aal W
(3c)
The intermediate result 22 R is the th column of the matrix Z. In MHSA, the softmax function is scaled, which normalizes its input by VD, i.e., the square root of the latent dimension. The scaling
"THSA improves MHSA and only modifies the intra-layer structure. It does not change the inter-layer structure outside MHSA.
5

---

Published as a conference paper at ICLR 2022
ensures that the distribution of the attention matrix approximates standard Gaussian distribution. To achieve similar effect, in THSA, we change the constant to vR accordingly.
THSA generalizes MHSA. THSA has learnable C. IfC = I (1p15), i... a Kronecker product of an all-ones matrix 1 E RDDD and an identity matrix Iи E RHxH the THSA reduces to an MHSA with H heads and latent dimension D, as illustrated in Figure 4a We call a Transformer with THSA module as Tunable Transformers (Tuformers). The tensor diagram representation of THSA is given in Figure 4b It has learnable core matrix C' compared to MHSA in Figure 4a which has preset core matrix.
Tuformers can be initialized with Pre-trained Transformers. Although different in the structure, Tuformers can always be initialized with pre-trained Transformer models through re- parameterization. V 1 [A] 1=1- (W Jh=1- [b] h=1 and W while in THSA MHSA is updated by weight matrices (W K
we have Wº, W WV, Wo and C. initialize THSA with MHSA, we follow the scheme proposed in Proposition [ C is the kronecker product between D x D all-one matrix and an H X H identity matrix, where is the latent dimension and H is the number of heads MHSA. 0o remains the same as that of MHSA. As for W W and W e FXX, we concatenate [ 1 (WQ 1 [b] h= (W
(W k
V. H M1=1 E Fx along mode D.
The notion of heads in THSA. Due to the learnable C' and data-dependent weight sharing mecha- nism, the notion of heads in THSA not as obvious as in MHSA. Although we have an upper bound R of the number of heads, which is achieved when the learned C' decides not to share any weights across heads, the actual effective number of heads depends on C. As result, we define the number of heads in THSA the stable rank of the core matrix C. which is Σ; of/ max; ds, where o is the singular value of C. We use the stable rank because it is largely unaffected by tiny singular values.
THSA flexible design space: C in THSA Figure 4b
naturally reveals design space. In MHSA, C' is completely fixed and in THSA C is fully trainable. We can also have flexible C' designs for user-specific needs of either increasing generalization or improving efficiency.
(a) skip-left C In @ (CFCs) C.)
(b) skip-right C₁ @ (Iblp)
(c) trainable heads-only
(d) heads-only
C C₁
1s
Figure 5: C naturally reveals design space. We can obtain better generalization or hetter efficiency by allowing fully trainable C. partially trainable C. pre-set fixed C' and etc. We demonstrate some of them extending from the structure of MHSA where C' takes the form of Kronecker product as shown in Figure 4c The relationship
between the generalization and efficiency of these designs is summarized in Figure 6
Better generalization. Better generalization can be obtained by increasing the expressive power if there is no over-fitting 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. We find that a more effective way to increase the expressive power for better gen- eralization to find more expressive structures. We demonstrate a few of them in Figure 5]and Figure 4b In MHSA, C is pre-set fixed where C
More Gereratzation
THSA
skip-left
skip-right
MHSA
Heads-only
Heads-only
More Efficency
In & (1
Figure 6: We can divide all structures listed in Figure 4 and Figure into different groups by their expressive power. The higher group structure is in, the higher expressive it has. Better generalization can be obtained in more expressive structures. We can also go for very efficient ones if we are willing to sacrifice some generalization. More details in Section 4
In Figure 5a (skip-left C In Φ (CJ C²) ) and Fig- ure (skipprigth C IH @ (C) Cz) ), we maintain the form a Kronecker product but allow either C, or C2. Ca trainable. We can think of skip-left and skip- right as skip-connection designs in C. In THSA, C is fully trainable. It is easy to observe
increasing in expressive power from MHSA to leip-left, skip-right, and finally THSA, together with an increasing in generalization in the same order. Our empirical study

---

Published as a conference paper at ICLR 2022
also show that these structurally more expressive models more effectively improve the generalization than naively increase the number of parameters.
More Efficiency. We can obtain different levels of efficient structures in the design space by adjusting the expressive power. Compared to MHSA, skip-left and skip-right have more generalization in various tasks but almost the same number of parameters. Being special cases of skip-right, Figure 5c ( trainable heads-only C = C₁ and Figure [d] heads-only C = TH ) are less expressive. However, trainable heads-only only has around 6% of the parameters compared MHSA while heads-only only has 1.5% with most performance maintained.
There is trade-off between generalization and efficiency in the design space C demonstrated in Figure 6 We empirically verify that without over-fitting generalization and efficiency have similar hierarchy as expressive power. Thus we can design structures different level of expressive power, tailoring needs for either better generalization or more efficiency. THSA has the best generalization, but 125% parameters compared to MHSA. Since skip-left and skip-right are special cases of THSA, it is natural that their expressive power and generalization are lower. However, they are more efficient. With almost the same number of parameters, they have better generalization than MHSA. Heads-only and trainable heads-only are very efficiency sacrificing an acceptable amount generalization. More details can be found in Section
4 EXPERIMENTS
Datasets, tasks and evaluation metrics. We evaluate Tuformers on 7 datasets for 5 tasks: word-level Language Modeling (LM) on Penn Treebank (PTB) (Mareus 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 QNLL(Wang et al. (018) and Image Generation on CIFAR10 (Krizhevsky 2009)
and MNIST Deng 2012) datasets.
Baselines. We compare the performance of MHSA module with THSA module in the following 5
backbone architectures: (1) a vanilla Transformer model (Vaswani et al. 2017). (2) an Transformer-
XL model (Dai et al., 2019). (3) BERT-large model (Devlin et al. 2019). (4) RoBERTa-
model (Liu et al. 2019). and (5) an ALBERT model (Lan et al.) 2019p. In all baselines, we adopt all default settings used in the original paper (hyperparameters specified in Appendix E). and only
aataacc/re-prrrmeterii the MHSA to THSA.
4.1 EMPIRICAL RESULTS OF TUFORMERS.
In this part, we first show that THSA (used in Tuformers) universally outperforms MHSA (used in Transformers) diverse tasks under different model scales and THSA is simple to use in practice.
(I) THSA consistently oulperforms MHSA on diverse datasets for a variety of tasks under different model scales. In LM, NMT and ASR tasks, we compare THSA against the MHSA models under different model scales. MHSA is parameterized by four weight matrices Wº, WK WV and Wo whereas THSA adds an additional core matrix C of the same size. Adopting the size of the four original weight matrices in MHSA to THSA incurs an 25% increase in the number of parameters. For fair comparison, we first compare MHSA against THSA with the same number of parameters (thus with smaller query//kyyvallueouupu weight matrices). To get a full picture, we also compare MHSA against THSA with the same sized four original weight matrices. As shown Figure 7 under the same number of parameters, curve) under all model scales. By increasing the number of parameters of the attention units by 25% (i.e., adding the core matrix C while keeping the size of the query/key/valne/ottuu weight matrices the same), our THSA (purple curve) further improves the performance universally under all model scales.
THSA (blue curve) universally outperforms MHSA (red
(2) Tuformers initialized with pre-trained strong models fodelsfuurther improve the performance.
In NLJI task, initialize Tuformers with pre-trained BERT-argee, RoBERTa-large and ALBERT models, and perform the same fine-tuning tasks. As shown in Table [ THSA is compatible with
As will be discussed later, the overall model size is not drastically increased.
7

---

Published as a conference paper at ICLR 2022
Larguage Modellnng
Machire Tranchtion WMF-15 Reggle--Grmenn
Jatoaaii Speeca Recogaiikas on Lesigeede Class
MIHSA (sanee oun poran)
MHSA (some
T (25% mere paras)
THS (same num of poran)
BLEU shrelight the bette))
20
THSA
THSA
CO (te Lwwce barker
12-
MRSSA (aae eam pancasee))
THSA (25% mase paransterr
TISA Gare came poranreees
tes (de Jover
56
28
64-
1054
20%
50%
100%
150
20066
10%4
20%
505
100%
150%
2009
10%
20%
50%
190%
200%
meddel scale
moddd ande
Figure 7: Performance comparison of THSA and MHSA in (left) language modeling on Penn Treebank dataset on TransformeXLL, (middle) neural machine translation on WMT16 English-German German dataset and (right) automatic speech recognition on Librispeech on Transformers. Under the same number of parameters, our THSA universally outperforms MHSA under all model sizes. The error bars displayed in shades are generated using 10 random runs in LM and NMT and random runs in ASR.
Table 1: (Left) THSA initialized with pre-trained models. Comparison of accuracies obtained for language inference task on MNLI and QNLI datasets, the higher the better. Results show that Tuformers can be conveniently initialized with pre-trained models for downstream tasks, and fine-tuning obtains impressive performance. (Right) THSA +kernels. Comparison of bits per dimension (BPD) for image generation task on MNIST and CIFAR-10 datasets, the lower the better. Results show that Tuformers be extended to image generation tasks and improve the performance of other efficient designs with linear computational and memory complexities.
Models
MNLI
QNLI
Models
MNIST
CIFAR 10
BERT (Devin et 2019
84.1%
92.1%
BERT + THSA
84.9%
92.8%
ELU kernel Katharopoolos et al. 2020
U.72
3.51
RoBERTa Liu al. 2019
89.4%
94.3%
ELU kemel + THSA
0.65
3.42
RoBERTa THSA
90.2%
94.7%
Polynomial kemel (Tai al. 2019
0.64
3.47
ALBERT [Lan al 2219
89.1%
97.9%
Polynomial kernel THSA
0.57
3.37
ALBERT THSA
89.5%
98.1% 1
existing pre-trained Transformer models and there is no need to train Tuformers from scratch: fine- tuning THSA initialized from pre-trained BERT and its variants further improves their performance.
(3) Combing THSA with kernels obtain linear computation and memory complexities in sequence length. Tuformers can also be incorporated into efficient variants of Transformers for lower computa- tion and memory complexities. We experiment on image generation tasks on CIFARIO (Krizhevsky) 2009) and MNIST (Deng) 2012) datasets as in Table | and show that Tuformers obtain strong performance and linear space complexity when combined with kernels.
4.2 FLEXIBLE DESIGN SPACE
(1) THSA reveals a flexible design space for Transformer. We empirically demonstrate that there exists trade-off between generalization and efficiency in this design space, which could be adjusted by the expressive power. By making C' aalia-train, partially trainable or pre-set fixed as shown in Figure 5 Figure 4 we demonstrate that some structures in C improves the generalization while some other designs improve the efficiency of MHSA.
(2) Models with higher expressive power in the de- sign space have better generalization. As shown in Table and Figure we compare two ways that in- crease the expressive power for better generalization: naively increasing the number of parameters and de- signing more expressive structures. Both approaches increase generalization. However, naively increas- ing the number of parameters is not as effective as adopting more expressive structures. With the same number of parameters, more expressive structures, such as THSA, lkip-left and skip-right, universally outperform MHSA.
Table 2: Performance in the flexible design space.
Models
NMT (BLEU)
parums
Train
Test
MHSA
28.9
100%
10.6
3.2
MIRR with 125% peams
29.0
125%
11.2
3.6
MISA with 1.5% prams
20.2
1.5%
6.4
2.5
THSA skip-ight C. C: a (1b)t)
29.4
125%
11.5
3.9
20.3
1018
10.6
3,2
skip left C- In (C))
29.2
100%
10.6
3.2
inánnille heads only C
27.4
6%
7.4
28
head--oll C In
27.1
1.5%
6.1
23
8

---

for more Published as a conference paper at ICLR 2022
(3) Models can also be made very efficient. Moreover, in Figure trainable heads-only C= we trade generalization for efficiency. With 6% parameters, we maintain 95% performance in the NMT task. In extreme case, if only maintain the head contraction as shown in Figure 5d the number of parameters goes down to 1.5% with a 42% faster training and 30% faster testing. Surprisingly, almost 94% performance is obtained with this light-weighted model.
4.3 ABLATION STUDY: PHECK-POINTIIG TO ALLEVIATE MEMORY OVERHEAD WITHOUT KERNELS.
Although an effective design with guaranteed higher expressive power, Tuformers come with an extra memory overhead when calculating values within the softmax. In Tuformers, if done naively, the memory overhead is MN R where N, . denote the sequence length and R = DH is the rank the
core tensor, incurring a D times larger memory overhead than vanilla Transformers in calculating softmax. Combining Tuformers with some tfat-off-hheart kernel-based efficient Transformers such Linear Transformers using ELU (Katharopoulos et al. 2020) and polynomial (T'i 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.
We can also alleviate the D times larger memory overhead compared with vanilla Transformers in calculating softmax without using SprnelsSSpecificalll we use Sheck-pointing to eliminate the memory overhead in THSA. Since no intermediate results are saved in the forward pass using check- pointing, we recalculate the intermediate results the softmax box the back prop-propagation which introduces some computational overhead. We observe that applying check-pointin to Tuformers leads to only a slight increase in training time while relieving the memory overhead. We can reduce or even avoid the repetitive calculations by developing a proper backward module, although such development beyond the scope of this paper and deferred for future works. See Figure 12 details.
5 RELATED WORKS
Analysis of MHSA. Given the popularity of Transformers in wide range of domains, a line of works focuses on understanding and improving MHSA. Voita et al. (2019) proposed a practical scheme to prune less informative heads. Michel et al. (2019) also showed that most attention heads
are removable in testing without much performance compromise. [Cordonnier et al. (0200) found empirical evidence of redundancy in key/query projections and proposed re-parameterization scheme. These works bring valuable insights into MHSA, but a consistent theoretical analysis for these findings is missing. Our work differs from these works in that we focus on the theoretical understanding of MHSA in terms of expressive power, and we provide theoretical guarantees for novel design, Tuformers, with higher expressive power and data-driven trainable heads.
Efficient Transformers. MHSA has a global receptive field, i.c., each token's representation is updated by attending to all other tokens, therefore incurring quadratic memory and computation complexities concerning the sequence length. An extensive line of works focusing on reliving such the dependencies the sequence length, such as Performer (Choromanski et al. 20220), Reformer (Kitaev et al. 2020). Linformer (Wang et al. 2020) and Linear Transformer (Katharopoulos et al.) 2020). Our proposed model, Tuformers, can be incorporated into some kernel-based efficient Transformers, which applies kernel-based methods to MHSA to remove the nonlinearity constraint brought by the softmax function, achieving the same state-of-the linear computation and memory complexities in the sequence length.
6 CONCLUSION
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, novel model design with a guaranteed higher expressive power. Furthermore, Tuformers have a data-driven structure where heads can be trained implicitly and initialized with pre-traine traitrd Transformer models. The introduction of the tensor diagram representations and the theory presented in this paper provide new tools and open new directions for future research searching for expressive and efficient architectures.

---

Published as a conference paper at ICLR 2022
REFERENCES
Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. arXiv preprint FrrXiv:4409000 2014.
Krzysztof Marcin Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Quincy Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Learning Representations. 2020.
Nadav Cohen, Or Sharir, and Amnon Shashua. On the expressive power of deep learning: A tensor analysis. In Conference on learning theory. pp. 698-728. PMLR, 2016.
Jean-Baptiste Cordonnier, Andreas Loukas, and Martin Jaggi. Multi-head attention: Collaborate
instead of concatenate. arXiv preprint ar0iv:2006. /6362, 2020.
Zihang Dai, Zhilin Yang, Yiming Yang. Jaime G Carbonell, Quoc Le. and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond fixed-length context. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 2978-29988 2019.
Li Deng. The mnist database of handwritten digit images for machine learning research. IEEE Signal Processing Magazine, 19(6):141118) 2012.
Jacob Devlin, Ming- Min Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chupter of the Association for Computational Linguistics: Human Language Technologies, Volume I (Long und Short Papers), pp. 4171-4186, 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations, 2020.
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference Machine Learning, pp. 5156-510 5165. PMLR, 2020.
Valentin Khrulkov, Alexander Novikov, and Ivan Oseledets. Expressive power of recurrent neural networks. In International Conference on Learning Representations, 2018.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Learning Representations, 2020.
Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
Zhenzhong Lan, Mingda Chen, Sebastian Goodman, Kevin Gimpel, Piyush Sharma, and Radu Soricut. Albert: A lite bert for self-supervised learning of language representations. arXiv preprint arXiv: 19422, 2019.
Yann LeCun, Yoshua Bengio, et al. Convolutional networks for images, speech, and time series. The handbook of brain theory and neural networks, (10):199)9995,, 1995.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann. wwwww...oo/eed/////tt../e/ 2, 2010.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach. arXiv preprint 190.v:1907.11660 16602,
Jerry Ma and Denis Yarats. On the adequacy of untuned warmup for adaptive optimization. In Proceedings of the AAAI Conference on Artificial Intelligence, 2021.
Mitchell P Marcus, Mary Ann Marcinkiewicz, and Beatrice Santorini. Building a large annotated corpus of english: the penn treebank. Computational Linguistics, 19(2)3313-330, 1993.
10

---

Published as a conference paper at ICLR 2022
Paul Michel, Omer Levy, and Graham Neubig. Are sixteen heads really better than one? Advances in Neural Information Processing Systems, 32: 32:14014-14024, 2019.
Tomas Mikolov, M. Karafiát, L. Burget, J. Cernocký, and S. Khudanpur. Recurrent neural network
based language model. In INTERSPEECH, 2010.
Niko Moritz, Takaaki Hori, and Jonathan Le. Streaming automatic speech recognition with the transformer model. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 0074-6078, IEEE, 2020.
Myle Ott, Sergey Edunov, Alexci Baevski, Angela Fan, Sam Gross, Nathan Ng. David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics (Demonstraaiioss), pp. 48 388-5,,22 2019.
Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur. Librispeech: an asr corpus based on public domain audio books. In 2015 IEEE international conference on acoustics, speech and signal processing (ICASSP), pp. 5206-50066666 5210. IEEE, 2015.
Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Noam Shazeer, Alexander Ku, and Dustin Tran. Image transformer. In International Conference Machine Learning, pp. 4055-4064.
PMLR, 2018.
Roger Penrose. Applications of negative dimensional tensors. Combinatorial mathematics and its applications, 1:221-244, 1971.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of bert: smaller, faster, cheaper and lighter. arXiv preprint 1rXiiv:910.00 000001108,,, 2019.
Rico Sennrich, Barry Haddow, and Alexandra Birch. Edinburgh neural machine translation systems for wmt 16. Proceedings of the First Conference on Machine Translation: Volume 2, Shared Task Papers, pp. 371-376, 2016.
Yao-Hung Hubert Tsai, Shaojie Bai, Makoto Yamada, Louis-Philippe Morency, and Ruslan Salakhut- dinov. Transformer dissection: An unified understanding for transformer's s attention via the lens of kernel. In Proceedings of the Conference on Empirical Methods in Natural Language Processing, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, pp. 5998-6008, 2017.
Elena Voita, David Talbot, Fedor Moiseev, Rico Sennrich, and Ivan Titov. Analyzing multi-head self-attention: Specialized heads do the heavy lifting, the rest can be pruned. In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, pp. 5797-5808. 2019.
Alex Wang, Amanpreet Singh, Julian Michael, Felix Hill, Omer Levy, and Samuel R Bowman. Glue: A multi-task benchmark and analysis platform for natural language understanding. arXiv preprint arXiv: 07461, 2018
Sinong Wang, Belinda Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 2020.
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technolo- gies, Volume 1 (Long Papers), pp. 11 1112-1122, Association for Computational Linguistics, 2018. URL http:://cc///c// org/ htw.://.....g...oooooooo///ll---
11

---

Published as a conference paper at ICLR 2022
A
SUPPLEMENTARY MATERIAL FOR TENSOR DIAGRAMS
In this section, we provide a comprehensive introduction to the tensor diagram representation.
A.1 COMPONENTS OF TENSOR DIAGRAM: ARRAYS AS NODES WITH LEG(S)
Arrays denoted as Nodes with Legs. An array is represented as a node with leg(s) in a tensor diagram shown in Figure T The order (the number of dimensions) of the array is denoted by the number of legs extending from the node. The legs do not need to be straight lines ,and the orientation does not matter. For example, matrices A and A' represents one mode of a tensor. Every mode of the tensor needs to be uniquely labeled. We usually use the dimension of the mode the label (i.e., associated positive integer written on top of each leg). If multiple modes of the same tensor have the same length, we use subscripts to differentiate them.
equivalent in this notation. Each labeled leg
A.2 OPERATIONS IN TENSOR DIAGRAM
A
M
C
A
N
M
A
8
M
N
R
exp(a Amn)
Aur Byr Cumr
Cmn
Amk B
Cun
En'=1 exp(a Amm'))
1=1
operation that generalizes the matrix multiplication. does summation on element-wise products along a mode in object and corresponding mode (with the same dimension) in object 2 (i.e., in lensor diagram language, along a leg in node and a correspording leg in node 2). 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). (2) Softmax as shown in Appendix AA2 is an element-wise exponential function normalized along a certain mode. Tensor diagram has convention to denote the contraction, but not the softmax. We propose to denote the a-scaled softmax function softmax( (a A) on A as a dotted box with labeled filled ball (to distinguish itself from tensor objects, i.e., nodes which are blank circles) attached to one leg. (3) Batch Multiplication, shown in appendix A.2 is an elementwise product along the connected legs.
Figure (1) Contraction as shown in Appendix A.2
There are three types of operationns.contruciionn softmax and batch multiplication, Contraction is operation that does summation on element-wise products along mode in object 1 and corresponding mode (with the same dimension) in object 2 (i.e., in tensor diagram language, along a leg in node and corresponding leg in node 2), whereas softmax is element-wise exponential function normalized along a certain mode. Tensor diagram has a convention to denote the contraction, but not the softmax. We will first introduce contraction in tensor diagram notation and then propose our design of softmax in tensor diagram language.
Contractions denoted as edges connecting the node legs. In tensor diagram, multiplying two matrices (or higher-order tensors with more than 2 legs) corresponds to "gluing" their corresponding legs (along certain mode) as shown in Appendix A.2 and it is called tensor contraction. Since the representation of arrays in the tensor diagram is orientation invariant, contractions are also orientation invariant. This is especially useful representing the self-attention unit. In some cases, the data object are represented as row vectors while in some other scenarios they are denoted as column vectors. If we use mathematical formula, we need take care of the ordering of the matrix multiplication because A B, AB" BBAT,, B" A and more are all different. In contrast, tensor diagram provides a universal graphical representation as long as the corresponding legs are connected correctly.
Nonlinear activation denoted as dotted box with a labelled filled ball attached. We propose to denote the a-scaled softmax function So)t on A dotted box with a labeled filled ball (to distinguish itself from tensor objects, i.c., nodes which are blank circles) attached to one leg. The label right above the filled ball denote the scaling parameter a, and the leg that the ball is attached to
12

---

Published as a conference paper at ICLR 2022
indicate the mode where normalization is implemented along in the softmax operation. As shown in Appendix A.2 let A denote the resultant of some operations (for example Q K [k)) the - -scaled vD softmax on A which normalize along leg N is depicted. Softmax does not change the shape of the input. That why we introduce a dotted box, which does not change the shape either, to denote it.
Batch Multiplication Given A and B, the batch multiplication is an element-wiss product along the connected legs. The resultant tensor dangling legs M, N and R. Since in MHSA, the attention calculation among key, query and value matrices are done separately within each head, batch multiplication a proper way to denote these "head-wise" operations.
obtained by merging connected nodes, maintaining the
A.3 ADVANTAGES TENSOR DIAGRAM
Tensor diagram representations enjoy few advantages compared with using flow charts together with math formulas.
(1) The notation is orientation invariant. The input embedding matrix X can be represented as either IR XFF or P with rows or columns being the embedding of each token in the sequence respectively. Although we take the former convention in our mathematical notations in Equation B.I there are papers that use the latter, creating discrepancies between notations. However, in tensor diagram, since the node legs are orientation invariant, obtain universal graphical representation of mead self-attention, irrespective of how the input embedding matrix represented, as shown in Figure
(2) The representation is both precise math and intuitive, making analysis easier. The tensor diagram itself is mathematically rigorous.
With the labels of the legs, the model complexity (number of parameters) is explicitly displayed, requiring supplementary information such as the size of the parameters W W and WV as in mathematical formulas. In addition,the resultant from an operation or a sequence of operations (i.e., connected sub-graph of the entire tensor diagram graph) can be treated as "merging" any connected nodes in the sub-graph. This is particularly convenient to obtain the size of the resultant: any the dangling leg becomes a leg of the resultant.
(3) Multi-linear operations are concisely represented and interpreted. Using tensor diagram, we concisely represent the multi-head self-attention rigorously in diagram as shown in Figure 9c More important, the multi-linear interaction between the latent features Q Qp- Kinj and Via. which essentially attribute to the multi-linear interaction between the weight tensors Wº, W and B clearly illustrated in the tensor diagram. This illustration, which is not achieved by any existing representations, is crucial for our principled understanding of the multi-head self-attention.
(4) The information flow clear. For example, in Figure Da there is no direct information pass from W.K Q to W while in Figure 9c such exchange is done through edge Ds and Hs. such [kj [h]
[M]
observation (there are multiple paths that connect the weight parameter matrices) that leads to our idea of comparing expressive power between different structures.
13

---

Published as a conference paper at ICLR 2022
B SUPPLEMENTARY MATERIAL FOR MULTI-HEAD SELF-ATTENTION
In the section, we prove the equivalence between the tensor diagram and the equations for multi-head self-attention (MHSA). To establish such equivalence, will need show that both tensor diagram and the matrix equations lead to the same result in element-wise notation.
X
(a) Single-Head Self-Attention
(b) Single-Head to Multi-Head
(c) Multi-head self-attention
Figure 9: Tensor diagrams from single-head self-attention to multi-head self-attention. Figure 9a is the lensor diagram representation of single-head self-attention. Figure 9c is the tensor-diagram representation of multi-head self-atiention.
For convenience of reference, we recap both representations in Equation B.1
and Figure
Q [+]
x [k]
K(A]
X
VIN
X
(B.la)
head
[|]
softmax
Q
D
Vini,
(B.
M head nJ) head
head [i] w
(B.1c)
Note that the query weight tensor 4 comes from concatenation of query weight matrices [k] Similarly, W WV are concatenated from (WK)) H h==
(W-Q [M] Mh=1 such that W hii. w
[W]h)] such that W
hiih WK [h]
[h) W. W. Furthermore, denote Q. K. as the
reshaped tensors from (Q M= [M V Finally, W is reshaped from Wo
#
To differentiate between equations and tensor diagrams in the proof, we add an over-line to every object in the tensor diagram throughout this section.
B.1 SINGLE-HEAD SELF-ATTENTIO (SHSA)
We will prove that the tensor diagram in Figure 9a is equivalent to Equations
B.la and B.Ib
Proof. For simplicity, we omit the subscript [h] this subsection.
We first show that a contraction of the nodes within the softmax box is equivalent to QKT Let a A QK Q (WK X where X N X E Mxx WQ, W FxD
We subscripts in the element- wise notation to differentiate edges when multiple edges with same dimension appear in the same equation.
Qnd
Σ h=l
x 1 fid
(B.2a)
nf.
K Kmdd
x ff
W
(B.2b)
fzd:
f==1
D
D
F
F
Anm ΣQ Kum
ΣΣ
x
we
W W X
(B.2c)
hh
hid fad mf-
d= d=
Denote be the result the subset of nodes in the softmax box as A.
A ΣΣΣX
f2d m/22
(B.3)
14

---

Published as a conference paper at ICLR 2022
Comparing Equations (B2cc) and (B.3), have A A. Let B softmax(A A) and B softmax(A) A), further have B softmax( (A) softmax(A) B.
Let M be the result of the tensor diagram and M be the matrix representation of the single-head attention. We have
Mnd Σ
Σ
nm X
W
(B.4)
1=1
F
M nd
B
mf W
(B.5)
nm
m=1
Comparing Equations B.4 and (B.5 we have C' is equivalent to C, which completes the proof.
B.2 MULTI-HEAD SELF- ATTENTITTNN (MHSA)
We now prove that tensor diagram Figure 9c is equivalent to Equation
Ic
Proof. Let be the number of heads. Let E N be the result of contraction the nodes
within the softmax box. From the last section, know that An...: A Q() (M)-
Applying
[N]
the softmax function B softmax (A), we have B softmax (eanxx(Qu pM K mj)
Let T be result of the tensor diagram and be that of Equation B.lc
Tafa
Σ
softmax Q ud Q (h). Mml.mdd WO
(hl.dass
(B.6)
m,h.da da
T n1/
Σ
softmax Sland, Kland, ) Wmmdy Wh
handg
(B.7)
m.
Comparing T and T, we complete the proof.
L
C CLOSER INVESTIGATION OF MHSA AND ITS COMPARISON TO THSA
Rethinking MHSA. There are two types contractions in MHSA as shown in Figure 4a Head- Contraction, a global contraction among all weights ? y WV and W along mode H, and Intent-Contractionn a local contraction between 8 and w well between V and Wo both along mode D. If we remove the Head-Contraction and maintain the Latent-Contraction only, multi-head is reduced to single-head with an additional contraction with the weight matrix W To understand the role of the two types contractions, we propose ablation studies by removing one of them and study the changes in the structure and the expressive power. It turns out that in both situation, MHSA falls into special case of THSA.
Repeated Proposition I with two special cases. Given R DH and other hyperparameters being the same, THSA module includes MHSA as special case. Specifically. (I) If C 1p1l D. i.e., D x D all-ones matrix, the THSA module reduces to MHSA with single head and latent dimension D (ie. a single-head self-attention with latent dimension D). (2) CfC Ih. i.e., an H x H identity matrix, the THSA reduces MHSA with H heads and latent dimension (i.e., a heads-only self-attention with heads).
(3) C Is @ (1p1,), i.e., a Kronecker product of an all-ones matrix 1p1 DxD and an identity matrix In E BRHH the THSA reduces to an MIISA with H heads and latent dimension D.
The proof of the above propositions is given in Appendix D
Definition (Stable Rank). Let A E N M the stable rank of A of is defined as E. o?/ max, 0; where a; is the singular value of A.
15

---

Published as a conference paper at ICLR 2022
X
W
W
X
x
w
W
stable rank
stable
1
w
W
X
w
W
X
Figure (Left) THSA includes a single-head self-attention. In this case, stable rank of the core tensor C' is 1. It is interesting to observe that this a single-head self-attention and this structure can be obtained by removing the head contraction from Figure 4a (Right) THSA includes MIIA with H heads and latent dimension 1. In this case, the stable rank equals the rank of the identical matrix, which is H, corresponding number heads of MHSA. Also, the structure on the right can be obtained by removing the latent contraction edges D from Figure 4a
X
F
B
Z
X
C= aa
D
i
T
X
Figure 11: This figure shows how to initialize THSA with MHSA. MHSA a special of THSA when the core tensor takes the Kronecker product of all-one matrix x H identity matrix. Note that in this case, the number of heads in MHSA equivalent to the stable rank of C.
The notion heads THSA. The concept of the number of heads in THSA generalizes to the stable rank of the matrix, allowing data-driven implicit training. The number of heads H in MHSA corresponds to the stable rank of the core matrix C, defined as Σ, of max; o; where a; is the singular value. is a useful surrogate for the rank because is largely unaffected by tiny singular values.
As shown in Appendix C when C is constructed by the outer product of two vectors, the stable rank of C' is 1, while in this case, the whole structure is equivalent to a single-head self-attention. When C is an identical matrix as shown in appendix C the stable rank of C equals to the rank of the identical matrix, which is H, corresponding to the number of heads of MHSA. Also, the structure on the right can be obtained by removing the latent contraction edges D from Figure 4a When C' is the Kronecker product of the all-one matrix and the identical matrix, which corresponds to MHSA, the stable rank C' is H, and this number does not change in the training because that is fixed matrix in MHSA.
D SUPPLEMENTARY MATERIAL FOR TUFORMER AND PROOFS
We first formally describe how to compare the expressive power between two models.
Definition 3 (Expressive Power). Suppose we have two function classes F, G with the same source domain X and target domain y, i.e.. each function f E FF (or g E G) is a mapping from to D. We say G is more expressive than F if G F: for any f E F, there exists
G such that g(z) = f()), Vx E X. Furthermore, we say is strictly more expressive than if G 5 E: besides
G 2 F, there exists g G such that given any I there exists x E X and g(x) + fxx
16

---

Published as a conference paper at ICLR 2022
According to the definition, expressive power is partial order set. Therefore, expressive power of different models usually not comparable as two sets might have overlap and non-overlap regions. However, when one is superset or subset of the other, can unambiguously compare their expressive power. In the following theorem, we formally show that a Tunable-head self-attention (THSA) module has higher expressive than multi-head self- (MHSA) if the rank R in THSA is equal to the product of the number of heads H and the latent dimension D in MHSA.
Theorem (THSA is more expressive than MHSA). Given R = DH and other hyper-paramete parameters being the same, Tunable-head self-attention (THSA) module with rank R is expressive than multi-head self-attention (MHSA) module with heads and latent dimension D.
Theorem proved in Appendix D.2 shows that any function realizable by MHSA can be realized by THSA.
In this section, we will prove Proposition I and and Theorem in Section 3 In both theorems, we assume the rank in the latent dimension D in multi-head self-attention (MHSA), i.e., R DH.
self-attention (THSA) equals the product of the number of heads and
For convenience of reference, we recap the mathematical expressions and tensor diagrams for both multi-head self-attention (MHSA) and Tunable-head self-attention (THSA) in the following.
Multi-head Self-Attention (MHSA). An MHSA module has learnable parameters as H sets W.Y 1H [k] hmm (with 8 W.K Ahp
of (query, key, value) weight matrices { WWQ [h]: Wh's
[k] [b] RFx
for each h) and output weight matrix W RxD The module maps three input matrices X& Q E R XF X- XV E R Mxx into an output matrix M E IR NxF.
[b]
X W [Aj
KIN
X
[M]S Vin X W
(D.la)
head [b]
softmax (Q
K [b] Vmg.
(D.1b)
M
head 1],
head 2] head H) W
(D.lc)
In Equation D.la). Q (k)] E NXD Kmy Vi E R MxD are query, key, value matrices respectively.
In Equation D. we omit the scaling factor 1/vD for the softmax function (the scalar can be merged into the learnable parameters), and head [6] E NxD is the resulted matrix for head
Tunable-Head Self-Attention. A THSA module is parameterized by five matrices: query weight key weight matrix W E FxR a value weight matrix W FxR
matrix Tº E FXR
an output weight matrix W E RxF and additional core matrix C R RxR The module has the same input and output domains as in MHSA, i.e., takes three matrices X Q E RN x X RM as inputs and returns matrix T N. as output.
Q Q= K XKW V XV
(D.2a)
head,- = softmax Crs
V
(D.2b)
T head W head, head,,
head
(D.2c)
To distinguish THSA from MHSA, we use over-scored symbols whenever needed. In Equation (D.2a) Q N K, V MxR are query, key, value matrices respectively. In Equation (D.2b), again omit the scaling factor /R in the softmax function. Equation (D.2b) leads to head matrix head FRF where head,- R denotes the th column of the matrix head.
17

---

Published as a conference paper at ICLR 2022
D.1 PROOF OF PROPOSITION
We constructively prove that THSA reduces to three cases of MHSA if the query/key/value weight matrices THSA concatenate the corresponding matrices MHSA:
Proof of Proposition 7
W
W
w Pap
W
(D.3a)
[2]
[
W
WKK
W
W
(D.3b)
[IJ
[2
[
W
W
W
W
(D.3c)
[2
[
In addition, W WO We have W
W
E RFxR and W E RRRF
(1) For C= = 11lp,i i.. an R x all-ones matrix, we aim to prove that the THSA module reduces to single-head self-attention with H 1 and D R. Since there is only one head in the module, we have W = [i] w W and W = W As an immediate result,
[1]
Q xº = @ Q
Q [1]+
(D.4a)
[1]
K x x W [1]
(D.4b)
V x WW x W ViI-
(D.4c)
Since C' is all-ones matrix, i... Cro 1, Vr, rewrite Equation D.2b
R [qK
head, softmax
V, softmax
Q K V
(D.5)
Notice that the equation holds for each column r, we further write all R equations jointly as:
head = softmax (Q K
V
softmax (Q (T)) Viu = head [I]-
(D.6)
Finally, we express the output matrix as:
head W head W = M,
(D.7)
which concludes the reduction for the special case (1).
(2) For C JR. an R R identity matrix, we aim to prove that the THSA module reduces heads-only self-attention with
and D 1. Since the latent dimension D 1, i.c., each W [b] (or W.K. K [A) (1)) is vector, we have W w Q [r] /
WH and W W. [r] Therefore,
Q, = x
x@wo Q(r):
(D.8a)
K, x" y"WK [rl K(r):
(D.8b)
V, = X = - Virl-
(D.8c)
Since C is identity matrix, i.e., Crr and C.. 0, Vr s, rewrite Equation D.2b)
head, softmax @ K V,
softmax (Qi))KT)) Vir head
vj.
(D.9)
Finally, we express the output matrix as:
head W [head head ...; head m]] W M, which concludes the reduction for the special case (2).
(D.10)
(3) For C Is (1n1]), i... Kronecker product between identity matrix In E H and an all-ones matrix 1 RDD we aim prove that the THSA module reduces to multi-head
18

---

Published as a conference paper at ICLR 2022
self-attention (MHSA) module with H heads and latent dimension D. In this general case, have
= = -x
W we
[Q(1],
Q(m]
(D.11)
[1]
[2]
[M]
K X W X
W
W
W
K [[]
K
(D.12)
[1]
[2]
[M]]
V X'WW -X
W W
W
V
[v].. (2)
Vinl
(D.13)
[2]
[8]
Notice that C' is a block diagonal matrix blkdiag(1 1 p 1 11 b)), where each block is an all-ones matrix. Therefore we can rewrite Equation D.2b for = (h - 1) as:
1 (b-1)DeeshD (Q.k. V.
head (1-) softmax
(D.14)
softmax
Σ
Vin.a
(D.15)
softmax
[h] (b)
(D.16)
Since Equation (D16) holds for each column d, we write equations jointly as:
head 1)D+ds
head
softmax Q (T)
VIn = head (A]
(D.17)
Again, since Equation (.17) holds for each block h, we combine all equations as:
head [head]]] head
head
(D.18)
Finally, we express the output matrix as:
T head W heaa head [2) ".. head [|]
W M,
(D.19)
which concludes the reduction for the case (3).
D.2 PROOF OF HBOREM 4
We have shown that THSA reduces to MHSA when takes specific forms. This shows that THSA is more expressive than MHSA any mapping by MHSA module can also be realized by a THSA module.
Proof of Theorem 4
To prove that THSA is strictly more expressive than MHSA, need find mapping by THSA module that can not be realized by an MHSA module. It suffices to show that the matrix M is rank-deficient, while the matrix T can be full-rank. We assume MH min(N,
For convenience, we divide the matrix D for MHSA into blocks.
W where WO E R DxF is the matrix for the hth head. We now rewrite the output matrix M as:
T
WOT 22
(D.20)
M Σ softmax (0
KT)
ViN
W M8
(D.21)
RNX
RMXF
Since the matrix inside summation is product of two matrices of size Nx and MxF its rank is at most M. Use the property rank(A + B) < rank(A) + rank(B), we have rank(M) < MH <
min(N, (N,F) F). Similarly, we divide the matrix w for THSA into rows and rewrite the output matrix T as:
T Σ softmax Σ
Cr. Cr.
K. 1n
V., W
(D.22)
RMXX
RNXM
Using the same argument, we have rank(T) < MR. Since H < R (in fact, = DH), we can always find example such that rank(T") rank rank(M) M). This completes the proof.
19

---

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

---

Published as a conference paper at ICLR 2022
Furthermore, the extra memory overhead be eliminated from the engineering side using the check- preecci-iintig technique: intermediate values are not saved in the memory but recalculated when needed. We compare the runtime in both training and inference phase with check-pointing. As shown in fig. 12 under the same number of parameters, there is only slight increase in the training and basically the same runtime in the inference phase. With direct transformation from Transformer to Tuformer (for example, when initialized with pre-trained models), the runtime is still under reasonable range.
Language Modeling on Penn Treebank
Language Modeling on Penn Treebank
MHSA
1.4
MHSA
THSA
THSA
THSA with chockpointing
1.2
THSA with checkpointing
second per iteration
secomd iteratioon
1.0
2.5
0.8
0.4
1.5
0.2
10%
20%
50%
100%
150%
200%
10%
20%
50%
100%
150%
200%
model scale
model scale
Figure 12: (Left) Training and (Right) inference run-time comparison between MHSA, THSA with the same number of parameters and THSA with the number of parameters with check-pointing The batch size 256.
21


Reviews:
[{"review_id": "V_MjPwJB1Ah", "reviewer": "Reviewer_R5aP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The main idea of the paper, to interpret the weight tensors of MHSA as factors of a decomposition, is new and THSA appears to be supported by experimental results (although the results are so small that I wonder whether they are significant relative to hyper-parameter tuning and stochastic fluctuations of the training). However, no evidence is provided to justify the choice - why should it be Tucker Decomposition, and not Tensor Train or Tensor Ring decomposition? And how does that choice influence the properties of the attention mechanism?\n\nThe authors also claim the following contributions:\n1. tensor diagram notation for attention mechanism\n2. design space for attention block together with expressiveness analysis framework\n3. initialisation scheme from pre-trained Transformer\n\nWhile (3) has important applications in reusing pre-trained transformer models, (1) doesn't help much to understand the attention mechanism, as it will be readable only for people whose multi-linear algebra background is advanced enough to also reach the same understanding without diagrams. Whereas (2) effectively proves that, given some network, adding an additional parametric matrix somewhere in the middle of that network does not decrease its expressiveness, which is trivial\n\nIn my view, the general idea is valid and potentially interesting, but it would make the paper much stronger to revise it in terms of methodology and experiments, so as to actually study how different tensor decompositions influence the attention mechanism, and which ones might be suited best for specific layouts and tasks. However, this would be a completely different paper that would have to be reviewed again, so I am on the fence about the present version.\n\nThe presentation is mostly clear, but in a few places it is sloppy and appears unfinished and rushed. This would need to be reworked:\npage 4 - duplicate paragraphs at the end of the page\npage 5 - duplicate paragraphs in 3. and 3.1.\nDefinition 2: f \\in X should be probably f \\in F ?\n(page 3: For instant -> For instance)", "summary_of_the_paper": "The paper attempts for a better understanding of the mechanism of multi-head self-attention (MHSA). It proposes to interpret MHSA weight matrices as factors of the Tucker Decomposition and adds an additional core weight tensor to MHSA to obtain Tucker-head self-attention (THSA).", "summary_of_the_review": "The paper address an interesting and relevant problem, namely generalisation of the MHSA mechanism, and it presents interesting, if somewhat terse and preliminary, experimental results. However, while the view as tensor decomposition is promising, the seemingly ad-hoc commitment to Tucker without a clear rationale or systematic study is unsatisfactory and leaves an unfinished impression. The paper does put forward a novel idea and has high potential, but it would require a major overhaul to fulfil that potential and turn it into a possibly really strong and useful contribution, rather than a fleeting quick-shot.\n\nI find it difficult to give a recommendation for this paper: If one follows the conventional criteria, there is novelty and the potential to lead to interesting follow-up, so one might argue for acceptance, even though there isn't a strong case for future impact. On the other hand, it feels a bit premature and a bit more work could make it a lot better and more significant.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "Ti43MyDteDd", "reviewer": "Reviewer_hfb5", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Tensor network diagram seems unnecessary to the paper. Using Eq. (2a,3b,2c) and Eq. (1a,1b,1c) is enough for readers to understand the difference between MHSA and THSA. The diagrams in figure 4 convey the same idea but involve more technical content. The first contribution claimed by the author could be safely erased without affecting the completeness of this work.  Let me know if I am wrong.\n\nIn Sec. 3.2, as a reader, I am more curious about the proof for theorem 3 than definition 2. But the authors put definition 2 in the main paper and proof for theorem 3 in the appendix -- why not do the opposite?\n\n\nThe author claimed \"THSA *significantly* improves MHSA\", I am wondering whether the author actually did the significant test (as I encourage the authors to do it), or it would better use another word like \"notably\" instead.\n\n\nPlease consider *quantitatively* showing the inference time.\n\nSubscripts in Figure (3) seem wrong, please have a check. A_mr B_nr not  A_mr B_nr. The same happens in the text related to Eq 2, please double-check the superscripts there (the shape of X^Q, W^Q, Q, Q_r).\n\n\nThe technical content is quite close to  Cordonnier et.al 2020 https://openreview.net/forum?id=bK-rJMKrOsm , please consider explaining the difference. See figure 2  (left) in Cordonnier et.al 2020, the standard SAN can be considered as a special case of theirs. This is quite similar to  Proposition 1 in this paper -- of course, they are not the same, but with some differences. \n\n\n\n", "summary_of_the_paper": "\nThe paper proposes a THSA that generalizes the standard self-attention. THSA empirically outperforms the standard self-attention.", "summary_of_the_review": "The paper introduces a general Self-attention, which is interesting and insightful. However, a general model that costs more (time or space) but slightly performs better is not that attractive.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "0FwBzwXSqL2", "reviewer": "Reviewer_cD8e", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Questions\n===\n\nHow can a \"super node\" in Fig.4 be seen as a single tensor if a contraction of it crosses the boundary of a non-linear operation? This overloads the tensor diagram notation in an undesirable way, violating the multilinearity assumption. Taking this notation to the extreme, an MLP can be represented as a single \"super-node\" matrix, which spans multiple nonlinearities. Such a convention for tensor diagrams would be not much more useful than a typical block scheme of most neural architectures.\n\nTucker decomposition of a 4d-tensor leads to a 4d core tensor. If the core is a matrix, the authors should either carefully specify which \"variant\" of the decomposition is used (with specific references to external sources) or use a different term for the proposed parameterization.\n\nThe notation of T-junctions used to connect \"stacked\" modes to other nodes (such as to the matrix C from the top in Fig.4b) is crucial for understanding the proposed module; however, it is missing. \n\n\"The concept of the number of heads in THSA generalizes to the stable rank of the core matrix\" -- why can't we draw a similar connection with the standard MHSA? \n\nRelated to the previous question, formulas (2) operate an integer number of heads (R) - how does this go along with the claim of a learned number of heads?\n\nFormula (2b) rewrites a product of QK^T matrices for a single head through indexing along the R mode using s subscript. This corresponds to the rank factorization of the said product. This also seems to be equivalent to a single-head self-attention with R features instead of D, due to the fact that the values C_{rs} can be absorbed into either W^Q or W^K (which are in turn head-specific and may need an _{[r]} subscript). Or, if matrices W^Q and W^K are not head-specific anymore, this looks like a weight sharing scheme, in which one of the matrices is head-specific, and the other one is reused across heads. The main questions are: (1) where is Tucker decomposition in this formulation, (2) how does this equation correspond to Fig.4b, and (3) how exactly is this different from a standard MHSA?\n\nCould the empirical advantage be due to a changed parameterization of W^K and W^Q matrices? I would like to see an experiment with the same configuration as MHSA in terms of dimensions of activations (probably corresponds to +25% configuration from the paper) but done in a way that the fifth matrix C is introduced into the QK product (using small surgical changes to the original MHSA module code). Such an added matrix experiment would level the changes with the proposed THSA and allow us to see whether the improvement is due to such parameterization of the QK product. Since such decomposition does not increase the expressive power of the QK product, any change in performance is due to either a different initialization or learning rate of the product elements resulting from such parameterization. If weight sharing is the case, additional studies of sharing W^Q or W^K in an MHSA setting are required.\n\nThe claim about the possibility of initializing THSA from MHSA should be followed with a remark that this is only possible in the case of matching dimensions. \n\nThe variance across runs in Fig 5 looks small - was RNG seed changed explicitly between runs to affect weight initializations, as opposed to relying on the intrinsic variance of SGD and randomness of multi-processing data loading in pytorch?\n\nAround formulas (2), the first dimensions of Q and X^Q are F and N, respectively, but how is this possible if Q=X^Q W^Q?\n\n\nDelivery\n===\nAbstract: \"the number of heads is trainable\" is misleading, as this number is discrete and is not directly optimized.\n\nAbstract: \"allows initialization from existing pre-trained models\" is too broad, given that this is only possible in cases with a specific R.\n\nThe first occurrence of \"expressive power\" (except abstract) requires a clarification of the subject (a class of functions) - either as a citation or as a forward reference to the definition from the supplementary materials. It should also be made clear how expressive power is compared to explain when used together with \"higher\".\n\n\"(2) The number of heads is not trainable\": not clear how this is a limitation or a problem.\n\n\"re-parameterization of the weight matrices via a variant of Tucker Decomposition or Tucker Representation\" - the provided reference is too broad given the word \"variant\". A default Tucker decomposition assumes a core tensor of the same dimensionality as the decomposed tensor, which is not the case in this paper.\n\nThe proposed visualization of the softmax operation in the tensor diagram is good but requires clarification that normalization is applied along the mode containing the \"alpha\" point.\n\nThe \"Evaluatuon order\" section would benefit from a remark that the order affects the computational complexity of the overall contraction.\n\nWhile the notation of stacking matrices to obtain a 3d-tensor is clear, connecting these stacked modes together requires additional explanation. Furthermore, if this notation is used purely to denote the stacked nature of nodes, it should be made clear, too (and perhaps such edges can be color-coded differently). \n\nThe related work section misses an overview of tensor decompositions used in machine learning.\n\n\nWriting\n===\n\nUsage of the word \"guaranteed\" throughout the text requires explaining: by what means? If through a proof - the word \"proven\" is better suited.\n\n\"(2) The number of heads is not trainable ... capture different context information through each head\": if this is a known fact, the authors should consider adding a reference where this is stated.\n\nThe first occurrence of \"tensor representation theory\" - a reference is required (e.g., Kolda)\n\n\"we also propose a novel data-driven structure that\" - too vague.\n\n\"re-parameterization of the weight matrices\" - which ones / of what?\n\n\"for instant, in Section 2.1\" -> \"for instance, in Fig. 2.3\"\n\nnotation for subscript \"_{[h]}\" not introduced\n\nTexts of \"Single-head self-attention in tensor diagram\" and \"Multi-head self-attention in tensor diagram\" sections are nearly identical; Fig. 3a is never referenced nor explained. \n\n\nCode\n===\nThe provided source code is not readily usable for reproducibility. A few problems and recommendations regarding the file with the implementation referred to in the readme:\n\n1. The code contains multiple python syntax errors\n2. Local variable tucker_heads is not defined in a few places; perhaps \"self.\" is missing\n3. Unused code clutter (all branches where self.tucker_heads is False) \n4. Unused variables (e.g., c = self.c_proj(...))", "summary_of_the_paper": "The paper investigates the multi-head self-attention mechanism (MHSA) of transformer networks through the lens of tensor decompositions via tensor diagram notation. The authors propose an extension to MHSA inspired by the Tucker decomposition (termed THSA), analyze its expressive power, and demonstrate that it belongs to a class of more expressive functions than MHSA. Further, the authors show the positive effect of this drop-in replacement on the downstream tasks.", "summary_of_the_review": "The approach to analyzing MHSA taken by the authors is interesting and worth digging deeper, especially due to the clear empirical advantage of the proposed modification. However, the paper contains several issues in claims, theory, and connections between them. On the side of delivery, the paper has good potential, but in its current state, it is not well-polished and requires a major improvement to facilitate better understanding.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "-5Y9M5nnOxU", "reviewer": "Reviewer_5BZt", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n- Clear writing.\n- The proposed idea is somewhat novel and interesting.\n\t\nCons:\n- Motivation of THSA is not clear. Although the paper claims \"inspired by a Tucker tensor form\", what is the advantage of the Tucker tensor form? And why MHSA can be benefited from the Tucker tensor form?\n- The ability of extended tensor diagram is limited. A traditional tensor diagram need not consider the combination order for diagram nodes. However, the proposed extension requires contracting the nodes in the softmax box firstly. I agree that the diagram extension is sufficient to present the proposed method. Nevertheless, it may need to elaborate the limitation clearly, for example,\n\t- Difficult to denote multi-layers due to the order of softmax box;\n\t- More attention when applying the tensor diagram extension to other cases besides the self-attention condition. \n- More situations are supposed to be considered as the diagram extension is firstly proposed:\n\t- Is there a probability to take a node out from the box? If possible, what transformation needs to be applied?\n\t- Will there be a case multi-box overlap? When this case happened, how to do it?\n- In Definition 2, the expressive power seems not to make sense. Under the definition, a model with higher expressive power may not perform well. For instance, an over-fitting model may have large expressive power, however, its performance is lower, compared to a suitable model. Therefore, it may not provide a guarantee to THSA performance. By the way, I guess there is a typo in Definition 2, \"$f \\in \\mathcal{X}$\" - > \"$f \\in F$\".\n- In Table 1, the paper seems not to mention that the results are all derived through re-implementing. And the results in the original paper are higher (e.g., ALBERT: MNLI-91.3%, QNLI-99.2%), which makes a concern that the experiment is not prepared well.\n- Insufficient experiments. There are plenty of new works proposed recently. And if THSA is really consistently better than MHSA, THSA is supposed to achieve SOTA results on a task.\n- More training run-time for THSA. Will THSA (same num of param) be slower when compared with MHSA?", "summary_of_the_paper": "Focusing on the multi-head self-attention (MHSA) structure, this paper proposes an extension of the tensor diagram to denote self-attention (SA) structures more intuitively. Then inspired by the Tucker format, this paper also proposes a new form of SA named Tucker-Head Self-Attention (THSA), which can also be illustrated with the tensor diagram extension. Although THSA derives some success, THSA may lack theoretical guarantee and experiments seems insufficient.", "summary_of_the_review": "Although this paper has a good presentation, for lacking theoretical guarantee and insufficient experiments, I regret to give a weak rejection.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents a tensor diagram view of the multi-headed self-attention (MHSA) mechanism used in Transformer architectures, and by modifying the tensor diagram, introduces a strict generalization of MHSA called the Tucker-head self attention (THSA) mechanism. While there is some concern regarding the incremental nature of the proposition, the identification of where to usefully add the additional parameter that converts from MHSA to THSA was nontrivial, and the experimental results on the performance benefits across multiple tasks is convincing.
Decision: Accept (Poster)