Autori:
(1) Albert Gu, oddelenie strojového učenia, Carnegie Mellon University s rovnakým príspevkom ([email protected]);
(2) Tri Dao, Katedra počítačovej vedy, Princetonská univerzita s rovnakým príspevkom ([email protected]).
Tabuľka odkazov
3 Selektívne modely stavového priestoru a 3.1 Motivácia: Selekcia ako prostriedok kompresie
3.3 Efektívna implementácia selektívnych SSM
3.4 Zjednodušená architektúra SSM
3.5 Vlastnosti selekčných mechanizmov
3.6 Ďalšie podrobnosti o modeli
4 Empirické hodnotenie a 4.1 Syntetické úlohy
4.4 Modelovanie a generovanie zvuku
A Diskusia: Mechanizmus výberu
B Súvisiace práce a B.1 Varianty a deriváty S4
B.4 Lineárna pozornosť a B.5 Modely s dlhým kontextom
D Hardvérový algoritmus pre selektívne SSM
E Experimentálne detaily a ďalšie výsledky a E.1 Syntetické úlohy
Abstraktné
Základné modely, ktoré teraz poháňajú väčšinu vzrušujúcich aplikácií v oblasti hlbokého učenia, sú takmer univerzálne založené na architektúre Transformer a jeho hlavnom module pozornosti. Na riešenie výpočtovej neefektívnosti transformátorov na dlhých sekvenciách boli vyvinuté mnohé architektúry subkvadratického času, ako sú lineárna pozornosť, hradlová konvolúcia a rekurentné modely a modely štruktúrovaného stavového priestoru (SSM), ale nefungovali tak dobre, ako je pozornosť v dôležitých modalitách, ako je jazyk. Zistili sme, že kľúčovou slabinou takýchto modelov je ich neschopnosť usudzovať na základe obsahu a urobiť niekoľko vylepšení. Po prvé, jednoducho nechať parametre SSM byť funkciami vstupu rieši ich slabosť pomocou diskrétnych modalít, čo umožňuje modelu selektívne šíriť alebo zabúdať informácie pozdĺž dimenzie dĺžky sekvencie v závislosti od aktuálneho tokenu. Po druhé, aj keď táto zmena bráni použitiu efektívnych konvolúcií, navrhujeme hardvérovo orientovaný paralelný algoritmus v rekurentnom režime. Tieto selektívne SSM integrujeme do zjednodušenej end-to-end architektúry neurónovej siete bez pozornosti alebo dokonca blokov MLP (Mamba). Mamba sa teší rýchlej inferencii (5× vyššia priepustnosť ako Transformers) a lineárnemu škálovaniu v dĺžke sekvencie a jej výkon sa zlepšuje na skutočných dátach až do miliónových sekvencií. Ako základná kosť všeobecného sekvenčného modelu dosahuje Mamba špičkový výkon v niekoľkých modalitách, ako je jazyk, zvuk a genomika. Pokiaľ ide o jazykové modelovanie, náš model Mamba-3B prekonáva transformátory rovnakej veľkosti a zodpovedá transformátorom dvojnásobnej veľkosti, a to v predtréningovom aj následnom hodnotení.
1 Úvod
Základné modely (FM) alebo veľké modely predtrénované na masívnych údajoch a následne prispôsobené pre následné úlohy sa ukázali ako účinná paradigma v modernom strojovom učení. Základom týchto FM sú často sekvenčné modely, fungujúce na ľubovoľných sekvenciách vstupov zo širokej škály domén, ako je jazyk, obrázky, reč, zvuk, časové rady a genomika (Brown a kol. 2020; Dosovitskiy a kol. 2020; Ismail Fawaz a kol. 2019; Oord2013 a kol.; 2016 Po. Sutskever, Vinyals a Quoc V Le 2014). Zatiaľ čo tento koncept je agnostický voči konkrétnemu výberu architektúry modelu, moderné FM sú prevažne založené na jedinom type sekvenčného modelu: Transformer (Vaswani et al. 2017) a jeho jadrová vrstva pozornosti (Bahdanau, Cho a Bengio 2015). Táto vlastnosť však prináša zásadné nevýhody: nemožnosť modelovať čokoľvek mimo konečného okna a kvadratické škálovanie vzhľadom na dĺžku okna. Objavilo sa obrovské množstvo výskumov o účinnejších variantoch pozornosti na prekonanie týchto nedostatkov (Tay, Dehghani, Bahri, et al. 2022), ale často na úkor samotných vlastností, ktoré ju robia efektívnou. Zatiaľ sa žiadny z týchto variantov nepreukázal ako empiricky účinný v rozsahu medzi doménami. Nedávno sa sekvenčné modely štruktúrovaného stavového priestoru (SSM) (Gu, Goel a Ré 2022; Gu, Johnson, Goel a kol. 2021) ukázali ako sľubná trieda architektúr pre sekvenčné modelovanie. Tieto modely možno interpretovať ako kombináciu rekurentných neurónových sietí (RNN) a konvolučných neurónových sietí (CNN), s inšpiráciou z klasických modelov stavového priestoru (Kalman 1960). Túto triedu modelov možno veľmi efektívne vypočítať buď ako opakovanie alebo konvolúciu, s lineárnym alebo takmer lineárnym škálovaním dĺžky sekvencie. Okrem toho majú principiálne mechanizmy na modelovanie závislostí na veľké vzdialenosti (Gu, Dao, et al. 2020) v určitých dátových modalitách a dominovali v benchmarkoch, ako je Long Range Arena (Tay, Dehghani, Abnar, et al. 2021). Mnohé príchute SSM (Gu, Goel a Ré 2022; Gu, Gupta a kol. 2022; Gupta, Gu a Berant 2022; Y. Li a kol. 2023; Ma a kol. 2023; Orvieto a kol. 2023; signál Linders in2, Warrington, kontinuálne údaje boli úspešné ako je zvuk a zrak (Goel a kol. 2022; Nguyen, Goel a kol. 2022; Saon, Gupta a Cui 2023). Boli však menej efektívne pri modelovaní diskrétnych a informačne hustých údajov, ako je text.
Navrhujeme novú triedu selektívnych modelov stavového priestoru, ktorá zlepšuje predchádzajúcu prácu na niekoľkých osiach, aby sa dosiahla modelovacia sila transformátorov pri lineárnom škálovaní v dĺžke sekvencie.
Mechanizmus výberu. Najprv identifikujeme kľúčové obmedzenie predchádzajúcich modelov: schopnosť efektívne vyberať údaje spôsobom závislým od vstupu (tj zamerať sa na konkrétne vstupy alebo ich ignorovať). Na základe intuície založenej na dôležitých syntetických úlohách, ako sú selektívne kopírovacie a indukčné hlavy, navrhujeme jednoduchý výberový mechanizmus parametrizáciou parametrov SSM na základe vstupu. To umožňuje modelu filtrovať nepodstatné informácie a pamätať si relevantné informácie na neurčito.
Algoritmus orientovaný na hardvér . Táto jednoduchá zmena predstavuje technickú výzvu pre výpočet modelu; v skutočnosti všetky predchádzajúce modely SSM musia byť časovo a vstupne invariantné, aby boli výpočtovo efektívne. Prekonáme to pomocou hardvérového algoritmu, ktorý počíta model opakovane so skenovaním namiesto konvolúcie, ale nezhmotňuje rozšírený stav, aby sa zabránilo prístupu IO medzi rôznymi úrovňami hierarchie pamäte GPU. Výsledná implementácia je rýchlejšia ako predchádzajúce metódy teoreticky (lineárne škálovanie v dĺžke sekvencie v porovnaní s pseudo-lineárnym pre všetky SSM založené na konvolúcii), ako aj na modernom hardvéri (až 3× rýchlejšie na GPU A100).
Architektúra. Zjednodušujeme predchádzajúce architektúry hlbokých sekvenčných modelov kombináciou dizajnu predchádzajúcich architektúr SSM (Dao, Fu, Saab, et al. 2023) s blokom MLP transformátorov do jedného bloku, čo vedie k jednoduchému a homogénnemu dizajnu architektúry (Mamba), ktorý zahŕňa selektívne stavové priestory. Selektívne SSM a rozšírením aj architektúra Mamba sú plne opakujúce sa modely s kľúčovými vlastnosťami, vďaka ktorým sú vhodné ako chrbtica všeobecných základných modelov fungujúcich na sekvenciách. (i) Vysoká kvalita: selektivita prináša vysoký výkon v hustých modalitách, ako je jazyk a genomika. (ii) Rýchle trénovanie a inferencia: výpočet a pamäť sa lineárne škálujú v dĺžke sekvencie počas trénovania a autoregresívne odvíjanie modelu počas inferencie vyžaduje iba konštantný čas na krok, pretože nevyžaduje vyrovnávaciu pamäť predchádzajúcich prvkov. (iii) Dlhý kontext: kvalita a efektívnosť spolu prinášajú zlepšenie výkonu na skutočných údajoch až do dĺžky sekvencie 1M.
Empiricky overujeme potenciál Mamby ako všeobecnej sekvenčnej FM chrbtice, v kvalite predtréningu aj vo výkone úloh špecifických pre doménu, na niekoľkých typoch modalít a nastavení:
• Syntetika. Pri dôležitých syntetických úlohách, ako sú kopírovacie a indukčné hlavy, ktoré boli navrhnuté ako kľúčové pre veľké jazykové modely, ich Mamba nielen ľahko vyrieši, ale dokáže extrapolovať riešenia na neurčito dlho (> 1 milión tokenov).
• Audio a genomika. Mamba prekonáva predchádzajúce najmodernejšie modely, ako sú SaShiMi, Hyena a Transformers, pokiaľ ide o modelovanie zvukových kriviek a sekvencií DNA, a to v kvalite predtrénovania aj následných metrikách (napr. zníženie FID na náročnom súbore údajov generovania reči o viac ako polovicu). V oboch nastaveniach sa jeho výkon zlepšuje s dlhším kontextom až do miliónových sekvencií.
• Jazykové modelovanie. Mamba je prvý sekvenčný model s lineárnym časom, ktorý skutočne dosahuje výkon v kvalite transformátora, a to ako pri predbežnom trénovaní, tak aj pri následných hodnoteniach. So zákonmi škálovania až do parametrov 1B ukazujeme, že Mamba prekračuje výkon veľkého rozsahu základných línií, vrátane veľmi silných moderných tréningových receptov Transformer založených na LLaMa (Touvron et al. 2023). Náš jazykový model Mamba má 5× generačnú priepustnosť v porovnaní s Transformermi podobnej veľkosti a kvalita Mamba-3B zodpovedá kvalite Transformers dvojnásobnej veľkosti (napr. o 4 body vyšší priemer na uvažovanie zdravým rozumom v porovnaní s Pythia-3B a dokonca prevyšuje Pythia-7B).
Modelový kód a vopred natrénované kontrolné body sú open source na https://github.com/state-spaces/mamba.
2 modely stavového priestoru
Modely všeobecného štátneho priestoru. Poznamenávame, že pojem model stavového priestoru má veľmi široký význam, ktorý jednoducho predstavuje pojem akéhokoľvek opakujúceho sa procesu s latentným stavom. Používa sa na označenie mnohých odlišných konceptov v rôznych disciplínach, vrátane Markovových rozhodovacích procesov (MDP) (posilňovacie učenie (Hafner et al. 2020)), dynamického kauzálneho modelovania (DCM) (výpočtová neuroveda (Friston, Harrison a Penny 2003)), Kalmanových filtrov (HMM190) dynamických modelov, skrytých a Markovových dynamických systémy (LDS) (strojové učenie) a rekurentné (a niekedy konvolučné) modely vo všeobecnosti (hĺbkové učenie).
V celom tomto článku používame výraz „SSM“ na označenie výlučne triedy štruktúrovaných modelov SSM alebo S4 (Gu, Goel a Ré 2022; Gu, Gupta a kol. 2022; Gupta, Gu a Berant 2022; Hasani a kol. 2023; Ma a kol. tieto pojmy sú zameniteľné. Pre pohodlie môžeme zahrnúť aj deriváty takýchto modelov, ako sú tie, ktoré sa zameriavajú buď na hľadisko lineárnej recidívy alebo globálnej konvolúcie (Y. Li a kol. 2023; Orvieto a kol. 2023; Poli a kol. 2023), a ak je to potrebné, objasníme nuansy.
Architektúra SSM. SSM sú samostatné sekvenčné transformácie, ktoré môžu byť začlenené do end-to-end architektúr neurónových sietí. (Architektúry SSM tiež niekedy nazývame SSNN, čo sú vrstvy SSM ako CNN vrstvy lineárnej konvolúcie.) Diskutujeme o niektorých z najznámejších architektúr SSM, z ktorých mnohé budú slúžiť aj ako naše primárne základné línie.
• Lineárna pozornosť (Katharopoulos et al. 2020) je aproximácia sebapozorovania zahŕňajúca recidívu, ktorú možno považovať za degenerovaný lineárny SSM.
• H3 (Dao, Fu, Saab a kol. 2023) zovšeobecnili túto recidívu na použitie S4; možno ju považovať za architektúru s SSM vloženým do dvoch hradlových spojení (obrázok 3). H3 tiež vloží štandardnú lokálnu konvolúciu, ktorú orámujú ako shift-SSM, pred hlavnú vrstvu SSM.
• Hyena (Poli et al. 2023) používa rovnakú architektúru ako H3, ale nahrádza vrstvu S4 globálnou konvolúciou parametrizovanou MLP (Romero et al. 2021).
• RetNet (Y. Sun et al. 2023) pridáva do architektúry ďalšiu bránu a používa jednoduchší SSM, umožňujúci alternatívnu paralelizovateľnú výpočtovú cestu, využívajúcu variant viachlavnej pozornosti (MHA) namiesto konvolúcií.
• RWKV (B. Peng et al. 2023) je nedávny RNN navrhnutý pre jazykové modelovanie založené na inej lineárnej aproximácii pozornosti (attention-free Transformer (S. Zhai et al. 2021)). Jeho hlavný mechanizmus „WKV“ zahŕňa recidívy LTI a možno ho považovať za pomer dvoch SSM.
Ďalšie úzko súvisiace SSM a architektúry sa ďalej rozoberajú v rozšírenej súvisiacej práci (príloha B). Vyzdvihujeme najmä metódy S5 (Smith, Warrington a Linderman 2023), QRNN (Bradbury et al. 2016) a SRU (Lei et al. 2017), ktoré považujeme za najpríbuznejšie metódy s naším základným selektívnym SSM.
Tento dokument je dostupný na arxiv pod licenciou CC BY 4.0 DEED.
