Transformer acceleration in heterogeneous platforms

Acele aci´on de T ans o me s sob e pla a o mas he e og´eneas T ans o me accele a ion in he e ogeneous pla o ms TRABAJO DE FIN DE GRADO CURSO 2021/22 UNIVERSIDAD COMPLUTENSE MADRID Facul ad de In o m´a ica Doble G ado Ingenie ´ıa In o m´a ica y Ma em´a icas F ancisco Bo ja Lozano del Mo al Tu o a: Ka zalin Olcoz He e o Table o Con en s 1 In oduc ion................................................... 2 2 P eamble: Model aining and Backp opaga ion . . . . . . . . . . . . . . . . . . . . 3 2.1 Models and Loss Func ions . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 3 2.2 Op imiza ion algo i hm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4 2.3 Backp opa ion algo i hm . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 4 3 T ans o me S uc u e.......................................... 6 4 Inne wo kings o a T ans o me . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 7 4.1 Linea Applica ion ........................................ 7 4.1.1 Fo wa d p opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8 4.1.2 Backp opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 8 4.1.3 Feed Fo wa d Laye .. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 9 4.2 So maxLaye ............................................ 10 4.2.1 Fo wa d p opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 10 4.2.2 Backp opaga ion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11 4.3 Masking Applica ion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 11 4.3.1 Fo wa d p opaga ion: Condi ional Subs i u ion Mask. . . . . 11 4.3.2 Fo wa d p opaga ion: Summing Mask. . . . . . . . . . . . . . . . . . . 12 4.3.3 Backp opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 12 4.4 Scaled Do -P oduc A en ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13 4.4.1 Fo wa d p opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13 4.4.2 Backp opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 13 4.4.3 Masking case al e a ions. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 4.4.4 E ec o epea ed inpu s. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 16 4.5 Mul i-Headed A en ion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18 4.5.1 Fo wa d p opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 18 4.5.2 Backp opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 20 4.6 Laye No maliza ion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 4.6.1 Fo wa d p opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 21 4.6.2 Backp opaga ion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 22 4.7 Embedding............................................... 25 4.7.1 Fo wa d p opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 4.7.2 Backp opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 25 4.8 Posi ional Encoding . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 26 4.8.1 Fo wa d p opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 4.8.2 Backp opaga ion. . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 27 5 Ou pu o he T ans o me . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 28 6 Da ase s and da a p ocessing . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 32 6.1 Da ase selec ion. ......................................... 32 6.2 Tokeniza ion. ............................................. 32 6.3 Da ase o ma . ........................................... 33 6.4 Da aloade ............................................... 33 7 Implemen a ion................................................ 34 7.1 HELENNA............................................... 34 7.2 Limi a ions o HELENNA . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 34 7.3 T ans o me suppo . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35 7.4 Fo wa d p opaga ion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 35 7.5 Backp opaga ion .......................................... 35 8 Resul s ....................................................... 36 8.1 Compa ison be ween HELENNA and PyTo ch . . . . . . . . . . . . . . . . 36 8.2 Compa ison o memo y layou s in HELENNA . . . . . . . . . . . . . . . . . 37 9 Conclusions ................................................... 38 A Appendix..................................................... 39 A.1 Da ase and Tokeniza ion . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 39 A.2 Gene ic implemen a ion o a Bucke Da aloade . . . . . . . . . . . . . . . 39 A.3 Tes s o co ec ness . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . . 39 A.4 HELENNA............................................... 40 A.5 PyTo ch ................................................. 40 Abs ac Dise˜namos una las ope aciones necesa ias pa a ejecu a un T ans o me ´op ima- men e en pla a o mas he e og´eneas como FPGAs, GPUs y CPU. En HELENNA, po limi aciones de iempo, u imos que lim a la implemen aci´on a CPU, si iendo como p ueba de concep o pa a la aslaci´on a las o as pla a o mas, pe o sin ob ene un endimien o compe i i o median e OpenMP y con igu aci´on de la memo ia. 1 1 In oduc ion Na u al language p ocessing is o g ea in e es due o he di icul y o ha ing machines display ‘unde s anding’ o language. Inside his p oblem, ansla ion is one o he asks ha ecei e mos a en ion. In his pape , we will explo e a model ha e olu ionized ansla ion machine lea ning: T ans o me s. This model, appea ing o he i s ime in [AIAYN(2017)], has a epea ed s uc u e ha u ilizes an inno a i e me hod o ha e he model lea n ela ionships be ween wo ds wi hou equi ing ecu sion. Ou objec i e will be o implemen a T ans o me using low-le el p og am- ming languages using di e se echniques and he use o he e ogeneous pla o ms o ise he pe o mance o he model as much as we can. We chose o implemen i on HELENNA, a s uden esea ch ool being cu en ly de eloped o del e in o deep-lea ning models. A e explaining he main concep s we will depend on h oughou he pape in he i s chap e , we will del e in o he s uc u e o he model and how o bes implemen he compu ing he necessa y g adien s o each laye . On chap e 5, we will explo e he possible ou pu s o he model and how he aining may be accele a ed h ough he use o all posi ions o he ou pu . Nex , we will explo e how we may ob ain da a o s ess- es ou model wi h ansla ion da ase s and how bes o load he da a and p epa e he ba ches used o aining. Following ha , explo e he p oblems encoun e ed du ing he implemen a ion o T ans o me s in HELENNA and inishing wi h he pe o mance achie ed wi h i . 2 2 P eamble: Model aining and Backp opaga ion Be o e del ing in o wha makes he T ans o me excel in ansla ion and ex p edic ion, i is necessa y o i s discuss wha a model is, how i is e alua ed, and how he aining happens. 2.1 Models and Loss Func ions The e a e se e al ypes o models when i comes o machine lea ning. Fo his wo k, we will ocus on supe ised classi ica ion models. De ini ion 1 (Supe ised Classi ica ion Model.) [SC(2018)] Models used o classi y samples among p e iously de e mined ca ego ies. They ain using da ase s ha ha e wha he expec ed ca ego y should be. E.g. A classi ie ha de ec s whe he planes appea in an image would be a supe ised classi ica ion model. T aining a model is op imizing he pa ame e s o weigh s o a model o ob ain he desi ed esul s. The main me hod o op imize he weigh s is de ining a unc ion ha e alua es how badly he model is doing and inding a di ec ion ha dec eases i , in o de o ind local minima o he unc ion. By con en ion, we will call his unc ion he loss unc ion, and i will ecei e he ou pu o he model and a ec o ep esen ing he expec ed esul . Usu- ally, bo h ha e he same dimensions, bu his can di e in cases such as some classi ica ion loss unc ions, whe e he ou pu o he model is he p obabili y o choosing each class and he expec ed ou pu is ep esen ed by a single index o he co ec class. In decision- aking asks, such as ex ansla ion/p edic ion, he e is a small selec ion o commonly used loss unc ions. Among hem, he mos equen ly used unc ion is he C oss En opy Loss o CEL (see [CEL(2020)]). This is he unc ion ha we will use o ou benchma ks. Fo simplici y, we will ep esen he co ec class ou pu using a one-ho ec o , which is one only in he index o he expec ed class and ze o e e ywhe e else. Following con en ion, we will de ine he unc ion so ha i ecei es he ou pu o he model as he i s a gumen and he co ec ou pu as he second a gumen . De ini ion 2 C oss En opy Loss.1 CEL : (y,ˆ y)7→ − M X i=1 ˆ yilog(yi) (1) 1Because CEL is widely used igh a e a so max applica ion, some imes CEL will be de ined as so max ollowed by he loss shown. This is done because he de i a i e o he inpu w. . . he loss can be collapsed in o one sub ac ion. This occu s in es ablished ML lib a ies such as PyTo ch. 3 2.2 Op imiza ion algo i hm Among he op imiza ion algo i hms, howe e , he a ailable selec ion o possibil- i ies is much mo e di e se. In ou case, because we will be ocusing on e alua ing he ime each aining s ep akes a he han he numbe o s eps un il he ain- ing is comple e, we will use Ba ch G adien Descen o BGD, as i is simple, as and e ec i e. A ba ch, o g oup o samples, will no mally be de ined as a sequence o pai s {(xn, yn)}B n=1, whe e xnis he inpu sample, and yn he expec ed ou pu ; o as a pai o ma ices (X, Y ) whe e each sample is s acked in o ows o columns. Using his, we can now de ine BGD o mally. De ini ion 3 Ba ch G adien Descen . Le X∈RB×N,Y∈RB×M, whe e Bis he ba ch size. Gi en a unc ion :RN→RM ha depends on a weigh w, and a loss unc ion Loss :RM×RM→R, he op imiza ion echnique, each i e a ion, he weigh is modi ied ollowing he ule: w( + 1) 7−→ w( )−λ1 M B X b=1 ∂ ∂w Loss( (Xb|w), Yb) (2) whe e λ > 0is he lea ning a e. This me hod e ec i ely mo es he weigh in he di ec ion o a local minimum o he a e age loss o he aining da ase , in he hopes o op imizing he expec ed alue E[Loss( (X), Y )|w]. Fo simplici y, we will deno e Loss( (X|w), Y ) as Loss(Inpu ) when we com- pu e g adien s in ollowing sec ions. 2.3 Backp opa ion algo i hm Backp opa ion is a di e en ia ion s a egy used o as de i a ion o nes ed unc ions o ob ain he de i a i e w. . . a iables inside he model. This is done by p opaga ing backwa ds h oughou he model, om he ou - pu laye o he inpu laye , he g adien s o he loss w. . di e en middle alues o he model in an a emp o euse as much compu a ion as possible du ing he compu a ion o he g adien s w. . he weigh s. Fo cla i y we will use he ollowing de ini ion in his documen . De ini ion 4 Because all g adien s a e w. . . ma ices, we will use in e change- ably he no a ion ∂ (X) ∂X c,d ≡∂ (X) ∂Xc,d o deno e he g adien w. . he elemen Xc,d o (X)gi en a unc ion : RN×M→R. One o he p incipal p ope ies o g adien s is he ollowing: 4 Lemma 1. Gi en h ee unc ions w,hand gw: gw:RN×M→RN′×M′(3) h:RN′×M′→R(4) w=h◦gw:RN×M→R(5) whe e wis a weigh pa ame e . Then, by he chain ule: ∂ (X) ∂Xc,d =X l,n ∂ (X) ∂g(X)l,n ∂g(X)l,n ∂Xc,d (6) ∂ (X) ∂w =X l,n ∂ (X) ∂g(X)l,n ∂g(X)l,n ∂w (7) Because hey sha e he g adien o (X) w. . . g(X), he compu a ions can be eused. This is he p ope y a he cen e o Backp opa ion. To compu e he g adien s o he loss w. . . a laye ’s weigh s, i is su icien o know he g adien o he loss w. . . he laye ’s ou pu s, which coinciden ally can be compu ed by knowing he g adien o he loss w. . . he nex laye ’s ou pu s. This causes a p opaga ion o he g adien o he loss w. . . each laye s’ ou pu s ha go om he las laye o he i s , making use o p e ious laye ’s calcula ions. Fo cla i y, we will o en abuse no a ion and ea a iables as unc ions o he inpu . I.e. Y:= g(X) and ∂ (X) ∂Xc,d =X l,n ∂ (X)) ∂Yl,n ∂Yl,n ∂Xc,d (8) In a way, ea ing he ou pu as a eplacemen o he unc ion. In he ollowing sec ions, i will be common o name he inpu o a speci ic block inside he model as Xand o i o be a eplacemen o a unc ion o he model’s inpu . E.g. X:= (Inpu ) whe e Inpu is he model’s inpu . Since ma ix ope a ions a e al eady implemen ed in a e y op imized way on mos machine lea ning lib a ies, we will se ou o simpli y he exp ession o g adien s and comp ess i as much as possible as e icien ma ix ope a ions. Fo ha pu pose, we will use he ollowing echniques: 1. Undo do -p oduc PkAa,kBk,b = (AB)a,b 2. Apply K onecke ’s del a Pn k=1 δk,c (k) = (c) when n≥c. 3. Spli a summa ion’s inne sum. X x (x)(g(x) + h(x)) = X x (x)g(x) + X x (x)h(x) 4. Simpli y in o comp ehensible ma ix ope a ions using: (a) Elemen -wise p oduc : (A⊙B)a,b =Aa,bBa,b (b) Sum o columns: SumCols(A)a=PkAa,k (c) Sum o columns: SumRows(A)b=PkAk,b (d) Cus om p oduc : MulCols(A, b)c,d =Ac,dbc 5 3 T ans o me S uc u e In o de o ansla e a sen ence o a iable leng h, models used o be ecu si e, equi ing as one execu ion pe sen ence, bu equi ing making many ecu si e s ages wi hin he model. This caused he backp opaga ion o he model o equi e aking in o accoun la ge amoun s o middle esul s, which made hese models ha e highe ime complexi y han desi ed. A T ans o me pe o ms ansla ion o sen ences by using he sou ce sen ence and he cu en ansla ion in o de o decide he nex wo d o oken o he ansla ion. T ansla ion equi es he execu ion o he model as many imes as he numbe o okens o he inal ansla ion. This migh seem ha i makes he model jus as ecu si e, bu i is no . The di e ence is in he cos o pe o ming he op imiza ion s ep. Wi h a ans o me , we can ain he model o ansla e co ec ly each a ge oken wi h single non- ecu si e backp opaga ions o g adien s, which is a sho e han he backp opaga ion o an equi alen ecu si e model. The model, as seen in igu e 1, has h ee dis inc pa s, he inpu p ocessing, he encode s and decode s. The wo sen ences he model ecei es, he sou ce and he a ge , a e ini ially exp essed wi h a ec o o okens, in ege s ep esen ing he wo ds o pa ial wo ds ha make hem up. In o de o wo k wi h hem, i s we ec o ize each oken using he embeddings ha hold maps om oken o ec o , and hen sends he ec o ized sou ce and a ge sen ences o, espec i ely, he encode and decode . The encode s a e a se ies o modules made up o inne laye s ha a e mean o ‘encode’ all he meaning o he okens and he ela ionships be ween hem and send he esul o he encode . The decode , howe e , will ge bo h he esul o he encode s and he ec o - ized a ge and will use ela ionship be ween wo ds wi hin he a ge and wi h he sou ce sou ce in o de o c ea e an ou pu . Wha makes he T ans o me ind ela ionships be ween posi ions so well a e he Mul i-Head A en ion modules (see 4.4-4.5). They a e an inno a i e way o ha e he model ocus on he ela ionship be ween speci ic posi ions o he sen ences. 6 ∂σ(ˆ Y)l,k ∂ˆ Ym, =δm,lσ(ˆ Y)l,kδn, −σ(ˆ Y)l,  =δm,lIl≤kσ(ˆ Y)l,kδk, −Il≤ σ(ˆ Y)l,  (27) And hus, o i o no be ze o i mus be ue ha m=l,l≤kand ei he : 1. k= , in which case m=l≤k= which con adic s m > . 2. l≤ , in which case m=l≤ which con adic s m> . And bo h con adic m> . Lemma 4. I is ze o i any o he ollowing is sa is ied: 1. l > n 2. m> 3. m=l 4.4 Scaled Do -P oduc A en ion The key aspec ha allowed he T ans o me o ise in enown was he inno a i e way hey used o ha e he model p ocess he ela ion be ween wo ds be ween wo gi en sen ences. The Scaled Do -P oduc A en ion, i s ob ains a ma ix o a en ion be- ween wo ds by pe o ming a scala p oduc be ween all pai s o wo d- ec o s o he gi en ‘que ies’ and ‘keys’ h ough a ma ix p oduc and applies i o he ‘ alues’ wi h ye again o he ma ix mul iplica ion. In he cases whe e we wish o p e en o wa d p opaga ion o in o ma ion wi hin a sen ence, he eby o cing he model o ain only depending on p e ious wo ds wi hin he sen ence, we will use he mask al eady explained in he sec ion abo e. 4.4.1 Fo wa d p opaga ion. A en(Q, K, V ) = σ(QKT √dk +Mask)V(28) whe e σis he so max unc ion and Mask can be op ional. In cases whe e i is no used, i will jus be he Ze o ma ix. 4.4.2 Backp opaga ion. To compu e he g adien o he loss w. . he in- pu s, we will wo k ou way backwa ds om he ou pu o he inpu ope a ion by ope a ion h ough he scaled do -p oduc . Fo cla i y, le us de ine some a iables which should be use ul la e : Y:= QKT √dk (29) U:= σ(Y) (30) H:= A en(Q, K, V ) = UV (31) 13 Fig. 2: Scaled Do -P oduc A en ion Fig. 3: Mul i-Headed A en ion The g adien s w. . Uand Va e ollow he same pa e n shown in 4.1. Fo he i s , we will go mo e in-dep h as an example. ∂Loss(Inpu ) ∂Uc,d =X l,n ∂Loss(Inpu ) ∂Hl,n ∂Hl,n ∂Uc,d (1) =X l,n ∂Loss(Inpu ) ∂Hl,n ∂PkUl,kVk,n ∂Uc,d =X l,n ∂Loss(Inpu ) ∂Hl,n X k Vk,n ∂Ul,k ∂Uc,d =X l,n ∂Loss(Inpu ) ∂Hl,n X k Vk,nδl,cδk,d (2) =X n ∂Loss(Inpu ) ∂Hc,n Vd,n (1) =∂Loss(Inpu ) ∂H VTc,d (32) 14 ∂Loss(Inpu ) ∂Vc,d =X l,n ∂Loss(Inpu ) ∂Hl,n ∂(PkUl,kVk,n) ∂Vc,d (1)(2) =UT∂Loss(Inpu ) ∂H c,d (33) The nex g adien o he loss, he one w. . Yis mo e complex. ∂Loss(Inpu ) ∂Yc,d =X l,n ∂Loss(Inpu ) ∂Ul,n ∂σ(Y)l,n ∂Yc,d (Eq.21) =X l,n ∂Loss(Inpu ) ∂Ul,n δc,lσ(Y)l,nδn,d −σ(Y)l,d (2) =X n ∂Loss(Inpu ) ∂Uc,n σ(Y)c,nδn,d −σ(Y)c,d (3) =X n ∂Loss(Inpu ) ∂Uc,n σ(Y)c,nδn,d −X n ∂Loss(Inpu ) ∂Uc,n σ(Y)c,nσ(Y)c,d (a) = (∂Loss(Inpu ) ∂U ⊙σ(Y))c,d −X n (∂Loss(Inpu ) ∂U ⊙σ(Y))c,nσ(Y)c,d (c) = (∂Loss(Inpu ) ∂U ⊙σ(Y))c,d −σ(Y)c,d ·SumCols(∂Loss(Inpu ) ∂U ⊙σ(Y))c (d) =∂Loss(Inpu ) ∂U ⊙σ(Y)) −MulCols(σ(Y),SumCols(∂Loss(Inpu ) ∂U ⊙σ(Y)))c,d (34) 15 The las wo, once again, ollow he same pa e n om sec ion 4.1. ∂Loss(Inpu ) ∂Qc,d =X l,n ∂Loss(Inpu ) ∂Yl,n ∂Yl,n ∂Qc,d =1 √dkX l,n ∂Loss(Inpu ) ∂Yl,n ∂(PkQl,kKn,k) ∂Qc,d =1 √dkX l,n ∂Loss(Inpu ) ∂Yl,n X k δc,lδd,kKn,k (2) =1 √dkX n ∂Loss(Inpu ) ∂Yc,n Kn,d (1) =1 √dk (∂Loss(Inpu ) ∂Y K)c,d (35) ∂Loss(Inpu ) ∂Kc,d =X l,n ∂Loss(Inpu ) ∂Yl,n ∂Yl,n ∂Kc,d =1 √dkX l,n ∂Loss(Inpu ) ∂Yl,n ∂(PkQl,kKn,k) ∂Kc,d =1 √dkX l,n ∂Loss(Inpu ) ∂Yl,n X k Ql,kδc,nδd,k 2 =1 √dkX lQl,d ∂Loss(Inpu ) ∂Yl,c  1 =1 √dk(∂Loss(Inpu ) ∂Y )TQc,d (36) This es ablishes a way o compu e he g adien quickly, as shown in igu e 4. 4.4.3 Masking case al e a ions. When he scaled do -p oduc is masked as discussed in sec ion 4.3, Uis a iangula ma ix, and as such, all he ope a ions ha include i o an elemen -wise p oduc can be hal ed in cos doing he necessa y changes o igno e he ze o posi ions. 4.4.4 E ec o epea ed inpu s. One ques ion ha will a ise is wha o do when some o e e y inpu , i.e. Q, K and Va e one and he same. In his case, we can ea hem as ou pu s o a elemen -wise unc ion o X ha is ei he he iden i y o he ze o unc ion. I.e. o he case whe e K=V=Q: Qc,d := q(Xc,d) = 0 (37) Kc,d := k(Xc,d) = Xc,d (38) Vc,d := (Xc,d) = Xc,d (39) 16 We can use he lemma 1 o compu e he g adien . ∂Loss(Inpu ) ∂Xc,d =∂Loss(Inpu ) ∂Qc,d ∂Qc,d ∂Xc,d +∂Loss(Inpu ) ∂Kc,d ∂Kc,d ∂Xc,d +∂Loss(Inpu ) ∂Vc,d ∂Vc,d ∂Xc,d = 0 + ∂Loss(Inpu ) ∂Kc,d +∂Loss(Inpu ) ∂Vc,d (40) Thus, he g adien o an inpu is he sum o all inpu s i is equal o. We will see in sec ion 4.5 ha i he linea applica ions o each head inside he MHA a e no he same, which equi es linked weigh s. Since his is no common in he implemen a ion o T ans o me s, we will wo k unde he assump ion ha hey a e di e en in each head o he MHA. Table 2: G adien o mulas o he Scaled Do -P oduc . W.R.T. G adien Fo mula U∂Loss(Inpu ) ∂H VT V UT∂Loss(Inpu ) ∂H Y∂Loss(Inpu ) ∂U ⊙σ(Y)−MulColsσ(Y),SumCols(∂Loss(Inpu ) ∂U ⊙σ(Y)) Q1 √dk ∂Loss(Inpu ) ∂Y K K1 √dk (∂Loss(Inpu ) ∂Y )TQ No e. X∈RN×din , W ∈Rdin×dou , b ∈R1×dou , L ∈RN×dou 17 Fig. 4: O de o ope a ions g adien s w. . he inpu s 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) Mul Cols O(LM) ∂Loss(X) ∂Y O(LM) ∂Loss(X) ∂K O(dkLM)∂Loss(X) ∂Q O(dkLM) 4.5 Mul i-Headed A en ion The mul i-head a en ion unc ion is he co e module inside a ans o me . I is also he mos cos ly pa o he model. The inpu s o i a e h ee ma ices Q, K and V. These a e he keys and alues and que ies. In a ans o me , we always ha e ha he inpu Kand Va e iden ically he same inpu , his can be obse ed in 1, we ep esen ha wi h K≡V. As seen in igu es 2 and 1, in mos cases all h ee coincide, bu in he decode , in he second Mul i-Head A en ion hey di e . Because o he di e ences in dimensions depending on whe e he MHA is loca ed we will wo k wi h K, V ∈RL×dmand Q∈RM×dmwhe e N, M ∈ {S, T }. 4.5.1 Fo wa d p opaga ion. I is mul i-head, because inside i , hheads pe - o m independen ly “Scaled Do -P oduc A en ion” o hei inpu s (see igu e 3). Since he e will no be cases o ma ix powe s, we will use supe indices o 18 di e en ia e be ween ma ices associa ed wi h each head, i.e. Qiwould be he inpu Q(see igu e 2) ecei ed by he head i. The ou pu o he mul i-head a en ion is a linea applica ion o he con- ca ena ion o he ou pu s o e e y head, which we will call Hi∈Rdk×T o he ou pu o head i. M:= Mul iHeadA en(Q, K, V ) = Linea (Conca (H1, . . . , Hh), Wo, bo) (41) whe e: Conca (H1, . . . , Hh) :=      H1 . . . Hh      (42) To abb e ia e, we will use C o ep esen Conca (H1, . . . , Hh) om his poin onwa ds. We will call he inpu s o head ias Qi, Ki, V i. Each o hese, linea appli- ca ions o he o iginal Q, K, V . This dec eases he dimensions o he inpu s o L×dkand M×dk espec i ely. Qi= Linea (Q, Wi q, bq) (43) Ki= Linea (K, W i k, bk) (44) Vi= Linea (V, W i , b ) (45) Using hese inpu s, each heads pe o ms he “Scaled Do -P oduc A en ion”: Hi:= A en(Qi, Ki, V i) A de ail ha is ele an o he implemen a ion o he o wa d pass is he possibili y o implemen ing he hlinea p ojec ions associa ed wi h each inpu a once. Lemma 5. Comp ession o linea applica ions. Linea (X,      W1 . . . Wh      ,(b1···bh)) = Linea (X, W1, b1),··· ,Linea (X, Wh, bh) (46) whe e X∈RN×din , W i∈Rdou ×din and bi∈R1×dou . 19 P oo . Le n′=i·dou +n, hen: Linea (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 = Linea (X, Wh, bh)l,n (49) hus ob aining he desi ed ma ices in one do p oduc i desi ed. This is used in he implemen a ion o he o wa d pass ??. 4.5.2 Backp opaga ion. As always, we assume ha we ha e ∂Loss(Inpu ) ∂M and we wan o ha e ∂Loss(Inpu ) ∂X o X∈ {Q, K, V }. To do ha , we will use backp opaga ion inside he MHA. Due o he de ini ion o M, we can use he o - mulas o sec ion 4.1 o compu e ∂Loss(Inpu ) ∂C . Now, we would equi e he g adien w. . each head. Using he de ini ion o C, i is i ial ha i equi es undoing he conca ena ion on he g adien w. . C: ∂Loss(Inpu ) ∂C =     ∂Loss(Inpu ) ∂H1 . . . ∂Loss(Inpu ) ∂Hh      (50) Now, using he equa ions o sec ion 4.4, we can ob ain ∂Loss(Inpu ) ∂Xi o Xi∈ {Qi, Ki, V i}. And, once again, using sec ion 4.1 we can compu e he g adien ∂Loss(Inpu ) ∂X om ∂Loss(Inpu ) ∂Xi. ∂Loss(Inpu ) ∂Xc,d =X i ∂Loss(Inpu ) ∂Xi l,n ∂Xi l,n ∂Xc,d (51) =X iX l,n ∂Loss(Inpu ) ∂Xi l,n ∂Xi l,n ∂Xc,d (52) Eq.12 =X i (∂Loss(Inpu ) ∂XiWi)c,d (53) 20 Thus being he sum o he g adien s w. . each linea applica ion’s inpu . Table 3: G adien o mulas speci ic o Mul i-Head A en ion. W.R.T. Complexi y G adien Fo mula HiO(1) Rows [(i−1)dk, i dk] o ∂Loss(Inpu ) ∂C Xi∈ {Qi, Ki, V i} O(hdmdk·max(M, L)) Pi ∂Loss(Inpu ) ∂XiWi No e 1. Re e o Linea Applica ion and Scaled Do -P oduc o middle-poin g adien s. No e 2. Xi∈Rdm×M∪Rdm×L, W ∈Rdm×dk, C ∈Rhdk×T 4.6 Laye No maliza ion Laye no maliza ion, no o be mis aken o ba ch no maliza ion, no malizes ac oss he posi ion ea u es ins ead o he samples. Le us assume ha he inpu is in he shape (B, N, E) whe e Bis he ba ch size, N he numbe o posi ions, and Eis he numbe o ea u es pe posi ion. The di e ence would be ha while Ba ch No maliza ion would no malize B alues o (N, E) andom a iables, Laye No maliza ion would no malize E alues o (B, N) andom a iables. Du ing he ollowing compu a ions, due o all ope a ions being done i e a ing ac oss E, we can, wi hou loss o gene aliza ion, ea he inpu as ha ing shape (B·N, E) o simpli y he esul ing exp essions. No e ha in a T ans o me , E always coincides wi h dmodel. 4.6.1 Fo wa d p opaga ion. The ma hema ical exp ession o he o wa d pass o weigh ed laye no maliza ion is: Yi,j =ˆ Yi,jγj+βj=Xi,j −µi pσ2 i+εγj+βj(54) Whe e: –X∈RB·N×Eis he inpu , –ˆ Y∈RB·N×Eis he no malized inpu , –γ,  β∈REa e he weigh s, –µiand σ2 ia e he mean and a iance o he ow ec o Xi, –and ε > 0 will be a small quan i y used in o de o a oid cases wi h ze o a iance. 21 4.6.2 Backp opaga ion We will i s compu e he g adien o he e o w. . he inpu . The exp ession, will be as ollows: ∂Loss(Inpu ) ∂Xc,d =X i,j ∂Loss(Inpu ) ∂Yi,j ∂Yi,j ∂Xc,d (55) whe e ∂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, o we will equi e he de i a i e o he nume a o and denomina o o he no maliza ion w. . . he inpu . The i s is s aigh o wa d. ∂(Xi,j −µi) ∂Xc,d =δi,c(δj,d −1 E) (57) Fo he denomina o , we will do he de i a i e in wo s eps. Fi s , he de i a- i e o he a iance w. . . he inpu . ∂σ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) Wi h his, he de i a i e o he denomina o o he no maliza ion w. . . he inpu 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 Le us ob ain wha he ou pu Hwill be compa ed o H′. QKT=  Q′ q KT=  Q′KT qKT (74) The ma ix 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, he ou pu o he head mus be: H=  U′ σ(1 dkqKT) V=  H′ σ(1 dkqKT)V (77) Thus, he ou pu only changes in he las ow. This means ha he new oken does no a ec he ou pu associa ed wi h he p e ious okens. In he p e ious case, he inpu Qi, which we eplaced wi h Q, was he only ma ix which changed in size. Le us see he di e ence in he ou pu s o he i s MHA o he decode when adding an addi ional oken. We will use he same s a egy o dec ease supe indices and we will wo k wi h he ou pu o head i. He e, he head’s h ee inpu s Q, K, V a e posi ion- wise linea applica ions o he MHA inpu s, which a e all one and he same (see igu e 1). Because o ha , we ha e ha he inpu s o he head iwill be: Q=  Q′ q =  Linea (X′, Wi q, bi q) Linea (x, Wi q, bi q) (78) and simila ly, omi ing he linea applica ions o X: K=  K′ k and V=  V′  (79) whe e q,k, ∈R1×dk. The deno a ion o wha he ou pu o he head would be using p e ious ma ices will be deno ed as: H′=U′V′=σ(Y′)V′=σ(Q′(K′)T √dk +Mask)V′(80) 29 Le us ob ain wha he ou pu Hwill be compa ed o H′. QKT=  Q′ q   K′ k  T =  Q′(K′)TQ′kT q(K′)TqkT (81) And hus, (QKT √dk +Mask) = 1 √dk  Q′(K′)T+Mask′Q′kT −∞ qkT (82) The ma ix Ywill be Y=  Y′1 √dkQ′kT −∞ 1 √dkqkT (83) Lemma 6. I y∈Rnand y∈R hen, 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) hen he so max o he ec o ysa is ies: σ(y) = hσ(y)α βeyi=hσ(y′)α σ(y)n+1i(87) I.e. he i s nposi ions a e he scaled so max o y’. Using his lemma, using he lemma on he ows o Y, and no ing ha U′ k=σ(Y′ k) (88) hen, o j:= exp( 1 √dkQ′ jkT) and αj:= Pp i=1 exp Y′ j,i Pp i=1 exp Y′ j,i + j (89) βj:= 1 Pp i=1 exp Y′ j,i + j (90) we ha e: U=         (U′)1α1β1exp( 1 √dkQ′ 1kT) . . .. . . (U′)pαpβpexp( 1 √dkQ′ pkT) 01         (91) 30 And hus, he ou pu is: A en(Q, K, V ) =         (U′)1α1β1exp( 1 √dkQ′ 1kT) . . .. . . (U′)pαpβpexp( 1 √dkQ′ pkT) 01           V′   =         U′ 1V′α1+β1exp( 1 √dkQ′ 1kT) . . . U′ pV′αp+βpexp( 1 √dkQ′ pkT)         =         H′ 1α1+β1exp( 1 √dkQ′ 1kT) . . . H′ pαp+βpexp( 1 √dkQ′ pkT)         (92) Then he dis ance o ow Hj om he ow H′ jis ∥H′ j−(H′ jαj+βj j )∥ =∥H′ j(1 −αj)−βj j ∥(93) =∥H′ j j Pp i=1 exp Y′ j,i + j− j Pp i=1 exp Y′ j,i + j ∥(94) = j Pp i=1 exp Y′ j,i + j∥H′ j− ∥(95) ≤ ∥H′ j− ∥(96) In conclusion, he dis ance be ween he ows om ma ix H′and he i s p ows o His a mos he dis ance o he ec o and depends on j,1≤j≤p. Due o ha , he i s p ows o he ou pu when he inpu o he decode is p+ 1 ows will mo e away om he ou pu i p ows had been used. In p ac ical e ms, his means ha op imizing ou pu o he ansla ion o he (p+ 1)- h oken does no necessa ily imply he op imiza ion o he model o he case wi h p okens. The decode , howe e , uses he Add & No m modules, and his addi ion will dec ease he dis ance be ween inpu and ou pu o he MHAs. As we ha e seen, he place whe e he added (p+ 1)- h ow a ec s he gene al ou pu o he model is in he sel -a en ion o he decode . The e, an addi ion module sums he inpu and ou pu o he sel -a en ion, se ing as a limi e on how changed he ou pu o he i s p ows can be due o he added ow. This inc eases posi ion 31 independence and can help in achie ing be e ansla ions o he i s p okens when op imizing he ou pu o he ansla ion o he (p+ 1)- h oken when op imizing all he ou pu s o he model and no only he las . In conclusion, lea ning o ansla e one oken o he ansla ion may help o ain all p e ious okens o he ansla ion, which would accele a e lea ning. 6 Da ase s and da a p ocessing In o de o benchma k he model’s ue pe o mance, we se ou o make use o eal da ase s used o amous ansla ion asks and we p ocessed i in o de o use i o es ima e he a e age ime he model akes in each laye aking in o accoun he oscilla ion be ween sen ence leng hs and ba ch sizes ha a e na u al in ansla ion aining. 6.1 Da ase selec ion. In NLP asks, inding a la ge numbe o alid samples be a di icul ask due o he la ge amoun o ypos ound in human-p oduced ex . Because o his, he da ase s used o ain models end o be ew in numbe and wi h sen ences a ying wildly in leng h and complexi y. Some o he da ase s mo e used o his a e he WMT da ase s. We used he WMT 2014 English o Ge man da ase o ou measu emen s. 6.2 Tokeniza ion. The o ma o he da ase s a e usually plain ex con aining he sen ences, and w i en in an speci ic language. Ou model uses as inpu s in ege s ep esen ing he sen ence in okens, hus, some da a p ocessing has o ake place. The p ocess o ans o ming a ex in o a se ies o okens, each ha ing a unique assigned in ege , is called okeniza ion. One o he possible ways o okenize he ex is o ha e e e y wo d become one oken and assign each as a di e en in ege . This has been ound o be sub- op imal, because i does no co ec ly cap u e he simila i ies be ween wo ds. This leads o he use o sub-wo d okeniza ion. The me hod we will use will be By e-Pai Encoding. This algo i hm inds commonly ound g oups o le e s ha end o appea oge he and assigns hem each a oken. This leads o he endency o assigning p e ixes, su ixes and wo d oo s as unique okens. Because we a e wo king p ima y wi h a model used o ansla ion asks, we ha e wo op ions a hand. We can c ea e wo di e en okeniza ions, which would cause he ocabula y size o he o iginal and a ge languages be asymme ical, and ela i ely small, o me ge all wo ds in a single co pus and ha e one la ge ocabula y o okens. The i s op ion allows us o keep he embeddings ela i ely small in ou model. The second, howe e , allows ou model he possibili y o using he same embedding in he encode and decode , and, i he languages a e closely ela ed, 32 could lead o u he comp ession and a o al embedding size smalle han wo di e en ones. In ou case, we selec ed he second app oach. Due o his algo i hm being ou o he scope o his wo k, we chose o use an exis ing implemen a ion o he algo i hm o he okeniza ion o ou da ase . 2 6.3 Da ase o ma . Once he da ase has been okenized, we un in o he p oblem o da ase size and eading. Depending on he me hod we use o w i e he da ase in o a ile, he eading could become a bo leneck du ing ou aining. The p ima y issue we had o deal wi h was i egula sample leng h. I each sample has a di e en leng h, shu ling and picking andom samples would cause a la ge eading ime o e head. To sol e his, we decided o s o e each sample padded o each he same leng h and s o ing he o iginal leng h (see 7). Ano he issue was which ile o ma o use. Due o he BPE implemen a ion being in Py hon, we decided o use he NumPy ile o ma and add suppo NumPy ile eading in C o he p ojec . Table 7: Da ase s uc u e. S T s1s2··· sN−1sN 1 2··· N−1 N Sou ce leng h Ta ge leng h Padded sou ce Padded a ge 6.4 Da aloade In NLP asks, samples can a y in size anywhe e be ween single wo ds o en i e pa ag aphs. Because o his, a ba ch o andomly picked samples o a da ase can po en ially ha e widely di e en sen ence leng hs. In o de o keep he ba ches samples’ leng hs simila , we used a bucke da aloade . Fi s ly, we s o ed he samples in bucke s o simila amoun s o okens. Using his, he da aloade can shu le each bucke in e nally and p io i ize picking samples o he same bucke o each ba ch. Using his me hod, we can assu e ha all samples wi hin he same ba ch ha e simila leng h and we can ha e each ba ch a e age he same amoun o okens. Howe e , in o de o ake ad an age o he ull esou ces a ou disposal du ing aining, when a bucke lacks enough samples o comple e a ba ch, we will comple e i wi h samples wi h om simila bucke s. One las hing o no e is he eading me hod used o ba ch loading. The e a e wo main op ions when i comes o loading da a, one is o mo e he en i e da ase in o RAM memo y, making eading as e , and he o he is o ead each ba ch 2Mo e in o ma ion can be ound in [BPE(1999)]. 33 om disk as needed, dec easing memo y use. In ou case, because ansla ion da ase s can be e y la ge and he ime equi ed o ead a ba ch om memo y is a he small, we will ead each ba ch om disk di ec ly. 7 Implemen a ion 7.1 HELENNA He e ogeneous Lea ning Neu al Ne wo k Applica ion o HELENNA, is a neu- al ne wo k amewo k w i en p ima ily in C which allows o de ine and use neu al ne wo k models wi h he pu pose o eaching and esea ching compu e a chi ec u e concep s. We chose HELENNA o implemen T ans o me s because he pu pose o his pape is esea ching he op ions o implemen a ion and hei pe o mance o deep-lea ning models, which closely ma ches he pu pose o HELENNA, and seemed o al eady ha e he majo i y o he ools necessa y o implemen a T ans- o me , which would allow mo e ime o ine uning op imiza ions speci ically o T ans o me s and he explo a ion o implemen a ions in o he pla o ms. 7.2 Limi a ions o HELENNA HELENNA was o iginally implemen ed wi h a ocus on suppo ing ANNs. The i s p oblem ha a ises om his is ha i caused hei implemen a ions o p ima ily expec linea inpu s o many laye s, such as he Linea and So max laye s. Which leads o no ha ing a clea suppo o mul idimensional sam- ples. This o ced he implemen a ion o new laye s and inc eased he wo kload signi ican ly. The main objec i e was he compa ison o implemen a ions mixing o he pla o ms (GPUs, FPGAs,. . . ) and echnologies such as MKL. Bu hese a enues could no be explo ed a his ime, as ba ely any pa o he T ans o me did could be execu ed in HELENNA wi hou hea y modi ica ions and his inc eased d ama ically he expec ed amoun o wo k. The second issue was ha he s anda d memo y layou used o s o age o middle esul s, inpu s and ou pu s has he sample size as las dimension. Wha would his en ail? In ou case, he inpu o he model would be (S, B) o he encode and (T, B) o he decode . In each o hem, he embedding would use as ou pu layou (N, dmodel, B) and his would emain he layou o all MHA, Laye No maliza ions and Feed Fo wa d laye s. The issue wi h his is ha on he ope a ions equi ed on mos laye s, he main i e a ion occu s o e he dmodel dimension, which is no con iguous in memo y. As a esul , we can p edic ha an implemen a ion using his memo y layou s a egy will wo sen he pe o mance. Du ing pa ial es s, his was p o en, which led o he a new implemen a ion using dmodel as ou las dimension and see how much pe o mance we can ob ain om a CPU using HELENNA. This, again, equi ed u he wo k and emo ed he possibili y o implemen ing he model on FPGAs o GPUs. 34 7.3 T ans o me suppo Among he laye s needed o implemen a T ans o me ha we e al eady imple- men ed in HELENNA, we ha e he ReLU, Addi ion and D opou laye s. In addi ion o hese, e en hough we will equi e hea y modi ica ions a imes, we could e-pu pose: – he exis ing Linea laye o c ea e a Posi ion-wise Linea laye , – he So max laye in o a Posi ion-Wise So max laye , –and he Ba ch No maliza ion laye in o a Laye No maliza ion laye . Wi h his, he only missing laye s ha would ha e o be buil om he g ound up would be: – he MHA laye , – he Embedding laye , –and he Posi ional Encoding laye . Se e al possible a ia ions o he implemen a ions will be es ed in an a - emp o ully u ilize he pa alleliza ion capabili ies h ough he use o he li- b a y OpenMP. Due o his being mainly changes in he i e a ion o de s and he use o auxilia a iables, we will no be del ing in o he op ions es ed. 3 E en hough one o he main goals was he implemen a ion o he T ans- o me ’s laye s in di e en pla o ms wi hin HELENNA, such as GPUs o using MKL. The amoun o unimplemen ed backbone unc ionali ies and he obscu i y o he lib a ies used caused us o limi he implemen a ion o CPU only. 7.4 Fo wa d p opaga ion Due o he way HELENNA alloca es he all he memo y ha will be equi ed a he beginning, all laye s a e ini ialized o he maximum ba ch size and sen ence leng h possible du ing he aining. Howe e , each ba ch will ha e di e ing size and dimensions. Because o his, on each o wa d p opaga ion, a he s a , a new unc ionali y had o be added o suppo a cons an ly changing he dimensions o all inpu s and ou pu s o laye s o accu a ely send in o ma ion be ween laye s. 7.5 Backp opaga ion In HELENNA, c oss-en opy loss is al eady implemen ed and unc ions well wi h hei so max laye . Bu , as we commen ed p e iously, HELENNA does no suppo so max ac oss one dimension. Because o his, a new implemen a ion o c oss-en opy loss and hei g adien had o be implemen ed by modi ying exis ing unc ions. 3Should he eade be in e es ed in op ions explo ed, some examples a e a ailable in he code ha a e hidden h ough he use o de ined global cons an s in he cpu.c ile. 35 8 Resul s 8.1 Compa ison be ween HELENNA and PyTo ch In o de o compa e he pe o mance o he implemen a ions o each laye bo h in he backp opaga ion and o wa d pass, we will p o ile each laye independen ly, in o de o compa e each laye p ecisely wi hou unwan ed in e ac ions wi h o he unning compu a ion. As we can compa e be ween he ables 8 and 9, he implemen a ion o each laye using HELENNA’s memo y layou and OMP o pa allelize ope a ions is a slowe han he p o essionally c ea ed implemen a ions ound in PyTo ch. Due o he ac ha we can see a se e e inc ease in ime o e e y laye wi h la ge inpu s, including a simple addi ion om 0.02 ms o 0.47 ms, i leads o he conclusion ha a a subop imal h ead managemen and memo y layou /lookup s a egy is is used in HELENNA, which is independen o how op imal he complexi y o he ope a ions a e supposed o be in heo y. This is u he co obo a ed by he me ics o he Embedding laye . This laye equi es o look, o ew inpu s, o i e a e o e posi ions in memo y which a e con iguous in memo y. This leads o a be e use o each h ead and ewe cache misses and achie es simila pe o mance o PyTo ch’s (0.05 ms compa ed o 0.13 ms). Table 8: Time me ics o HELENNA’s implemen a ion wi h NBE memo y layou . Laye Fo wa d Backwa d Time/call Time sha e Time/call Time sha e MHA 21.08 ms 6.53% 92.80 ms 28.74% Linea 24.46 ms 10.52% 101.64 ms 43.72% D opou 0.53 ms 0.29% 0.03 ms 0.02% Laye No maliza ion 2.88 ms 2.99% 0.24 ms 0.25% Embedding 0.05 ms 0.00% 2.04 ms 0.05% Posi ional Encoding 0.04 ms 0.00% 0.02 ms 0.00% So max 1.53 ms 0.01% 0.45 ms 0.01% Addi ion 0.47 ms 0.01% 1.14 ms 0.01% No e. Me ics using he T ans o me p esen ed in [AIAYN(2017)] wi h ba ches con aining 150 −250 okens. 36 Table 9: Time me ics o PyTo ch’s implemen a ion Laye Fo wa d Backwa d Time/call Time sha e Time/call Time sha e MHA 1.47 ms 18.78% 2.54 ms 33.73% Linea 1.30 ms 12.20% 2.22 ms 13.44% D opou 0.05 ms 2.06% 0.08 ms 3.65% Laye No maliza ion 0.05 ms 1.04% 0.08 ms 1.94% Embedding 0.13 ms 0.14% 4.23 ms 4.68% Posi ional Encoding 0.04 ms 0.15% - - So max 0.94 ms 0.32% - - Addi ion 0.02 ms 0.58% - - No e 1. Addi ion, Posi ional Encoding and So max do no allow s and-alone, backp opaga ion as hey a e mean o be op imized when connec ed o o he laye s. No e 2. Me ics using he T ans o me used in [AIAYN(2017)] and ba ches con aining 150 −250 okens. 8.2 Compa ison o memo y layou s in HELENNA As obse ed in he p e ious sec ion, he cu en bo leneck in he execu ion o a T ans o me is he ca e ul managemen o memo y and pa alleliza ion. In o de o inc ease he pe o mance o he implemen a ion in HELENNA, we se ou o c ea e an implemen a ion using he dimension o he ea u es, dmodel o E, as he las dimension in he memo y layou o inpu s and ou pu s o each laye . This allows us o u he ake ad an age o each h ead and manage memo y use, by changing he i e a ion s a egy o maximize he numbe o ope a ions done in pa allel on con iguous a eas o he memo y. As seen in able 10, he pe o mance changed g ea ly on e e y laye wi h a la ge sha e o he compu a ion ime, dec easing by be ween 30% and 60% he ime spen on each laye on he o wa d and backwa d passes. In o al, due o hese laye s appea ing epea edly h oughou he model, he o al aining ime was dec eased by 40%. In spi e o his, he pe o mance inc ease ob ained h ough he op imiza ion o he memo y usage, he ope a ions pe o med, and he pa alleliza ion s a egy h ough OpenMP, i is s ill a away om he pe o mance displayed by he p o essionally c ea ed PyTo ch lib a ies o ma ix ope a ions and au o-g adien compu a ion. 37 Table 10: Compa ison o HELENNA’s implemen a ions. O iginal New Pass Laye ms/call % ms/call % Time Dec eaase MHA 29.06 5.85 28.27 7.98 2.73 Linea 35.02 0.98 33.78 13.24 3.54 D opou 0.73 0.26 0.71 0.36 3.54 Laye No m. 4.10 1.38 1.69 0.29 58.90 Fo wa d Embedding 0.06 0.00 0.07 0.01 −12.58 Pos. Encoding 0.05 0.00 0.04 0.00 8.47 So max 2.34 0.03 1.10 0.02 52.97 Addi ion 0.02 0.01 0.02 0.01 4.72 MHA 170.59 34.32 100.04 28.27 41.36 Linea 153.32 42.82 105.97 41.54 30.88 D opou 0.04 0.01 0.04 0.02 −1.52 Laye No m. 0.62 0.21 0.61 0.29 1.93 Backwa d Embedding 1.44 0.03 1.48 0.05 −3.09 Pos. Encoding 0.02 0.00 0.02 0.00 −15.37 So max 0.64 0.03 0.34 0.01 47.12 Addi ion 0.04 0.01 0.04 0.02 −1.32 No e. We used he T ans o me om [AIAYN(2017)] wi h ba ches con aining 250 −450 okens. 9 Conclusions T ans o me s, jus like he g ea majo i y o machine lea ning models, depends g ea ly ope a ions ha equi e i e a ing o e ows o columns o la ge ma ices. Because o his, e en i he o mulas used o he o wa d and backwa d passes o he model had op imal heo e ical ime complexi y, wi hou ca e ul memo y layou planning he memo y and h ead scheduling, memo y look-ups would become he bo leneck o he model. In o de o u he op imize he model in HELENNA o each o su pass he pe o mance in displayed in PyTo ch, we would need o cus omize mo e closely bo h he da a memo y layou and he exac beha io o each h ead du ing 38