Princeton dhe CMU shtyjnë kufijtë e AI me modelin e sekuencës Mamba

Tabela e lidhjeve

Abstrakt dhe 1 Hyrje

2 Modele të Hapësirës Shtetërore

3 Modele selektive të hapësirës shtetërore dhe 3.1 Motivimi: Përzgjedhja si mjet ngjeshjeje

3.2 Përmirësimi i SSM-ve me përzgjedhje

3.3 Zbatimi efikas i SSM-ve selektive

3.4 Një arkitekturë e thjeshtuar SSM

3.5 Vetitë e Mekanizmave të Përzgjedhjes

3.6 Detaje shtesë të modelit

4 Vlerësim empirik dhe 4.1 Detyra sintetike

4.2 Modelimi i gjuhës

4.3 Modelimi i ADN-së

4.4 Modelimi dhe gjenerimi i audios

4.5 Standardet e shpejtësisë dhe kujtesës

4.6 Modeli i Ablations

5 Diskutim

6 Përfundime dhe Referenca

Një Diskutim: Mekanizmi i Përzgjedhjes

B Punë të ngjashme

C Mekanika e SSM-ve selektive

D Algoritmi i vetëdijshëm për harduerin për SSM-të selektive

E Detaje eksperimentale dhe rezultate shtesë

Abstrakt

Modelet e fondacionit, që tani fuqizojnë shumicën e aplikacioneve emocionuese në mësimin e thellë, bazohen pothuajse në mënyrë universale në arkitekturën e Transformer-it dhe modulin e tij kryesor të vëmendjes. Shumë arkitektura të kohës nën-kuadratike si vëmendja lineare, konvolucioni i mbyllur dhe modelet e përsëritura, dhe modelet e hapësirës së strukturuar të gjendjes (SSM) janë zhvilluar për të trajtuar joefikasitetin llogaritës të transformatorëve në sekuenca të gjata, por ato nuk kanë performuar aq mirë vëmendjen në modalitete të rëndësishme si p.sh. si gjuhë. Ne identifikojmë se një dobësi kryesore e modeleve të tilla është paaftësia e tyre për të kryer arsyetim të bazuar në përmbajtje dhe për të bërë disa përmirësime. Së pari, thjesht lënia e parametrave SSM të jenë funksione të hyrjes adreson dobësinë e tyre me modalitete diskrete, duke lejuar modelin të përhapë ose harrojë informacionin në mënyrë selektive përgjatë dimensionit të gjatësisë së sekuencës në varësi të shenjës aktuale. Së dyti, edhe pse ky ndryshim parandalon përdorimin e konvolucioneve efikase, ne projektojmë një algoritëm paralel të vetëdijshëm për harduerin në modalitetin e përsëritur. Ne i integrojmë këto SSM selektive në një arkitekturë të thjeshtuar të rrjetit nervor nga skaji në fund pa vëmendje apo edhe blloqe MLP (Mamba). Mamba gëzon konkluzion të shpejtë (5× xhiros më të lartë se Transformers) dhe shkallëzim linear në gjatësinë e sekuencës, dhe performanca e tij përmirësohet në të dhënat reale deri në sekuenca miliona gjatësi. Si një bazë e modelit të sekuencës së përgjithshme, Mamba arrin performancën më të fundit në disa modalitete si gjuha, audio dhe gjenomika. Sa i përket modelimit të gjuhës, modeli ynë Mamba-3B tejkalon transformatorët me të njëjtën madhësi dhe përputhet me Transformatorët dyfishin e madhësisë së tij, si në para-stërvitje ashtu edhe në vlerësimin e mëvonshëm.

1 Hyrje

