Princeton a CMU posouvají hranice umělé inteligence pomocí sekvenčního modelu Mamba

Tabulka odkazů

Abstrakt a 1 Úvod

2 modely stavového prostoru

3 Selektivní modely stavového prostoru a 3.1 Motivace: Selekce jako prostředek komprese

3.2 Zlepšení SSM s výběrem

3.3 Efektivní provádění selektivních SSM

3.4 Zjednodušená architektura SSM

3.5 Vlastnosti selekčního mechanismu

3.6 Další podrobnosti o modelu

4 Empirické hodnocení a 4.1 Syntetické úlohy

4.2 Jazykové modelování

4.3 Modelování DNA

4.4 Modelování a generování zvuku

4.5 Srovnávací testy rychlosti a paměti

4.6 Modelové ablace

5 Diskuse

6 Závěr a odkazy

A Diskuse: Mechanismus výběru

B Související práce

C Mechanika selektivních SSM

D Hardwarový algoritmus pro selektivní SSM

E Experimentální detaily a další výsledky

Abstraktní

Základní modely, které nyní pohánějí většinu vzrušujících aplikací v hlubokém učení, jsou téměř univerzálně založeny na architektuře Transformer a jejím hlavním modulu pozornosti. Bylo vyvinuto mnoho architektur subkvadratického času, jako je lineární pozornost, hradlová konvoluce a rekurentní modely a modely strukturovaného stavového prostoru (SSM), aby se vypořádaly s výpočetní neefektivitou transformátorů na dlouhých sekvencích, ale nevedly tak dobře jako pozornost v důležitých modalitách, jako je např. jako jazyk. Zjistili jsme, že klíčovou slabinou těchto modelů je jejich neschopnost usuzovat na základě obsahu, a provedli jsme několik vylepšení. Za prvé, prosté ponechání parametrů SSM jako funkcí vstupu řeší jejich slabost pomocí diskrétních modalit, což modelu umožňuje selektivně šířit nebo zapomínat informace podél dimenze délky sekvence v závislosti na aktuálním tokenu. Zadruhé, i když tato změna brání použití efektivních konvolucí, navrhujeme hardwarově orientovaný paralelní algoritmus v rekurentním režimu. Tyto selektivní SSM integrujeme do zjednodušené end-to-end architektury neuronové sítě bez pozornosti nebo dokonce bloků MLP (Mamba). Mamba se těší rychlé inferenci (5× vyšší propustnost než Transformers) a lineárnímu škálování v délce sekvence a její výkon se zlepšuje na skutečných datech až do sekvencí o délce milionu. Jako obecný sekvenční model páteře dosahuje Mamba nejmodernějšího výkonu v několika modalitách, jako je jazyk, zvuk a genomika. Pokud jde o jazykové modelování, náš model Mamba-3B překonává transformátory stejné velikosti a odpovídá transformátorům dvojnásobné velikosti, a to jak v předtrénování, tak v následném hodnocení.

1 Úvod

Základní modely (FM), neboli velké modely předem připravené na masivních datech a následně upravené pro následné úkoly, se ukázaly jako efektivní paradigma v moderním strojovém učení. Páteří těchto FM jsou často sekvenční modely, fungující na libovolných sekvencích vstupů z celé řady oblastí, jako je jazyk, obrázky, řeč, zvuk, časové řady a genomika (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz a kol. 2019; Sutskever, Vinyals a Quoc V Le 2014). Zatímco tento koncept je agnostický vůči konkrétní volbě modelové architektury, moderní FM jsou převážně založeny na jediném typu sekvenčního modelu: Transformeru (Vaswani et al. 2017) a jeho základní vrstvě pozornosti (Bahdanau, Cho a Bengio 2015). účinnost sebepozornosti je přisuzována její schopnosti směrovat informace hustě v rámci kontextového okna, což jí umožňuje modelovat složitá data. Tato vlastnost však přináší zásadní nevýhody: nemožnost modelovat cokoli mimo konečné okno a kvadratické škálování s ohledem na délku okna. Objevilo se obrovské množství výzkumů o účinnějších variantách pozornosti k překonání těchto nevýhod (Tay, Dehghani, Bahri, et al. 2022), ale často na úkor samotných vlastností, díky nimž je efektivní. Dosud se žádná z těchto variant neprokázala jako empiricky účinná v měřítku napříč doménami.

Nedávno se jako slibná třída architektur pro sekvenční modelování objevily strukturované stavové sekvenční modely (SSM) (Gu, Goel a Ré 2022; Gu, Johnson, Goel a kol. 2021). Tyto modely lze interpretovat jako kombinaci rekurentních neuronových sítí (RNN) a konvolučních neuronových sítí (CNN), s inspirací klasickými modely stavového prostoru (Kalman 1960). Tuto třídu modelů lze velmi efektivně vypočítat buď jako rekurentní nebo konvoluční, s lineárním nebo téměř lineárním škálováním délky sekvence. Navíc mají principiální mechanismy pro modelování závislostí na dlouhé vzdálenosti (Gu, Dao, et al. 2020) v určitých datových modalitách a dominují benchmarkům, jako je Long Range Arena (Tay, Dehghani, Abnar, et al. 2021). Mnoho pří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) byli úspěšní v oblastech zahrnujících spojitá data signálu jako je zvuk a vidění (Goel a kol. 2022; Nguyen, Goel a kol. 2022; Saon, Gupta a Cui 2023). Byly však méně účinné při modelování diskrétních dat s hustotou informací, jako je text.

