Qorayaasha:
(1) Albert Gu, Waaxda Barashada Mashiinka, Jaamacadda Carnegie Mellon iyo wax ku biirinta siman;
(2) Tri Dao, Waaxda Sayniska Kombiyuutarka, Jaamacadda Princeton iyo wax ku biirinta siman.
3 Qaababka Meelaynta Gobolka iyo 3.1 Dhiirigelin: Xulashada Macnaha Cadaadiska
3.2 Horumarinta SSM-yada oo leh Xulasho
3.3 Hirgelinta hufan ee xulashada SSM-yada
3.4 Dhismaha SSM la Fududeeyay
3.6 Faahfaahinta Qaabka Dheeraadka ah
4 Qiimaynta Dhabta ah iyo 4.1 Hawlaha Dabiiciga ah
4.4 Qaabaynta Maqalka iyo Jiilka
4.5 Xawaaraha iyo Tilmaamaha Xusuusta
Algorithm-ka Hardware-ka og ee SSM-yada Xulashada ah
E Faahfaahinta Tijaabada iyo Natiijooyinka Dheeraadka ah
Nashqadaha saaxiibtinimo ee qalabka sida convolutions (Krizhevsky, Sutskever, and Hinton 2012) iyo Transformers (Vaswani et al. 2017) waxay ku raaxaystaan codsi baahsan. Halkan waxaan hiigsaneynaa inaan ka dhigno SSM-yada xulashada ah kuwo ku hufan qalabka casriga ah (GPU) sidoo kale. Habka xulashada waa mid dabiici ah, oo hore ayaa la isku dayay in lagu daro kiisaska gaarka ah ee xulashada, sida u ogolaan ∆ kala duwanaansho waqti ka dib SSM-yada soo noqnoqda (Gu, Dao, et al. 2020). Si kastaba ha ahaatee, sidii hore loo sheegay xaddidaadda aasaasiga ah ee isticmaalka SSM-yada ayaa ah waxtarkooda xisaabinta, taas oo ahayd sababta S4 iyo dhammaan noocyada kala duwan ay u adeegsadeen moodooyinka LTI (aan la dooran), oo u badan qaabka isku-dhafan ee caalamiga ah.
3.3.1 Dhiirigelinta Modelladii Hore
Waxaan marka hore dib u eegeynaa dhiirigelintan oo aan dulmar ku samaynay habkeena aan uga gudubno xaddidaadaha habab hore.
• Marka loo eego heerka sare, moodooyinka soo noqnoqda sida SSM-yada had iyo jeer waxay isku dheelitiraan isdhaafsiga u dhexeeya muujinta iyo xawaaraha: sida looga hadlay Qaybta 3.1, moodooyinka leh cabbir heer qaran oo qarsoon waa inay noqdaan kuwo waxtar badan leh laakiin gaabis ah. Markaa waxaan rabnaa inaan sare u qaadno cabbirka qarsoon ee gobolka anagoon bixinin xawaaraha iyo kharashka xusuusta.
• Ogsoonow in habka soo noqnoqda uu ka dabacsan yahay habka convolution, maadaama kan dambe (3) uu ka yimid ballaarinta kii hore (2) (Gu, Goel, and Ré 2022; Gu, Johnson, Goel, et al. 2021). Si kastaba ha ahaatee, tani waxay u baahan doontaa xisaabinta iyo xaqiijinta xaaladda daahsoon ℎ oo leh qaab (B, L, D, N), aad uga weyn (dhinaca N, cabbirka gobolka SSM) marka loo eego gelinta x iyo soo-saarka y ee qaabka (B, L, D). Sidaas awgeed habka isku-dhafan ee waxtarka badan ayaa la soo bandhigay kaas oo dhaafi kara xisaabinta dawladda oo rumayn kara kernel convolution (3a) oo keliya (B, L, D).
• SSM-yadii hore ee LTI waxay ka faa'idaystaan labada foom ee soo noqnoqda-convolutional si ay u kordhiyaan cabbirka gobolka wax ku oolka ah ee Nx (≈ 10 - 100), aad uga weyn RNN-yada dhaqanka, iyada oo aan lahayn ganaaxyo waxtar leh.
3.3.2 Dulmarka Sawirka Xulashada: Balaadhinta Gobolka Hardware-Aware
Habka xulashada waxaa loogu talagalay in laga gudbo xaddidaadaha moodooyinka LTI; isla markaa, waxaan u baahanahay inaan dib u eegno dhibaatada xisaabinta ee SSM-yada. Waxaan tan ku wajahaynaa saddex farsamooyin qadiimi ah: fusion kernel, scan is barbar socda, iyo dib u xisaabinta. Waxaan samaynaa laba indho-indheyn oo waaweyn:
• Xisaabinta soo noqnoqota ee naive waxay isticmaashaa O(BLDN) FLOPs halka xisaabinta convolutional ay isticmaasho O(BLD log(L)) FLOPs, tan horena waxay leedahay arin joogto ah oo hoose. Sidaa darteed taxanaha dheer iyo cabbirka gobolka ee aan aad u weynayn, qaabka soo noqnoqda wuxuu runtii isticmaali karaa FLOPs ka yar.
Labada caqabadood waa dabeecadda isku xigxiga ee soo noqoshada, iyo isticmaalka xusuusta weyn. Si aan wax uga qabto kan dambe, si la mid ah habka iswaafajinta, waxaan isku dayi karnaa inaanaan dhab ahaan u hirgelin xaaladda buuxda ℎ.
Fikradda ugu weyn waa in laga faa'iidaysto sifooyinka dardar-geliyayaasha casriga ah (GPUs) si loo hirgeliyo gobolka kaliya heerar wax ku ool ah oo kala sareynta xusuusta ah. Gaar ahaan, hawlgallada intooda badan (marka laga reebo isku-dhufashada matrix) waxay ku xidhan yihiin xajmiga xajmiga xusuusta (Dao, Fu, Ermon, et al. 2022; Ivanov et al. 2021; Williams, Waterman, and Patterson 2009). Tan waxaa ka mid ah hawlgalkayaga iskaanka, waxaanan isticmaalnaa isku-dhafka kernel si loo yareeyo xaddiga xusuusta IO-yada, taasoo horseedaysa xawaare weyn marka la barbar dhigo hirgelinta caadiga ah.
Si looga fogaado soo noqnoqoshada isku xigxiga, waxaan aragnaa in inkasta oo aysan ahayn mid toosan, haddana waxaa lagu barbar dhigi karaa algorithm isbarbar-bar-ku-ool ah (Blelloch 1990; Martin and Cundy 2018; Smith, Warrington, and Linderman 2023).
Ugu dambeyntii, waa inaan sidoo kale ka fogaannaa badbaadinta dawladaha dhexe, kuwaas oo lagama maarmaan u ah dib-u-celinta. Waxaan si taxadar leh u isticmaalnaa farsamada caadiga ah ee xisaabinta si aan u yareyno shuruudaha xusuusta: dowladaha dhexe lama kaydiyo laakiin waxaa lagu xisaabiyaa kaarka dambe marka agabka laga soo raro HBM ilaa SRAM. Natiijo ahaan, lakabka iskaanka xulashada ee isku dhafan ayaa leh shuruudaha xusuusta la mid ah sida hirgelinta isbeddelka la hagaajiyey ee FlashAttention.
Faahfaahinta kernel-ka isku qasan iyo dib u xisaabinta waxay ku taal Lifaaqa D. Lakabka xulashada SSM oo buuxa iyo algorithm ayaa lagu muujiyey sawirka 1.
Warqadan waxaa laga heli karaa arxiv iyadoo la raacayo shatiga CC BY 4.0 DEED.