Modelet e fondacionit (FM), ose modelet e mëdha të paratrajnuara në të dhëna masive, të përshtatura më pas për detyrat e rrjedhës së poshtme, janë shfaqur si një paradigmë efektive në mësimin modern të makinerive. Shtylla kurrizore e këtyre FM-ve janë shpesh modele sekuence, që veprojnë në sekuenca arbitrare të hyrjeve nga një gamë e gjerë fushash si gjuha, imazhet, fjalimi, audio, seritë kohore dhe gjenomika (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al. 2019; 2023, Sutskever, Vinyals dhe Quoc V Le 2014). Ndërsa ky koncept është agnostik për një zgjedhje të veçantë të arkitekturës së modelit, FM-të moderne bazohen kryesisht në një lloj modeli të vetëm sekuence: Transformer (Vaswani et al. 2017) dhe shtresa e tij kryesore e vëmendjes (Bahdanau, Cho dhe Bengio 2015). Efikasiteti i vetë-vëmendjes i atribuohet aftësisë së tij për të drejtuar informacionin në mënyrë të dendur brenda një dritareje konteksti, duke e lejuar atë të modelojë të dhëna komplekse. Megjithatë, kjo veti sjell të meta themelore: një pamundësi për të modeluar ndonjë gjë jashtë një dritareje të kufizuar, dhe shkallëzim kuadratik në lidhje me gjatësinë e dritares. Një grup i madh kërkimesh është shfaqur në variante më efikase të vëmendjes për të kapërcyer këto mangësi (Tay, Dehghani, Bahri, et al. 2022), por shpesh në kurriz të vetive që e bëjnë atë efektive. Deri më tani, asnjë nga këto variante nuk është treguar të jetë empirikisht efektiv në shkallë në të gjithë domenet.

Kohët e fundit, modelet e sekuencave të hapësirës së strukturuar të gjendjes (SSM) (Gu, Goel dhe Ré 2022; Gu, Johnson, Goel, et al. 2021) janë shfaqur si një klasë premtuese e arkitekturave për modelimin e sekuencave. Këto modele mund të interpretohen si një kombinim i rrjeteve nervore të përsëritura (RNN) dhe rrjeteve nervore konvolucionale (CNN), me frymëzim nga modelet klasike të hapësirës shtetërore (Kalman 1960). Kjo klasë modelesh mund të llogaritet në mënyrë shumë efikase ose si një përsëritje ose konvolucioni, me shkallëzim linear ose afërsisht linear në gjatësinë e sekuencës. Për më tepër, ata kanë mekanizma parimorë për modelimin e varësive me rreze të gjatë (Gu, Dao, et al. 2020) në modalitete të caktuara të të dhënave dhe kanë dominuar standarde të tilla si Arena me rreze të gjatë (Tay, Dehghani, Abnar, et al. 2021). Shumë shije të SSM-ve (Gu, Goel dhe Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu dhe Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith, Warrington dhe Linderman 2023) kanë qenë të suksesshëm në fusha që përfshin të dhëna të vazhdueshme të sinjalit si audio dhe vizion (Goel et al. 2022; Nguyen, Goel, et al. 2022; Saon, Gupta dhe Cui 2023). Megjithatë, ato kanë qenë më pak efektive në modelimin e të dhënave diskrete dhe të dendura me informacion, si teksti.

Ne propozojmë një klasë të re të modeleve të hapësirës selektive të gjendjes, që përmirëson punën e mëparshme në disa akse për të arritur fuqinë modeluese të transformatorëve ndërsa shkallëzohet në mënyrë lineare në gjatësinë e sekuencës.

Mekanizmi i përzgjedhjes. Së pari, ne identifikojmë një kufizim kyç të modeleve të mëparshme: aftësinë për të përzgjedhur në mënyrë efikase të dhënat në një mënyrë të varur nga inputet (p.sh. fokusimi ose injorimi i inputeve të veçanta). Duke u bazuar në intuitën e bazuar në detyra të rëndësishme sintetike si kokat selektive të kopjimit dhe induksionit, ne hartojmë një mekanizëm të thjeshtë përzgjedhjeje duke parametrizuar parametrat SSM bazuar në hyrjen. Kjo i lejon modelit të filtrojë informacione të parëndësishme dhe të mbajë mend informacionin përkatës për një kohë të pacaktuar.

