Princeton na CMU Push Mipaka ya AI na Mfano wa Mfuatano wa Mamba

Jedwali la Viungo

Muhtasari na 1 Utangulizi

Miundo 2 ya Nafasi za Jimbo

Miundo 3 ya Anga ya Hali Teule na 3.1 Motisha: Uteuzi kama Njia ya Mgandamizo

3.2 Kuboresha SSM kwa Chaguo

3.3 Utekelezaji Bora wa SSMs Teule

3.4 Usanifu Uliorahisishwa wa SSM

3.5 Sifa za Mbinu za Uteuzi

3.6 Maelezo ya Ziada ya Mfano

4 Tathmini ya Kijaribio na 4.1 Kazi za Sintetiki

4.2 Kuiga Lugha

4.3 Uundaji wa DNA

4.4 Uundaji wa Sauti na Uzalishaji

4.5 Vigezo vya Kasi na Kumbukumbu

4.6 Matoleo ya Mfano

5 Mazungumzo

6 Hitimisho na Marejeo

Majadiliano: Utaratibu wa Uchaguzi

B Kazi Zinazohusiana

C Mitambo ya SSMs Teule

Algorithm ya kufahamu maunzi kwa SSMs Teule

E Maelezo ya Majaribio na Matokeo ya Ziada

Muhtasari

Miundo ya msingi, ambayo sasa inaendesha programu nyingi za kusisimua katika kujifunza kwa kina, karibu kote inategemea usanifu wa Transfoma na moduli yake kuu ya umakini. Usanifu mwingi wa wakati wa subquadratic kama vile umakini wa mstari, mabadiliko ya gated na mifano ya kawaida, na miundo ya nafasi ya serikali (SSMs) imetengenezwa ili kushughulikia kutofaulu kwa hesabu kwa Transfoma kwenye mlolongo mrefu, lakini haijafanya kazi na umakini juu ya njia muhimu kama hizo. kama lugha. Tunatambua kuwa udhaifu mkuu wa miundo kama hii ni kutokuwa na uwezo wa kutoa hoja kulingana na maudhui, na kufanya maboresho kadhaa. Kwanza, kuruhusu tu vigezo vya SSM kuwa kazi za ingizo hushughulikia udhaifu wao kwa njia tofauti, ikiruhusu kielelezo kueneza kwa hiari au kusahau habari kando ya kipimo cha urefu wa mlolongo kulingana na ishara ya sasa. Pili, ingawa mabadiliko haya yanazuia utumiaji wa michanganyiko ifaayo, tunatengeneza algoriti sawia inayofahamu maunzi katika hali ya kujirudia. Tunaunganisha hizi SSM zilizochaguliwa katika usanifu wa mtandao wa neural uliorahisishwa kutoka mwisho hadi mwisho bila umakini au hata vizuizi vya MLP (Mamba). Mamba inafurahia makisio ya haraka (5× ya juu zaidi kuliko Transfoma) na kuongeza mstari katika urefu wa mfuatano, na utendakazi wake unaboresha kwenye data halisi hadi mifuatano ya urefu wa milioni. Kama uti wa mgongo wa kielelezo cha mfuatano, Mamba hufanikisha utendakazi wa hali ya juu katika mbinu kadhaa kama vile lugha, sauti na genomics. Kwenye uundaji wa lugha, muundo wetu wa Mamba-3B hupita Transfoma za ukubwa sawa na kulinganisha Transfoma mara mbili ya ukubwa wake, katika mafunzo ya awali na tathmini ya chini.

1 Utangulizi

