Svět umělé inteligence má nového miláčka – a není to transformátor

Tabulka odkazů

Abstrakt a 1. Úvod

2 modely stavového prostoru

3 Selektivní stavové modely 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, poděkování a odkazy

A Diskuse: Mechanismus výběru

B Související práce a B.1 Varianty a deriváty S4

B.2 Architektury SSM

B.3 Vztah k RNN

B.4 Lineární pozornost a B.5 Modely s dlouhým kontextem

C Mechanika selektivních SSM

D Hardwarový algoritmus pro selektivní SSM

E Experimentální detaily a další výsledky a E.1 Syntetické úlohy

E.2 Jazykové modelování

E.3 Modelování DNA

E.4 Podrobnosti o zvuku

E.5 Benchmark účinnosti

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. Mnoho architektur subkvadratického času, jako je lineární pozornost, hradlová konvoluce a rekurentní modely a modely strukturovaného stavového prostoru (SSM), bylo vyvinuto pro řešení výpočetní neefektivnosti transformátorů na dlouhých sekvencích, ale nevedly tak dobře jako pozornost v důležitých modalitách, jako je 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. Za druhé, 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 široké škály domén, jako je jazyk, obrázky, řeč, zvuk, časové řady a genomika (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al. 2019; Oord2023 et al.; 2016 et al. Sutskever, Vinyals a Quoc V Le 2014). Zatímco tento koncept je agnostický ke konkrétní volbě architektury modelu, 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). 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, et al. 2022; Gupta, Gu a Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; signál Linders in3, Warrington, kontinuální data byly úspěšné 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 filtrovat 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 sekvencí tím, že kombinujeme návrh dřívějších architektur SSM (Dao, Fu, Saab, et al. 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ě vedou ke 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 u náročné datové sady 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 uvažování selským rozumem ve srovnání s Pythia-3B a dokonce přesahující Pythia-7B).

Modelový kód a předtrénované kontrolní body jsou open-source na https://github.com/state-spaces/mamba.

2 modely stavového prostoru

Modely obecného stavového prostoru. Všimli jsme si, že pojem stavový model má velmi široký význam, který jednoduše představuje pojem jakéhokoli rekurentního procesu s latentním stavem. Používá se k odkazování na mnoho nesourodých konceptů v různých disciplínách, včetně Markovových rozhodovacích procesů (MDP) (posílení učení (Hafner et al. 2020)), dynamického kauzálního modelování (DCM) (výpočetní neurověda (Friston, Harrison a Penny 2003)), Kalmanových filtrů (HMM190) dynamických modelů (HMM190) dynamických systémy (LDS) (strojové učení) a rekurentní (a někdy konvoluční) modely obecně (hluboké učení).

V celém tomto článku používáme termín „SSM“ k označení výhradně třídy strukturovaných SSM nebo modelů S4 (Gu, Goel a Ré 2022; Gu, Gupta a kol. 2022; Gupta, Gu a Berant 2022; použití Hasani a kol. 2023; Ma a kol.; Linderman) 20,302 tyto termíny jsou zaměnitelné. Pro usnadnění můžeme také zahrnout deriváty takových modelů, jako jsou ty, které se zaměřují buď na hledisko lineární recidivy nebo globální konvoluce (Y. Li et al. 2023; Orvieto et al. 2023; Poli et al. 2023), a v případě potřeby objasnit nuance.

Architektury SSM. SSM jsou samostatné sekvenční transformace, které lze začlenit do end-to-end architektur neuronových sítí. (Někdy také nazýváme architektury SSM SSNN, což jsou vrstvy SSM jako CNN vrstvy lineární konvoluce.) Diskutujeme o některých nejznámějších architekturách SSM, z nichž mnohé budou také sloužit jako naše primární základní linie.

• Lineární pozornost (Katharopoulos et al. 2020) je aproximací sebepozornosti zahrnující recidivu, kterou lze považovat za degenerovaný lineární SSM.

• H3 (Dao, Fu, Saab, et al. 2023) zobecnil tuto recidivu na použití S4; lze na něj nahlížet jako na architekturu s SSM vloženým do dvou hradlových spojení (obrázek 3). H3 také vloží standardní místní konvoluci, kterou rámují jako shift-SSM, před hlavní vrstvu SSM.

• Hyena (Poli et al. 2023) používá stejnou architekturu jako H3, ale nahrazuje vrstvu S4 globální konvolucí parametrizovanou MLP (Romero et al. 2021).

• RetNet (Y. Sun et al. 2023) přidává další bránu do architektury a používá jednodušší SSM, což umožňuje alternativní paralelizovatelnou výpočetní cestu využívající variantu vícehlavé pozornosti (MHA) namísto konvolucí.

• RWKV (B. Peng et al. 2023) je nedávný RNN navržený pro jazykové modelování založené na jiné lineární aproximaci pozornosti (attention-free Transformer (S. Zhai et al. 2021)). Jeho hlavní mechanismus „WKV“ zahrnuje opakování LTI a lze jej považovat za poměr dvou SSM.

Další úzce související SSM a architektury jsou dále diskutovány v rozšířené související práci (příloha B). Zdůrazňujeme zejména S5 (Smith, Warrington a Linderman 2023), QRNN (Bradbury et al. 2016) a SRU (Lei et al. 2017), které považujeme za nejblíže související metody s naším základním selektivním SSM.