Algoritmi i vetëdijshëm për harduerin. Ky ndryshim i thjeshtë paraqet një sfidë teknike për llogaritjen e modelit; në fakt, të gjitha modelet e mëparshme të SSM-ve duhet të jenë të pandryshueshme në kohë dhe në hyrje, në mënyrë që të jenë efikase llogaritëse. Ne e kapërcejmë këtë me një algoritëm të vetëdijshëm për harduerin që llogarit modelin në mënyrë periodike me një skanim në vend të konvolucionit, por nuk e materializon gjendjen e zgjeruar për të shmangur aksesin IO midis niveleve të ndryshme të hierarkisë së memories GPU. Zbatimi që rezulton është më i shpejtë se metodat e mëparshme si në teori (shkallëzim linear në gjatësinë e sekuencës, krahasuar me pseudo-linear për të gjitha SSM-të e bazuara në konvolucion) dhe në pajisje moderne (deri në 3× më shpejt në GPU-të A100).

Arkitekturë . Ne thjeshtojmë arkitekturat e modelit të sekuencës së thellë të mëparshme duke kombinuar dizajnin e arkitekturave të mëparshme SSM (Dao, Fu, Saab, et al. 2023) me bllokun MLP të transformatorëve në një bllok të vetëm, duke çuar në një dizajn të thjeshtë dhe homogjen të arkitekturës (Mamba) duke përfshirë hapësira selektive shtetërore.

SSM-të selektive, dhe si rrjedhojë arkitektura Mamba, janë modele plotësisht të përsëritura me vetitë kryesore që i bëjnë ato të përshtatshme si shtylla kurrizore e modeleve të përgjithshme të themelit që veprojnë në sekuenca. (i) Cilësi e lartë: selektiviteti sjell performancë të fortë në modalitete të dendura si gjuha dhe gjenomika. (ii) Trajnimi i shpejtë dhe konkludimi: llogaritja dhe memoria shkallëzohen në mënyrë lineare në gjatësinë e sekuencës gjatë trajnimit dhe zbërthimi i modelit në mënyrë autoregresive gjatë konkluzionit kërkon vetëm kohë konstante për hap pasi nuk kërkon një memorie të elementeve të mëparshme. (iii) Konteksti i gjatë: cilësia dhe efikasiteti së bashku japin përmirësime të performancës në të dhënat reale deri në gjatësinë e sekuencës 1M.

Ne vërtetojmë në mënyrë empirike potencialin e Mamba-s si një bazë FM e përgjithshme e sekuencës, si në cilësinë e trajnimit paraprak ashtu edhe në performancën e detyrave specifike për domenin, në disa lloje modalitetesh dhe cilësimesh:

• Sintetike. Për detyra të rëndësishme sintetike, të tilla si kokat e kopjimit dhe induksionit, të cilat janë propozuar si kyçe për modelet e mëdha të gjuhës, Mamba jo vetëm që i zgjidh ato lehtësisht, por mund të ekstrapolojë zgjidhje për një kohë të pacaktuar (> 1 milion shenja).

• Audio dhe Gjenomika. Mamba performon më mirë modelet e mëparshme më të avancuara si SaShiMi, Hyena dhe Transformers në modelimin e formave të valëve audio dhe sekuencave të ADN-së, si në cilësinë e para-stërvitjes ashtu edhe në metrikat e poshtme (p.sh. reduktimi i FID në një grup të dhënash sfiduese të gjenerimit të të folurit me më shumë se gjysmën ). Në të dy cilësimet, performanca e tij përmirësohet me kontekst më të gjatë deri në sekuenca të gjata miliona.

• Modelimi i gjuhës. Mamba është modeli i parë i sekuencës kohore lineare që me të vërtetë arrin performancën e cilësisë së transformatorit, si në hutimin e para-stërvitjes ashtu edhe në vlerësimet e rrjedhës së poshtme. Me ligjet e shkallëzimit deri në parametrat 1B, ne tregojmë se Mamba tejkalon performancën e një game të madhe vijash bazë, duke përfshirë receta shumë të forta moderne të trajnimit të Transformerëve të bazuara në LLaMa (Touvron et al. 2023). Modeli ynë i gjuhës Mamba ka xhiro 5× në krahasim me Transformatorët me madhësi të ngjashme dhe cilësia e Mamba-3B përputhet me atë të Transformers dyfishin e madhësisë së tij (p.sh. 4 pikë më i lartë mesatarisht në arsyetimin me sens të përbashkët në krahasim me Pythia-3B dhe madje tejkalon Pythia-7B ).