Princeton a CMU posúvajú hranice AI pomocou sekvenčného modelu Mamba

Tabuľka odkazov

Abstrakt a 1 Úvod

2 modely stavového priestoru

3 Selektívne modely stavového priestoru a 3.1 Motivácia: Selekcia ako prostriedok kompresie

3.2 Zlepšenie SSM s výberom

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.2 Jazykové modelovanie

4.3 Modelovanie DNA

4.4 Modelovanie a generovanie zvuku

4.5 Testy rýchlosti a pamäte

4.6 Modelové ablácie

5 Diskusia

6 Záver a odkazy

A Diskusia: Mechanizmus výberu

B Súvisiace práce

C Mechanika selektívnych SSM

D Hardvérový algoritmus pre selektívne SSM

E Experimentálne detaily a ďalšie výsledky

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 bolo vyvinutých mnoho architektúr subkvadratického času, ako je lineárna pozornosť, hradlovaná konvolúcia a rekurentné modely a modely štruktúrovaného stavového priestoru (SSM), ale nefungujú tak dobre, ako sú dôležité spôsoby, ako sú napr. ako 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 et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz a kol., 2016; Sutskever, Vinyals a Quoc V Le 2014). Aj keď je tento koncept 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). účinnosť sebapozorovania sa pripisuje jeho schopnosti husto smerovať informácie v rámci kontextového okna, čo mu umožňuje modelovať komplexné údaje. 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). Mnoho chutí 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; Smith, Warrington a Linderman 2023) boli úspešní v doménach zahŕňajúcich údaje o nepretržitom signáli 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 návrhu predchádzajúcich architektúr SSM (Dao, Fu, Saab a kol. 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 silný 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).

• Zvuk 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 sa zhoduje s kvalitou Transformers dvojnásobnou jeho veľkosťou (napr. o 4 body vyšší priemer na základe zdravého rozumu v porovnaní s Pythia-3B a dokonca prevyšuje Pythia-7B ).