Navrhujeme novou třídu selektivních stavových prostorových modelů, které vylepšují předchozí práci na několika osách, aby se dosáhlo modelovací schopnosti transformátorů při lineárním škálování v délce sekvence.

Výběrový mechanismus. Nejprve identifikujeme klíčové omezení předchozích modelů: schopnost efektivně vybírat data způsobem závislým na vstupu (tj. zaměřit se na konkrétní vstupy nebo je ignorovat). Na základě intuice založené na důležitých syntetických úlohách, jako jsou selektivní kopírovací a indukční hlavy, navrhujeme jednoduchý selekční mechanismus pomocí parametrizace parametrů SSM na základě vstupu. To umožňuje modelu odfiltrovat nepodstatné informace a pamatovat si relevantní informace donekonečna.

Algoritmus s ohledem na hardware. Tato jednoduchá změna představuje technickou výzvu pro výpočet modelu; ve skutečnosti všechny předchozí modely SSM musí být časově a vstupní invariantní, aby byly výpočetně efektivní. To jsme překonali pomocí hardwarového algoritmu, který počítá model opakovaně se skenováním namísto konvoluce, ale nerealizuje rozšířený stav, aby se zabránilo přístupu IO mezi různými úrovněmi hierarchie paměti GPU. Výsledná implementace je rychlejší než předchozí metody jak teoreticky (lineární škálování v délce sekvence, ve srovnání s pseudo-lineární pro všechny SSM založené na konvoluci), tak na moderním hardwaru (až 3× rychlejší na GPU A100).

Architektura . Zjednodušujeme dřívější architektury hlubokých sekvenčních modelů tím, že kombinujeme návrh dřívějších architektur SSM (Dao, Fu, Saab a kol. 2023) s blokem MLP transformátorů do jediného bloku, což vede k jednoduchému a homogennímu návrhu architektury (Mamba) zahrnujícímu selektivní stavové prostory.

Selektivní SSM a v rozšíření i architektura Mamba jsou plně rekurentní modely s klíčovými vlastnostmi, které je činí vhodnými jako páteř obecných základních modelů pracujících na sekvencích. (i) Vysoká kvalita: selektivita přináší vysoký výkon v hustých modalitách, jako je jazyk a genomika. (ii) Rychlé trénování a inference: výpočet a paměť lineárně škálují v délce sekvence během trénování a autoregresivní rozbalování modelu během inference vyžaduje pouze konstantní čas na krok, protože nevyžaduje vyrovnávací paměť předchozích prvků. (iii) Dlouhý kontext: kvalita a účinnost společně poskytují zlepšení výkonu na reálných datech až do délky sekvence 1M.

Empiricky ověřujeme potenciál Mamby jako obecné sekvenční páteře FM, a to jak v kvalitě předtréninku, tak v plnění úkolů specifických pro doménu, na několika typech modalit a nastavení:

• Syntetika. U důležitých syntetických úloh, jako jsou kopírovací a indukční hlavy, které byly navrženy jako klíčové pro velké jazykové modely, je Mamba nejen snadno řeší, ale může extrapolovat řešení neomezeně dlouho (>1 milion tokenů).

• Audio a genomika. Mamba překonává předchozí nejmodernější modely, jako jsou SaShiMi, Hyena a Transformers, pokud jde o modelování zvukových křivek a sekvencí DNA, a to jak v kvalitě předtrénování, tak v následných metrikách (např. snížení FID na náročném datovém souboru pro generování řeči o více než polovinu ). V obou nastaveních se jeho výkon zlepšuje s delším kontextem až do milionových sekvencí.

• Jazykové modelování. Mamba je první sekvenční model s lineárním časem, který skutečně dosahuje výkonu v kvalitě Transformer, a to jak při předtréninkové složitosti, tak při následném hodnocení. Se zákony škálování až do 1B parametrů ukazujeme, že Mamba překračuje výkon velkého rozsahu základních linií, včetně velmi silných moderních tréninkových receptur Transformer založených na LLaMa (Touvron et al. 2023). Náš jazykový model Mamba má 5× generační propustnost ve srovnání s Transformers podobné velikosti a kvalita Mamba-3B odpovídá kvalitě Transformers dvojnásobné velikosti (např. o 4 body vyšší průměr na základě zdravého rozumu ve srovnání s Pythia-3B a dokonce vyšší než Pythia-7B ).