Transformer acceleration in heterogeneous platforms

Aceleraci´on de Transformers sobre plataformas heterog´eneas Transformer acceleration in heterogeneous platforms TRABAJO DE FIN DE GRADO CURSO 2021/22 UNIVERSIDAD COMPLUTENSE MADRID Facultad de Inform´atica Doble Grado Ingenier´ıa Inform´atica y Matem´aticas Francisco Borja Lozano del Moral Tutora: Katzalin Olcoz Herrero Table of Contents 1 Introduction................................................... 2 2 Preamble: Model training and Backpropagation . . . . . . . . . . . . . . . . . . . . 3 2.1 Models and Loss Functions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 3 2.2 Optimization algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4 2.3 Backpropation algorithm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4 3 TransformerStructure.......................................... 6 4 Inner workings of a Transformer . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7 4.1 LinearApplication ........................................ 7 4.1.1 Forward propagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8 4.1.2 Backpropagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8 4.1.3 Feed Forward Layer.. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9 4.2 SoftmaxLayer ............................................ 10 4.2.1 Forward propagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 10 4.2.2 Backpropagation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11 4.3 Masking Application . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11 4.3.1 Forward propagation: Conditional Substitution Mask. . . . . 11 4.3.2 Forward propagation: Summing Mask. . . . . . . . . . . . . . . . . . . 12 4.3.3 Backpropagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12 4.4 Scaled Dot-Product Attention. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13 4.4.1 Forward propagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13 4.4.2 Backpropagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13 4.4.3 Masking case alterations. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 4.4.4 Effect of repeated inputs. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 4.5 Multi-Headed Attention . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18 4.5.1 Forward propagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18 4.5.2 Backpropagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 4.6 Layer Normalization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 4.6.1 Forward propagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 4.6.2 Backpropagation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22 4.7 Embedding............................................... 25 4.7.1 Forward propagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 4.7.2 Backpropagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 4.8 Positional Encoding . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 4.8.1 Forward propagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 4.8.2 Backpropagation. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 5 Output of the Transformer . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 28 6 Datasets and data processing . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32 6.1 Datasetselection. ......................................... 32 6.2 Tokenization. ............................................. 32 6.3 Datasetformat. ........................................... 33 6.4 Dataloader ............................................... 33 7 Implementation................................................ 34 7.1 HELENNA............................................... 34 7.2 Limitations of HELENNA . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34 7.3 Transformer support . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35 7.4 Forward propagation . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35 7.5 Backpropagation .......................................... 35 8 Results ....................................................... 36 8.1 Comparison between HELENNA and PyTorch . . . . . . . . . . . . . . . . 36 8.2 Comparison of memory layouts in HELENNA . . . . . . . . . . . . . . . . . 37 9 Conclusions ................................................... 38 A Appendix..................................................... 39 A.1 Dataset and Tokenization . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 39 A.2 Generic implementation of a Bucket Dataloader . . . . . . . . . . . . . . . 39 A.3 Tests for correctness . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 39 A.4 HELENNA............................................... 40 A.5 PyTorch ................................................. 40 Abstract Dise˜namos una las operaciones necesarias para ejecutar un Transformer ´optimamente en plataformas heterog´eneas como FPGAs, GPUs y CPU. En HELENNA, por limitaciones de tiempo, tuvimos que limtar la implementaci´on a CPU, sirviendo como prueba de concepto para la traslaci´on a las otras plataformas, pero sin obtener un rendimiento competitivo mediante OpenMP y configuraci´on de la memoria. 1 1 Introduction Natural language processing is of great interest due to the difficulty of having machines display ‘understanding’ of language. Inside this problem, translation is one of the tasks that receive most attention. In this paper, we will explore a model that revolutionized translation machine learning: Transformers. This model, appearing for the first time in [AIAYN(2017)], has a repeated structure that utilizes an innovative method to have the model learn relationships between words without requiring recursion. Our objective will be to implement a Transformer using low-level programming languages using diverse techniques and the use of heterogeneous platforms to rise the performance of the model as much as we can. We chose to implement it on HELENNA, a student research tool being currently developed to delve into deep-learning models. After explaining the main concepts we will depend on throughout the paper in the first chapter, we will delve into the structure of the model and how to best implement the computing the necessary gradients of each layer. On chapter 5, we will explore the possible outputs of the model and how the training may be accelerated through the use of all positions of the output. Next, we will explore how we may obtain data to stress-test our model with translation datasets and how best to load the data and prepare the batches used for training. Following that, explore the problems encountered during the implementation of Transformers in HELENNA and finishing with the performance achieved with it. 2 2 Preamble: Model training and Backpropagation Before delving into what makes the Transformer excel in translation and text prediction, it is necessary to first discuss what a model is, how it is evaluated, and how the training happens. 2.1 Models and Loss Functions There are several types of models when it comes to machine learning. For this work, we will focus on supervised classification models. Definition 1 (Supervised Classification Model.) [SC(2018)] Models used to classify samples among previously determined categories. They train using datasets that have what the expected category should be. E.g. A classifier that detects whether planes appear in an image would be a supervised classification model. Training a model is optimizing the parameters or weights of a model to obtain the desired results. The main method to optimize the weights is defining a function that evaluates how badly the model is doing and finding a direction that decreases it, in order to find local minima of the function. By convention, we will call this function the loss function, and it will receive the output of the model and a vector representing the expected result. Usually, both have the same dimensions, but this can differ in cases such as some classification loss functions, where the output of the model is the probability of choosing each class and the expected output is represented by a single index of the correct class. In decision-taking tasks, such as text translation/prediction, there is a small selection of commonly used loss functions. Among them, the most frequently used function is the Cross Entropy Loss or CEL (see [CEL(2020)]). This is the function that we will use for our benchmarks. For simplicity, we will represent the correct class output using a one-hot vector, which is one only in the index of the expected class and zero everywhere else. Following convention, we will define the function so that it receives the output of the model as the first argument and the correct output as the second argument. Definition 2 Cross Entropy Loss.1 CEL : (y,ˆ y)7→ − M X i=1 ˆ yilog(yi) (1) 1Because CEL is widely used right after a softmax application, sometimes CEL will be defined as softmax followed by the loss shown. This is done because the derivative of the input w.r.t. the loss can be collapsed into one subtraction. This occurs in established ML libraries such as PyTorch. 3 2.2 Optimization algorithm Among the optimization algorithms, however, the available selection of possibilities is much more diverse. In our case, because we will be focusing on evaluating the time each training step takes rather than the number of steps until the training is complete, we will use Batch Gradient Descent or BGD, as it is simple, fast and effective. A batch, or group of samples, will normally be defined as a sequence of pairs {(xn, yn)}B n=1, where xnis the input sample, and ynthe expected output; or as a pair of matrices (X, Y ) where each sample is stacked into rows or columns. Using this, we can now define BGD formally. Definition 3 Batch Gradient Descent. Let X∈RB×N,Y∈RB×M, where Bis the batch size. Given a function f:RN→RMthat depends on a weight w, and a loss function Loss :RM×RM→R, the optimization technique, each iteration, the weight is modified following the rule: w(t+ 1) 7−→ w(t)−λ1 M B X b=1 ∂ ∂w Loss(f(Xb|w), Yb) (2) where λ > 0is the learning rate. This method effectively moves the weight in the direction of a local minimum of the average loss of the training dataset, in the hopes of optimizing the expected value E[Loss(f(X), Y )|w]. For simplicity, we will denote Loss(f(X|w), Y ) as Loss(Input) when we compute gradients in following sections. 2.3 Backpropation algorithm Backpropation is a differentiation strategy used for fast derivation of nested functions to obtain the derivative w.r.t. variables inside the model. This is done by propagating backwards throughout the model, from the output layer to the input layer, the gradients of the loss w.r.t different middle values of the model in an attempt to reuse as much computation as possible during the computation of the gradients w.r.t the weights. For clarity we will use the following definition in this document. Definition 4 Because all gradients are w.r.t. matrices, we will use interchangeably the notation ∂f(X) ∂X c,d ≡∂f(X) ∂Xc,d to denote the gradient w.r.t the element Xc,d of f(X)given a function f: RN×M→R. One of the principal properties of gradients is the following: 4 Lemma 1. Given three functions fw,hand gw: gw:RN×M→RN′×M′(3) h:RN′×M′→R(4) fw=h◦gw:RN×M→R(5) where wis a weight parameter. Then, by the chain rule: ∂f(X) ∂Xc,d =X l,n ∂f(X) ∂g(X)l,n ∂g(X)l,n ∂Xc,d (6) ∂f(X) ∂w =X l,n ∂f(X) ∂g(X)l,n ∂g(X)l,n ∂w (7) Because they share the gradient of f(X) w.r.t. g(X), the computations can be reused. This is the property at the center of Backpropation. To compute the gradients of the loss w.r.t. a layer’s weights, it is sufficient to know the gradient of the loss w.r.t. the layer’s outputs, which coincidentally can be computed by knowing the gradient of the loss w.r.t. the next layer’s outputs. This causes a propagation of the gradient of the loss w.r.t. each layers’ outputs that go from the last layer to the first, making use of previous layer’s calculations. For clarity, we will often abuse notation and treat variables as functions of the input. I.e. Y:= g(X) and ∂f(X) ∂Xc,d =X l,n ∂f(X)) ∂Yl,n ∂Yl,n ∂Xc,d (8) In a way, treating the output as a replacement for the function. In the following sections, it will be common to name the input of a specific block inside the model as Xand for it to be a replacement for a function of the model’s input. E.g. X:= f(Input) where Input is the model’s input. Since matrix operations are already implemented in a very optimized way on most machine learning libraries, we will set out to simplify the expression of gradients and compress it as much as possible as efficient matrix operations. For that purpose, we will use the following techniques: 1. Undo dot-product PkAa,kBk,b = (AB)a,b 2. Apply Kronecker’s delta Pn k=1 δk,cf(k) = f(c) when n≥c. 3. Split a summation’s inner sum. X x f(x)(g(x) + h(x)) = X x f(x)g(x) + X x f(x)h(x) 4. Simplify into comprehensible matrix operations using: (a) Element-wise product: (A⊙B)a,b =Aa,bBa,b (b) Sum of columns: SumCols(A)a=PkAa,k (c) Sum of columns: SumRows(A)b=PkAk,b (d) Custom product: MulCols(A, b)c,d =Ac,dbc 5 3 Transformer Structure In order to translate a sentence of variable length, models used to be recursive, requiring as one execution per sentence, but requiring making many recursive stages within the model. This caused the backpropagation of the model to require taking into account large amounts of middle results, which made these models have higher time complexity than desired. A Transformer performs translation of sentences by using the source sentence and the current translation in order to decide the next word or token of the translation. Translation requires the execution of the model as many times as the number of tokens of the final translation. This might seem that it makes the model just as recursive, but it is not. The difference is in the cost of performing the optimization step. With a transformer, we can train the model to translate correctly each target token with single non-recursive backpropagations of gradients, which is far shorter than the backpropagation of an equivalent recursive model. The model, as seen in figure 1, has three distinct parts, the input processing, the encoders and decoders. The two sentences the model receives, the source and the target, are initially expressed with a vector of tokens, integers representing the words or partial words that make them up. In order to work with them, first we vectorize each token using the embeddings that hold maps from token to vector, and then sends the vectorized source and target sentences to, respectively, the encoder and decoder. The encoders are a series of modules made up of inner layers that are meant to ‘encode’ all the meaning of the tokens and the relationships between them and send the result to the encoder. The decoder, however, will get both the result of the encoders and the vectorized target and will use relationship between words within the target and with the source source in order to create an output. What makes the Transformer find relationships between positions so well are the Multi-Head Attention modules (see 4.4-4.5). They are an innovative way to have the model focus on the relationship between specific positions of the sentences. 6 ∂σ(ˆ Y)l,k ∂ˆ Ym,r =δm,lσ(ˆ Y)l,kδn,r −σ(ˆ Y)l,r =δm,lIl≤kσ(ˆ Y)l,kδk,r −Il≤rσ(ˆ Y)l,r (27) And thus, for it to not be zero it must be true that m=l,l≤kand either: 1. k=r, in which case m=l≤k=rwhich contradicts m > r. 2. l≤r, in which case m=l≤rwhich contradicts m>r. And both contradict m>r. Lemma 4. It is zero if any of the following is satisfied: 1. l > n 2. m>r 3. m=l 4.4 Scaled Dot-Product Attention The key aspect that allowed the Transformer to rise in renown was the innovative way they used to have the model process the relation between words between two given sentences. The Scaled Dot-Product Attention, first obtains a matrix of attention between words by performing a scalar product between all pairs of word-vectors of the given ‘queries’ and ‘keys’ through a matrix product and applies it to the ‘values’ with yet again other matrix multiplication. In the cases where we wish to prevent forward propagation of information within a sentence, hereby forcing the model to train only depending on previous words within the sentence, we will use the mask already explained in the section above. 4.4.1 Forward propagation. Atten(Q, K, V ) = σ(QKT √dk +Mask)V(28) where σis the softmax function and Mask can be optional. In cases where it is not used, it will just be the Zero matrix. 4.4.2 Backpropagation. To compute the gradient of the loss w.r.t the inputs, we will work our way backwards from the output to the input operation by operation through the scaled dot-product. For clarity, let us define some variables which should be useful later: Y:= QKT √dk (29) U:= σ(Y) (30) H:= Atten(Q, K, V ) = UV (31) 13 Fig. 2: Scaled Dot-Product Attention Fig. 3: Multi-Headed Attention The gradients w.r.t Uand Vare follow the same pattern shown in 4.1. For the first, we will go more in-depth as an example. ∂Loss(Input) ∂Uc,d =X l,n ∂Loss(Input) ∂Hl,n ∂Hl,n ∂Uc,d (1) =X l,n ∂Loss(Input) ∂Hl,n ∂PkUl,kVk,n ∂Uc,d =X l,n ∂Loss(Input) ∂Hl,n X k Vk,n ∂Ul,k ∂Uc,d =X l,n ∂Loss(Input) ∂Hl,n X k Vk,nδl,cδk,d (2) =X n ∂Loss(Input) ∂Hc,n Vd,n (1) =∂Loss(Input) ∂H VTc,d (32) 14 ∂Loss(Input) ∂Vc,d =X l,n ∂Loss(Input) ∂Hl,n ∂(PkUl,kVk,n) ∂Vc,d (1)(2) =UT∂Loss(Input) ∂H c,d (33) The next gradient of the loss, the one w.r.t Yis more complex. ∂Loss(Input) ∂Yc,d =X l,n ∂Loss(Input) ∂Ul,n ∂σ(Y)l,n ∂Yc,d (Eq.21) =X l,n ∂Loss(Input) ∂Ul,n δc,lσ(Y)l,nδn,d −σ(Y)l,d (2) =X n ∂Loss(Input) ∂Uc,n σ(Y)c,nδn,d −σ(Y)c,d (3) =X n ∂Loss(Input) ∂Uc,n σ(Y)c,nδn,d −X n ∂Loss(Input) ∂Uc,n σ(Y)c,nσ(Y)c,d (a) = (∂Loss(Input) ∂U ⊙σ(Y))c,d −X n (∂Loss(Input) ∂U ⊙σ(Y))c,nσ(Y)c,d (c) = (∂Loss(Input) ∂U ⊙σ(Y))c,d −σ(Y)c,d ·SumCols(∂Loss(Input) ∂U ⊙σ(Y))c (d) =∂Loss(Input) ∂U ⊙σ(Y)) −MulCols(σ(Y),SumCols(∂Loss(Input) ∂U ⊙σ(Y)))c,d (34) 15 The last two, once again, follow the same pattern from section 4.1. ∂Loss(Input) ∂Qc,d =X l,n ∂Loss(Input) ∂Yl,n ∂Yl,n ∂Qc,d =1 √dkX l,n ∂Loss(Input) ∂Yl,n ∂(PkQl,kKn,k) ∂Qc,d =1 √dkX l,n ∂Loss(Input) ∂Yl,n X k δc,lδd,kKn,k (2) =1 √dkX n ∂Loss(Input) ∂Yc,n Kn,d (1) =1 √dk (∂Loss(Input) ∂Y K)c,d (35) ∂Loss(Input) ∂Kc,d =X l,n ∂Loss(Input) ∂Yl,n ∂Yl,n ∂Kc,d =1 √dkX l,n ∂Loss(Input) ∂Yl,n ∂(PkQl,kKn,k) ∂Kc,d =1 √dkX l,n ∂Loss(Input) ∂Yl,n X k Ql,kδc,nδd,k 2 =1 √dkX lQl,d ∂Loss(Input) ∂Yl,c  1 =1 √dk(∂Loss(Input) ∂Y )TQc,d (36) This establishes a way to compute the gradient quickly, as shown in figure 4. 4.4.3 Masking case alterations. When the scaled dot-product is masked as discussed in section 4.3, Uis a triangular matrix, and as such, all the operations that include it or an element-wise product can be halved in cost doing the necessary changes to ignore the zero positions. 4.4.4 Effect of repeated inputs. One question that will arise is what to do when some or every input, i.e. Q, K and Vare one and the same. In this case, we can treat them as outputs of a element-wise function of Xthat is either the identity of the zero function. I.e. for the case where K=V=Q: Qc,d := q(Xc,d) = 0 (37) Kc,d := k(Xc,d) = Xc,d (38) Vc,d := v(Xc,d) = Xc,d (39) 16 We can use the lemma 1 to compute the gradient. ∂Loss(Input) ∂Xc,d =∂Loss(Input) ∂Qc,d ∂Qc,d ∂Xc,d +∂Loss(Input) ∂Kc,d ∂Kc,d ∂Xc,d +∂Loss(Input) ∂Vc,d ∂Vc,d ∂Xc,d = 0 + ∂Loss(Input) ∂Kc,d +∂Loss(Input) ∂Vc,d (40) Thus, the gradient of an input is the sum of all inputs it is equal to. We will see in section 4.5 that if the linear applications of each head inside the MHA are not the same, which requires linked weights. Since this is not common in the implementation of Transformers, we will work under the assumption that they are different in each head of the MHA. Table 2: Gradient formulas for the Scaled Dot-Product. W.R.T. Gradient Formula U∂Loss(Input) ∂H VT V UT∂Loss(Input) ∂H Y∂Loss(Input) ∂U ⊙σ(Y)−MulColsσ(Y),SumCols(∂Loss(Input) ∂U ⊙σ(Y)) Q1 √dk ∂Loss(Input) ∂Y K K1 √dk (∂Loss(Input) ∂Y )TQ Note. X∈RN×din , W ∈Rdin×dout , b ∈R1×dout , L ∈RN×dout 17 Fig. 4: Order of operations gradients w.r.t the inputs Q RM×dk U RM×L K RL×dk V RL×dk ∂Loss(X) ∂H RM×dk ∂Loss(X) ∂U O(MLdk) ∂Loss(X) ∂V O(LMdk) ∂Loss(X) ∂U ⊙U O(LM) SumCols O(LM) MultCols O(LM) ∂Loss(X) ∂Y O(LM) ∂Loss(X) ∂K O(dkLM)∂Loss(X) ∂Q O(dkLM) 4.5 Multi-Headed Attention The multi-head attention function is the core module inside a transformer. It is also the most costly part of the model. The inputs for it are three matrices Q, K and V. These are the keys and values and queries. In a transformer, we always have that the input Kand Vare identically the same input, this can be observed in 1, we represent that with K≡V. As seen in figures 2 and 1, in most cases all three coincide, but in the decoder, in the second Multi-Head Attention they differ. Because of the differences in dimensions depending on where the MHA is located we will work with K, V ∈RL×dmand Q∈RM×dmwhere N, M ∈ {S, T }. 4.5.1 Forward propagation. It is multi-head, because inside it, hheads perform independently “Scaled Dot-Product Attention” to their inputs (see figure 3). Since there will not be cases of matrix powers, we will use superindices to 18 differentiate between matrices associated with each head, i.e. Qiwould be the input Q(see figure 2) received by the head i. The output of the multi-head attention is a linear application of the concatenation of the outputs of every head, which we will call Hi∈Rdk×Tfor the output of head i. M:= MultiHeadAtten(Q, K, V ) = Linear(Concat(H1, . . . , Hh), Wo, bo) (41) where: Concat(H1, . . . , Hh) :=      H1 . . . Hh      (42) To abbreviate, we will use Cto represent Concat(H1, . . . , Hh) from this point onwards. We will call the inputs of head ias Qi, Ki, V i. Each of these, linear applications of the original Q, K, V . This decreases the dimensions of the inputs to L×dkand M×dkrespectively. Qi= Linear(Q, Wi q, bq) (43) Ki= Linear(K, W i k, bk) (44) Vi= Linear(V, W i v, bv) (45) Using these inputs, each heads performs the “Scaled Dot-Product Attention”: Hi:= Atten(Qi, Ki, V i) A detail that is relevant to the implementation of the forward pass is the possibility of implementing the hlinear projections associated with each input at once. Lemma 5. Compression of linear applications. Linear(X,      W1 . . . Wh      ,(b1···bh)) = Linear(X, W1, b1),··· ,Linear(X, Wh, bh) (46) where X∈RN×din , W i∈Rdout×din and bi∈R1×dout . 19 Proof. Let n′=i·dout +n, then: Linear(X,      W1 . . . Wh      ,(b1···bh))l,n′(47) =X(W1)T··· (Wh)Tl,n′+     b1. . . bh . . . b1. . . bh     l,n′ (48) =X(Wi)Tl,n +     bi . . . bi     l,n = Linear(X, Wh, bh)l,n (49) thus obtaining the desired matrices in one dot product if desired. This is used in the implementation of the forward pass ??. 4.5.2 Backpropagation. As always, we assume that we have ∂Loss(Input) ∂M and we want to have ∂Loss(Input) ∂X for X∈ {Q, K, V }. To do that, we will use backpropagation inside the MHA. Due to the definition of M, we can use the formulas of section 4.1 to compute ∂Loss(Input) ∂C . Now, we would require the gradient w.r.t each head. Using the definition of C, it is trivial that it requires undoing the concatenation on the gradient w.r.t C: ∂Loss(Input) ∂C =     ∂Loss(Input) ∂H1 . . . ∂Loss(Input) ∂Hh      (50) Now, using the equations of section 4.4, we can obtain ∂Loss(Input) ∂Xifor Xi∈ {Qi, Ki, V i}. And, once again, using section 4.1 we can compute the gradient ∂Loss(Input) ∂X from ∂Loss(Input) ∂Xi. ∂Loss(Input) ∂Xc,d =X i ∂Loss(Input) ∂Xi l,n ∂Xi l,n ∂Xc,d (51) =X iX l,n ∂Loss(Input) ∂Xi l,n ∂Xi l,n ∂Xc,d (52) Eq.12 =X i (∂Loss(Input) ∂XiWi)c,d (53) 20 Thus being the sum of the gradients w.r.t each linear application’s input. Table 3: Gradient formulas specific to Multi-Head Attention. W.R.T. Complexity Gradient Formula HiO(1) Rows [(i−1)dk, i dk] of ∂Loss(Input) ∂C Xi∈ {Qi, Ki, V i} O(hdmdk·max(M, L)) Pi ∂Loss(Input) ∂XiWi Note 1. Refer to Linear Application and Scaled Dot-Product for middle-point gradients. Note 2. Xi∈Rdm×M∪Rdm×L, W ∈Rdm×dk, C ∈Rhdk×T 4.6 Layer Normalization Layer normalization, not to be mistaken for batch normalization, normalizes across the position features instead of the samples. Let us assume that the input is in the shape (B, N, E) where Bis the batch size, Nthe number of positions, and Eis the number of features per position. The difference would be that while Batch Normalization would normalize Bvalues of (N, E) random variables, Layer Normalization would normalize Evalues of (B, N) random variables. During the following computations, due to all operations being done iterating across E, we can, without loss of generalization, treat the input as having shape (B·N, E) to simplify the resulting expressions. Note that in a Transformer, E always coincides with dmodel. 4.6.1 Forward propagation. The mathematical expression of the forward pass for weighted layer normalization is: Yi,j =ˆ Yi,jγj+βj=Xi,j −µi pσ2 i+εγj+βj(54) Where: –X∈RB·N×Eis the input, –ˆ Y∈RB·N×Eis the normalized input, –γ,  β∈REare the weights, –µiand σ2 iare the mean and variance of the row vector Xi, –and ε > 0 will be a small quantity used in order to avoid cases with zero variance. 21 4.6.2 Backpropagation We will first compute the gradient of the error w.r.t the input. The expression, will be as follows: ∂Loss(Input) ∂Xc,d =X i,j ∂Loss(Input) ∂Yi,j ∂Yi,j ∂Xc,d (55) where ∂Yi,j ∂Xc,d =γj· ∂(Xi,j −µi) ∂Xc,d ·pσ2 i+ε−∂(√σ2 i+ε) ∂Xc,d ·(Xi,j −µi) σ2 i+ε =γj·∂(Xi,j −µi) ∂Xc,d 1 pσ2 i+ε−∂(pσ2 i+ε) ∂Xc,d (Xi,j −µi) σ2 i+ε (56) Now, to we will require the derivative of the numerator and denominator of the normalization w.r.t. the input. The first is straightforward. ∂(Xi,j −µi) ∂Xc,d =δi,c(δj,d −1 E) (57) For the denominator, we will do the derivative in two steps. First, the derivative of the variance w.r.t. the input. ∂σ2 i ∂Xc,d =1 n−1 n X j=1 ∂(Xi,j −µi)2 ∂Xc,d =1 n−1 n X k=1 2(Xi,k −µi)∂(Xi,k −µi) ∂Xc,d Eq.57 =1 n−1 n X k=1 2(Xi,k −µi)δi,c(δk,d −1 E) (3) =2δi,c n−1 n X k=1 (Xi,k −µi)δk,d −1 E n X k=1 (Xi,k −µi) (2) =2δi,c n−1(Xi,d −µi)−1 E n X k=1 (Xi,k −µi) =2δi,c n−1(Xi,d −µi)−0=2δi,c n−1(Xi,d −µi) (58) With this, the derivative of the denominator of the normalization w.r.t. the input wil be: ∂(pσ2 i+ε) ∂Xc,d Eq.58 =1 2pσ2 i+ε·2δi,c n−1(Xi,d −µi) =δi,c (E−1)pσ2 i+ε·(Xi,d −µi) (59) 22 Let us obtain what the output Hwill be compared to H′. QKT=  Q′ q KT=  Q′KT qKT (74) The matrix Ywill be: Y=1 dk  Q′KT qKT =  Y′ 1 dkqKT (75) And in consequence: U=σ(1 dk  Q′KT qKT ) =   σ(Y′) σ(1 dkqKT) =  U′ σ(1 dkqKT) (76) Then, the output of the head must be: H=  U′ σ(1 dkqKT) V=  H′ σ(1 dkqKT)V (77) Thus, the output only changes in the last row. This means that the new token does not affect the output associated with the previous tokens. In the previous case, the input Qi, which we replaced with Q, was the only matrix which changed in size. Let us see the difference in the outputs of the first MHA of the decoder when adding an additional token. We will use the same strategy to decrease superindices and we will work with the output of head i. Here, the head’s three inputs Q, K, V are positionwise linear applications of the MHA inputs, which are all one and the same (see figure 1). Because of that, we have that the inputs of the head iwill be: Q=  Q′ q =  Linear(X′, Wi q, bi q) Linear(x, Wi q, bi q) (78) and similarly, omitting the linear applications of X: K=  K′ k and V=  V′ v (79) where q,k,v∈R1×dk. The denotation for what the output of the head would be using previous matrices will be denoted as: H′=U′V′=σ(Y′)V′=σ(Q′(K′)T √dk +Mask)V′(80) 29 Let us obtain what the output Hwill be compared to H′. QKT=  Q′ q   K′ k  T =  Q′(K′)TQ′kT q(K′)TqkT (81) And thus, (QKT √dk +Mask) = 1 √dk  Q′(K′)T+Mask′Q′kT −∞ qkT (82) The matrix Ywill be Y=  Y′1 √dkQ′kT −∞ 1 √dkqkT (83) Lemma 6. If y∈Rnand y∈Rthen, assuming y:= hy′yi(84) α:= Pn i=1 exp yi Pn+1 i=1 exp yi (85) β:= 1 Pn+1 i=1 exp yi (86) then the softmax of the vector ysatisfies: σ(y) = hσ(y)α βeyi=hσ(y′)α σ(y)n+1i(87) I.e. the first npositions are the scaled softmax of y’. Using this lemma, using the lemma on the rows of Y, and noting that U′ k=σ(Y′ k) (88) then, for rj:= exp( 1 √dkQ′ jkT) and αj:= Pp i=1 exp Y′ j,i Pp i=1 exp Y′ j,i +rj (89) βj:= 1 Pp i=1 exp Y′ j,i +rj (90) we have: U=         (U′)1α1β1exp( 1 √dkQ′ 1kT) . . .. . . (U′)pαpβpexp( 1 √dkQ′ pkT) 01         (91) 30 And thus, the output is: Atten(Q, K, V ) =         (U′)1α1β1exp( 1 √dkQ′ 1kT) . . .. . . (U′)pαpβpexp( 1 √dkQ′ pkT) 01           V′ v  =         U′ 1V′α1+β1exp( 1 √dkQ′ 1kT)v . . . U′ pV′αp+βpexp( 1 √dkQ′ pkT)v v         =         H′ 1α1+β1exp( 1 √dkQ′ 1kT)v . . . H′ pαp+βpexp( 1 √dkQ′ pkT)v v         (92) Then the distance of row Hjfrom the row H′ jis ∥H′ j−(H′ jαj+βjrjv)∥ =∥H′ j(1 −αj)−βjrjv∥(93) =∥H′ j rj Pp i=1 exp Y′ j,i +rj−rj Pp i=1 exp Y′ j,i +rj v∥(94) =rj Pp i=1 exp Y′ j,i +rj∥H′ j−v∥(95) ≤ ∥H′ j−v∥(96) In conclusion, the distance between the rows from matrix H′and the first p rows of His at most the distance to the vector vand depends on rj,1≤j≤p. Due to that, the first prows of the output when the input of the decoder is p+ 1 rows will move away from the output if prows had been used. In practical terms, this means that optimizing output for the translation of the (p+ 1)-th token does not necessarily imply the optimization of the model for the case with ptokens. The decoder, however, uses the Add & Norm modules, and this addition will decrease the distance between input and output of the MHAs. As we have seen, the place where the added (p+ 1)-th row affects the general output of the model is in the self-attention of the decoder. There, an addition module sums the input and output of the self-attention, serving as a limiter on how changed the output of the first prows can be due to the added row. This increases position 31 independence and can help in achieving better translations for the first ptokens when optimizing the output for the translation of the (p+ 1)-th token when optimizing all the outputs of the model and not only the last. In conclusion, learning to translate one token of the translation may help to train all previous tokens of the translation, which would accelerate learning. 6 Datasets and data processing In order to benchmark the model’s true performance, we set out to make use of real datasets used for famous translation tasks and we processed it in order to use it to estimate the average time the model takes in each layer taking into account the oscillation between sentence lengths and batch sizes that are natural in translation training. 6.1 Dataset selection. In NLP tasks, finding a large number of valid samples be a difficult task due to the large amount of typos found in human-produced text. Because of this, the datasets used to train models tend to be few in number and with sentences varying wildly in length and complexity. Some of the datasets more used for this are the WMT datasets. We used the WMT 2014 English to German dataset for our measurements. 6.2 Tokenization. The format of the datasets are usually plain text containing the sentences, and written in an specific language. Our model uses as inputs integers representing the sentence in tokens, thus, some data processing has to take place. The process of transforming a text into a series of tokens, each having a unique assigned integer, is called tokenization. One of the possible ways to tokenize the text is to have every word become one token and assign each as a different integer. This has been found to be suboptimal, because it does not correctly capture the similarities between words. This leads to the use of sub-word tokenization. The method we will use will be Byte-Pair Encoding. This algorithm finds commonly found groups of letters that tend to appear together and assigns them each a token. This leads to the tendency of assigning prefixes, suffixes and word roots as unique tokens. Because we are working primary with a model used for translation tasks, we have two options at hand. We can create two different tokenizations, which would cause the vocabulary size of the original and target languages be asymmetrical, and relatively small, or merge all words in a single corpus and have one large vocabulary of tokens. The first option allows us to keep the embeddings relatively small in our model. The second, however, allows our model the possibility of using the same embedding in the encoder and decoder, and, if the languages are closely related, 32 could lead to further compression and a total embedding size smaller than two different ones. In our case, we selected the second approach. Due to this algorithm being out of the scope of this work, we chose to use an existing implementation of the algorithm for the tokenization of our dataset. 2 6.3 Dataset format. Once the dataset has been tokenized, we run into the problem of dataset size and reading. Depending on the method we use to write the dataset into a file, the reading could become a bottleneck during our training. The primary issue we had to deal with was irregular sample length. If each sample has a different length, shuffling and picking random samples would cause a large reading time overhead. To solve this, we decided to store each sample padded to reach the same length and storing the original length (see 7). Another issue was which file format to use. Due to the BPE implementation being in Python, we decided to use the NumPy file format and add support NumPy file reading in C to the project. Table 7: Dataset structure. S T s1s2··· sN−1sNt1t2··· tN−1tN Source length Target length Padded source Padded target 6.4 Dataloader In NLP tasks, samples can vary in size anywhere between single words to entire paragraphs. Because of this, a batch of randomly picked samples of a dataset can potentially have widely different sentence lengths. In order to keep the batches samples’ lengths similar, we used a bucket dataloader. Firstly, we stored the samples in buckets of similar amounts of tokens. Using this, the dataloader can shuffle each bucket internally and prioritize picking samples of the same bucket for each batch. Using this method, we can assure that all samples within the same batch have similar length and we can have each batch average the same amount of tokens. However, in order to take advantage of the full resources at our disposal during training, when a bucket lacks enough samples to complete a batch, we will complete it with samples with from similar buckets. One last thing of note is the reading method used for batch loading. There are two main options when it comes to loading data, one is to move the entire dataset into RAM memory, making reading faster, and the other is to read each batch 2More information can be found in [BPE(1999)]. 33 from disk as needed, decreasing memory use. In our case, because translation datasets can be very large and the time required to read a batch from memory is rather small, we will read each batch from disk directly. 7 Implementation 7.1 HELENNA Heterogeneous Learning Neural Network Application or HELENNA, is a neural network framework written primarily in C which allows to define and use neural network models with the purpose of teaching and researching computer architecture concepts. We chose HELENNA to implement Transformers because the purpose of this paper is researching the options of implementation and their performance for deep-learning models, which closely matches the purpose of HELENNA, and seemed to already have the majority of the tools necessary to implement a Transformer, which would allow more time for fine tuning optimizations specifically for Transformers and the exploration of implementations in other platforms. 7.2 Limitations of HELENNA HELENNA was originally implemented with a focus on supporting ANNs. The first problem that arises from this is that it caused their implementations to primarily expect linear inputs for many layers, such as the Linear and Softmax layers. Which leads to not having a clear support for multidimensional samples. This forced the implementation of new layers and increased the workload significantly. The main objective was the comparison of implementations mixing other platforms (GPUs, FPGAs,. . . ) and technologies such as MKL. But these avenues could not be explored at this time, as barely any part of the Transformer did could be executed in HELENNA without heavy modifications and this increased dramatically the expected amount of work. The second issue was that the standard memory layout used for storage of middle results, inputs and outputs has the sample size as last dimension. What would this entail? In our case, the input of the model would be (S, B) for the encoder and (T, B) for the decoder. In each of them, the embedding would use as output layout (N, dmodel, B) and this would remain the layout of all MHA, Layer Normalizations and Feed Forward layers. The issue with this is that on the operations required on most layers, the main iteration occurs over the dmodel dimension, which is not contiguous in memory. As a result, we can predict that an implementation using this memory layout strategy will worsen the performance. During partial tests, this was proven, which led to the a new implementation using dmodel as our last dimension and see how much performance we can obtain from a CPU using HELENNA. This, again, required further work and removed the possibility of implementing the model on FPGAs or GPUs. 34 7.3 Transformer support Among the layers needed to implement a Transformer that were already implemented in HELENNA, we have the ReLU, Addition and Dropout layers. In addition to these, even though we will require heavy modifications at times, we could re-purpose: –the existing Linear layer to create a Position-wise Linear layer, –the Softmax layer into a Position-Wise Softmax layer, –and the Batch Normalization layer into a Layer Normalization layer. With this, the only missing layers that would have to be built from the ground up would be: –the MHA layer, –the Embedding layer, –and the Positional Encoding layer. Several possible variations of the implementations will be tested in an attempt to fully utilize the parallelization capabilities through the use of the library OpenMP. Due to this being mainly changes in the iteration orders and the use of auxiliar variables, we will not be delving into the options tested. 3 Even though one of the main goals was the implementation of the Transformer’s layers in different platforms within HELENNA, such as GPUs or using MKL. The amount of unimplemented backbone functionalities and the obscurity of the libraries used caused us to limit the implementation to CPU only. 7.4 Forward propagation Due to the way HELENNA allocates the all the memory that will be required at the beginning, all layers are initialized to the maximum batch size and sentence length possible during the training. However, each batch will have differing size and dimensions. Because of this, on each forward propagation, at the start, a new functionality had to be added to support a constantly changing the dimensions of all inputs and outputs of layers to accurately send information between layers. 7.5 Backpropagation In HELENNA, cross-entropy loss is already implemented and functions well with their softmax layer. But, as we commented previously, HELENNA does not support softmax across one dimension. Because of this, a new implementation of cross-entropy loss and their gradient had to be implemented by modifying existing functions. 3Should the reader be interested in options explored, some examples are available in the code that are hidden through the use of defined global constants in the cpu.c file. 35 8 Results 8.1 Comparison between HELENNA and PyTorch In order to compare the performance of the implementations of each layer both in the backpropagation and forward pass, we will profile each layer independently, in order to compare each layer precisely without unwanted interactions with other running computation. As we can compare between the tables 8 and 9, the implementation of each layer using HELENNA’s memory layout and OMP to parallelize operations is far slower than the professionally created implementations found in PyTorch. Due to the fact that we can see a severe increase in time for every layer with large inputs, including a simple addition from 0.02 ms to 0.47 ms, it leads to the conclusion that a far suboptimal thread management and memory layout/lookup strategy is is used in HELENNA, which is independent of how optimal the complexity of the operations are supposed to be in theory. This is further corroborated by the metrics of the Embedding layer. This layer requires to look, for few inputs, to iterate over positions in memory which are contiguous in memory. This leads to a better use of each thread and fewer cache misses and achieves similar performance to PyTorch’s (0.05 ms compared to 0.13 ms). Table 8: Time metrics of HELENNA’s implementation with NBE memory layout. Layer Forward Backward Time/call Time share Time/call Time share MHA 21.08 ms 6.53% 92.80 ms 28.74% Linear 24.46 ms 10.52% 101.64 ms 43.72% Dropout 0.53 ms 0.29% 0.03 ms 0.02% Layer Normalization 2.88 ms 2.99% 0.24 ms 0.25% Embedding 0.05 ms 0.00% 2.04 ms 0.05% Positional Encoding 0.04 ms 0.00% 0.02 ms 0.00% Softmax 1.53 ms 0.01% 0.45 ms 0.01% Addition 0.47 ms 0.01% 1.14 ms 0.01% Note. Metrics using the Transformer presented in [AIAYN(2017)] with batches containing 150 −250 tokens. 36 Table 9: Time metrics of PyTorch’s implementation Layer Forward Backward Time/call Time share Time/call Time share MHA 1.47 ms 18.78% 2.54 ms 33.73% Linear 1.30 ms 12.20% 2.22 ms 13.44% Dropout 0.05 ms 2.06% 0.08 ms 3.65% Layer Normalization 0.05 ms 1.04% 0.08 ms 1.94% Embedding 0.13 ms 0.14% 4.23 ms 4.68% Positional Encoding 0.04 ms 0.15% - - Softmax 0.94 ms 0.32% - - Addition 0.02 ms 0.58% - - Note 1. Addition, Positional Encoding and Softmax do not allow stand-alone, backpropagation as they are meant to be optimized when connected to other layers. Note 2. Metrics using the Transformer used in [AIAYN(2017)] and batches containing 150 −250 tokens. 8.2 Comparison of memory layouts in HELENNA As observed in the previous section, the current bottleneck in the execution of a Transformer is the careful management of memory and parallelization. In order to increase the performance of the implementation in HELENNA, we set out to create an implementation using the dimension of the features, dmodel or E, as the last dimension in the memory layout of inputs and outputs of each layer. This allows us to further take advantage of each thread and manage memory use, by changing the iteration strategy to maximize the number of operations done in parallel on contiguous areas of the memory. As seen in table 10, the performance changed greatly on every layer with a large share of the computation time, decreasing by between 30% and 60% the time spent on each layer on the forward and backward passes. In total, due to these layers appearing repeatedly throughout the model, the total training time was decreased by 40%. In spite of this, the performance increase obtained through the optimization of the memory usage, the operations performed, and the parallelization strategy through OpenMP, it is still far away from the performance displayed by the professionally created PyTorch libraries of matrix operations and auto-gradient computation. 37 Table 10: Comparison of HELENNA’s implementations. Original New Pass Layer ms/call % ms/call % Time Decreaase MHA 29.06 5.85 28.27 7.98 2.73 Linear 35.02 0.98 33.78 13.24 3.54 Dropout 0.73 0.26 0.71 0.36 3.54 Layer Norm. 4.10 1.38 1.69 0.29 58.90 Forward Embedding 0.06 0.00 0.07 0.01 −12.58 Pos. Encoding 0.05 0.00 0.04 0.00 8.47 Softmax 2.34 0.03 1.10 0.02 52.97 Addition 0.02 0.01 0.02 0.01 4.72 MHA 170.59 34.32 100.04 28.27 41.36 Linear 153.32 42.82 105.97 41.54 30.88 Dropout 0.04 0.01 0.04 0.02 −1.52 Layer Norm. 0.62 0.21 0.61 0.29 1.93 Backward Embedding 1.44 0.03 1.48 0.05 −3.09 Pos. Encoding 0.02 0.00 0.02 0.00 −15.37 Softmax 0.64 0.03 0.34 0.01 47.12 Addition 0.04 0.01 0.04 0.02 −1.32 Note. We used the Transformer from [AIAYN(2017)] with batches containing 250 −450 tokens. 9 Conclusions Transformers, just like the great majority of machine learning models, depends greatly operations that require iterating over rows or columns of large matrices. Because of this, even if the formulas used for the forward and backward passes of the model had optimal theoretical time complexity, without careful memory layout planning the memory and thread scheduling, memory look-ups would become the bottleneck of the model. In order to further optimize the model in HELENNA to reach or surpass the performance in displayed in PyTorch, we would need to customize more closely both the data memory layout and the exact behavior of each thread during 38