Miundo ya msingi (FMs), au miundo mikubwa iliyofunzwa awali kwenye data kubwa kisha ikabadilishwa kwa ajili ya kazi za chini, imeibuka kama dhana bora katika kujifunza mashine ya kisasa. Uti wa mgongo wa FMS hizi mara nyingi ni miundo ya mfuatano, inayofanya kazi kwa mpangilio holela wa ingizo kutoka kwa aina mbalimbali za vikoa kama vile lugha, picha, hotuba, sauti, mfululizo wa saa na genomics (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al. 2019; 2023; Sutskever, Vinyals, and Quoc V Le 2014). Ingawa dhana hii ni ya kiagnostiki kwa chaguo fulani la usanifu wa kielelezo, FM za kisasa zinategemea zaidi aina moja ya modeli ya mfuatano: Transfoma (Vaswani et al. 2017) na safu yake kuu ya uangalizi (Bahdanau, Cho, na Bengio 2015) The ufanisi wa kujiangalia unachangiwa na uwezo wake wa kuelekeza habari kwa wingi ndani ya dirisha la muktadha, kuiruhusu kuiga data changamano. Walakini, mali hii huleta shida za kimsingi: kutokuwa na uwezo wa kuiga kitu chochote nje ya dirisha lenye kikomo, na kuongeza ukubwa wa quadratic kwa heshima na urefu wa dirisha. Kundi kubwa la utafiti limeonekana kuhusu aina bora zaidi za umakini ili kuondokana na vikwazo hivi (Tay, Dehghani, Bahri, et al. 2022), lakini mara nyingi kwa gharama ya sifa zinazoifanya iwe na ufanisi. Kufikia sasa, hakuna lahaja mojawapo kati ya hizi ambazo zimeonyeshwa kuwa na ufanisi katika kiwango katika vikoa.

Hivi majuzi, miundo ya mpangilio wa anga za juu (SSM) (Gu, Goel, na Ré 2022; Gu, Johnson, Goel, et al. 2021) imeibuka kama aina ya usanifu wa kuahidi wa uundaji wa mfuatano. Miundo hii inaweza kufasiriwa kama mchanganyiko wa mitandao ya neural inayojirudia (RNNs) na mitandao ya neural convolutional (CNNs), yenye msukumo kutoka kwa miundo ya anga ya hali ya juu (Kalman 1960). Daraja hili la miundo linaweza kukokotwa kwa ufanisi sana kama ujirudiaji au ugeuzaji, kwa kuongeza mstari au karibu na mstari katika urefu wa mfuatano. Zaidi ya hayo, wana mbinu za kanuni za kuiga utegemezi wa masafa marefu (Gu, Dao, et al. 2020) katika mbinu fulani za data, na wametawala vigezo kama vile Long Range Arena (Tay, Dehghani, Abnar, et al. 2021). Ladha nyingi za SSM (Gu, Goel, and Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, and Berrant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith, Warrington, na Linderman 2023) wamefaulu katika vikoa vinavyohusisha data ya mawimbi endelevu kama vile sauti na maono (Goel et al. 2022; Nguyen, Goel, et al. 2022; Saon, Gupta, and Cui 2023). Hata hivyo, zimekuwa na ufanisi mdogo katika kuiga data bainifu na zenye habari nyingi kama vile maandishi.

Tunapendekeza aina mpya ya miundo ya anga ya hali iliyochaguliwa, ambayo inaboresha kazi ya awali kwenye shoka kadhaa ili kufikia uwezo wa kielelezo wa Transfoma huku ikipanua kimstari kwa urefu wa mfuatano.

Utaratibu wa Uchaguzi. Kwanza, tunatambua kizuizi kikuu cha miundo ya awali: uwezo wa kuchagua data kwa ufanisi kwa njia inayotegemea ingizo (yaani kuzingatia au kupuuza ingizo mahususi). Kwa kuzingatia angavu kulingana na kazi muhimu za sanisi kama vile nakala teule na vichwa vya utangulizi, tunabuni mbinu rahisi ya uteuzi kwa kuainisha vigezo vya SSM kulingana na ingizo. Hii inaruhusu kielelezo kuchuja taarifa zisizo muhimu na kukumbuka taarifa muhimu kwa muda usiojulikana.

Algorithm ya kufahamu vifaa. Mabadiliko haya rahisi yanaleta changamoto ya kiufundi kwa ukokotoaji wa modeli; kwa kweli, miundo yote ya awali ya SSM lazima ziwe za kutofautiana kwa wakati na pembejeo ili kuwa na ufanisi wa kimahesabu. Tunashinda hili kwa algoriti inayofahamu maunzi ambayo hukokotoa muundo mara kwa mara na uchanganuzi badala ya ubadilishaji, lakini haifanyiki hali iliyopanuliwa ili kuepuka ufikiaji wa IO kati ya viwango tofauti vya safu ya kumbukumbu ya GPU. Utekelezaji unaotokana ni wa haraka zaidi kuliko mbinu za awali katika nadharia (kuongeza kimstari kwa urefu wa mfuatano, ikilinganishwa na mstari wa uwongo kwa SSM zote zinazotegemea ubadilishaji) na kwenye maunzi ya kisasa (hadi 3× haraka zaidi kwenye A100 GPUs).

Usanifu . Tunarahisisha usanifu wa miundo ya mfuatano wa kina wa awali kwa kuchanganya muundo wa usanifu wa awali wa SSM (Dao, Fu, Saab, na wenzie. 2023) na muundo wa MLP wa Transfoma hadi kwenye kitalu kimoja, na hivyo kusababisha muundo rahisi na unaofanana wa usanifu (Mamba) unaojumuisha. nafasi za serikali zilizochaguliwa.

SSM teule, na kwa upanuzi usanifu wa Mamba, ni miundo inayojirudia rudia yenye sifa kuu zinazozifanya zifae kama uti wa mgongo wa miundo ya jumla ya msingi inayofanya kazi kwenye mfuatano. (i) Ubora wa juu: uteuzi huleta utendaji dhabiti kwenye mbinu mnene kama vile lugha na genomics. (ii) Mafunzo ya haraka na makisio: ukokotoaji na mizani ya kumbukumbu kwa mstari katika urefu wa mfuatano wakati wa mafunzo, na kufungua kielelezo kiotomatiki wakati wa makisio kunahitaji muda thabiti pekee kwa kila hatua kwa kuwa hauhitaji akiba ya vipengele vya awali. (iii) Muktadha mrefu: ubora na ufanisi kwa pamoja huleta maboresho ya utendakazi kwenye data halisi hadi urefu wa mfuatano wa 1M.

Tunathibitisha kwa uthabiti uwezo wa Mamba kama uti wa mgongo wa mfululizo wa FM, katika ubora wa mafunzo ya awali na utendakazi wa kazi mahususi wa kikoa, kwenye aina kadhaa za kanuni na mipangilio:

• Sintetiki. Juu ya kazi muhimu za usanifu kama vile kunakili na vichwa vya utangulizi ambavyo vimependekezwa kuwa muhimu kwa miundo mikubwa ya lugha, Mamba haisuluhishi kwa urahisi tu bali inaweza kutoa suluhu kwa muda usiojulikana (> tokeni 1M).

• Sauti na Genomics. Mamba huigiza zaidi miundo ya hali ya juu kama vile SaShiMi, Fisi, na Transfoma kwenye muundo wa mawimbi ya sauti na mifuatano ya DNA, katika mafunzo ya awali ya ubora na vipimo vya chini (kwa mfano, kupunguza FID kwenye mkusanyiko wa data wenye changamoto wa kuzalisha usemi kwa zaidi ya nusu. ) Katika mipangilio yote miwili, utendakazi wake huboreshwa kwa muktadha mrefu hadi mfuatano wa urefu wa milioni.

• Kuiga Lugha. Mamba ni kielelezo cha kwanza cha mfuatano wa wakati ambao hufanikisha utendakazi wa ubora wa Transfoma, katika mafunzo ya kutatanisha na tathmini za chini. Kwa kuongeza sheria za hadi vigezo vya 1B, tunaonyesha kuwa Mamba inazidi utendakazi wa anuwai kubwa ya msingi, ikijumuisha mapishi ya kisasa ya mafunzo ya Transformer kulingana na LLaMa (Touvron et al. 2023). Muundo wetu wa lugha ya Mamba una uboreshaji wa kizazi cha 5× ikilinganishwa na Transfoma za ukubwa sawa, na ubora wa Mamba-3B unalingana na ule wa Transfoma mara mbili ya ukubwa wake (km. pointi 4 juu zaidi. kwa hoja za kawaida ikilinganishwa na Pythia-3B na hata kuzidi Pythia-7B )