Die KI-wêreld het 'n nuwe liefling - en dit is nie 'n transformator nie

Tabel van skakels

Abstrakte en 1. Inleiding

2 Staatsruimtemodelle

3 Selektiewe Toestandsruimtemodelle en 3.1 Motivering: Seleksie as 'n middel van kompressie

3.2 Verbetering van SSM'e met seleksie

3.3 Doeltreffende implementering van selektiewe SSM'e

3.4 'n Vereenvoudigde SSM-argitektuur

3.5 Eienskappe van seleksiemeganismes

3.6 Bykomende modelbesonderhede

4 Empiriese evaluering en 4.1 Sintetiese take

4.2 Taalmodellering

4.3 DNA-modellering

4.4 Oudiomodellering en generering

4.5 Spoed en geheue maatstawwe

4.6 Model Ablasies

5 Bespreking

6 Gevolgtrekking, erkennings en verwysings

'n Bespreking: Keurmeganisme

B Verwante Werk en B.1 S4 Variante en Afgeleides

B.2 SSM-argitekture

B.3 Verwantskap met RNN'e

B.4 Lineêre Aandag- en B.5 Langkonteksmodelle

C Meganika van selektiewe SSM'e

D Hardeware-bewuste algoritme vir selektiewe SSM'e

E Eksperimentele besonderhede en bykomende resultate en E.1 Sintetiese take

E.2 Taalmodellering

E.3 DNA-modellering

E.4 Oudiobesonderhede

E.5 Doeltreffendheidsmaatstaf

Abstrak

Grondslagmodelle, wat nou die meeste van die opwindende toepassings in diepleer aandryf, is byna universeel gebaseer op die Transformer-argitektuur en sy kernaandagmodule. Baie subkwadratiese-tyd-argitekture soos lineêre aandag, omheinde konvolusie en herhalende modelle, en gestruktureerde toestandruimtemodelle (SSM's) is ontwikkel om Transformers se rekenaarondoeltreffendheid op lang rye aan te spreek, maar hulle het nie so goed presteer as aandag aan belangrike modaliteite soos taal nie. Ons identifiseer dat 'n belangrike swakheid van sulke modelle hul onvermoë is om inhoudgebaseerde redenasies uit te voer, en verskeie verbeterings aan te bring. Eerstens, om bloot die SSM-parameters te laat funksioneer van die invoer, spreek hul swakheid aan met diskrete modaliteite, wat die model in staat stel om inligting selektief te versprei of te vergeet langs die volgordelengte-dimensie, afhangende van die huidige teken. Tweedens, alhoewel hierdie verandering die gebruik van doeltreffende konvolusies verhoed, ontwerp ons 'n hardeware-bewuste parallelle algoritme in herhalende modus. Ons integreer hierdie selektiewe SSM's in 'n vereenvoudigde end-tot-end neurale netwerkargitektuur sonder aandag of selfs MLP-blokke (Mamba). Mamba geniet vinnige afleiding (5× hoër deurset as Transformers) en lineêre skaal in volgordelengte, en sy werkverrigting verbeter op werklike data tot miljoen-lengte rye. As 'n algemene volgordemodel-ruggraat, behaal Mamba die nuutste prestasie oor verskeie modaliteite soos taal, oudio en genomika. Wat taalmodellering betref, presteer ons Mamba-3B-model beter as Transformers van dieselfde grootte en pas Transformers twee keer sy grootte ooreen, beide in vooropleiding en stroomaf-evaluering.

1 Inleiding

Grondslagmodelle (FM's), of groot modelle wat vooraf opgelei is op massiewe data wat dan aangepas is vir stroomaftake, het na vore gekom as 'n effektiewe paradigma in moderne masjienleer. Die ruggraat van hierdie FM's is dikwels volgordemodelle, wat op arbitrêre reekse van insette van 'n wye verskeidenheid domeine werk, soos taal, beelde, spraak, oudio, tydreekse en genomika (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al. 2019; 201 et al. Poli et al. 201 et al. Sutskever, Vinyals en Quoc V Le 2014). Alhoewel hierdie konsep agnosties is vir 'n spesifieke keuse van modelargitektuur, is moderne FM's oorwegend gebaseer op 'n enkele tipe volgordemodel: die Transformer (Vaswani et al. 2017) en sy kern aandaglaag (Bahdanau, Cho, en Bengio 2015). Die doeltreffendheid van self-aandag word toegeskryf aan sy vermoë om dit in 'n komplekse datavenster te modelleer. Hierdie eienskap bring egter fundamentele nadele: 'n onvermoë om enigiets buite 'n eindige venster te modelleer, en kwadratiese skaal met betrekking tot die vensterlengte. 'n Enorme hoeveelheid navorsing het verskyn oor meer doeltreffende variante van aandag om hierdie nadele te oorkom (Tay, Dehghani, Bahri, et al. 2022), maar dikwels ten koste van die einste eienskappe wat dit effektief maak. Daar is nog nie getoon dat een van hierdie variante empiries effektief is op skaal oor domeine heen nie. Onlangs het gestruktureerde toestandruimtevolgordemodelle (SSM'e) (Gu, Goel en Ré 2022; Gu, Johnson, Goel, et al. 2021) na vore gekom as 'n belowende klas argitekture vir volgordemodellering. Hierdie modelle kan geïnterpreteer word as 'n kombinasie van herhalende neurale netwerke (RNN'e) en konvolusionele neurale netwerke (CNN's), met inspirasie van klassieke toestandruimtemodelle (Kalman 1960). Hierdie klas modelle kan baie doeltreffend bereken word as óf 'n herhaling óf konvolusie, met lineêre of naby-lineêre skaal in volgordelengte. Boonop het hulle beginselmeganismes vir die modellering van langafstandafhanklikhede (Gu, Dao, et al. 2020) in sekere datamodaliteite, en het maatstawwe soos die Long Range Arena oorheers (Tay, Dehghani, Abnar, et al. 2021). Baie geure van SSM's (Gu, Goel, en Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, en Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith, Warrington20 en23man) het deurlopende data in domvings getoon. soos oudio en visie (Goel et al. 2022; Nguyen, Goel, et al. 2022; Saon, Gupta en Cui 2023). Hulle was egter minder effektief met die modellering van diskrete en inligtingdigte data soos teks.

Ons stel 'n nuwe klas selektiewe toestandruimtemodelle voor, wat verbeter op vorige werk op verskeie asse om die modelleringskrag van transformators te bereik terwyl hulle lineêr in volgordelengte skaal.

Keuringsmeganisme. Eerstens identifiseer ons 'n sleutelbeperking van vorige modelle: die vermoë om data doeltreffend op 'n insetafhanklike wyse te selekteer (dws fokus op of spesifieke insette te ignoreer). Gebou op intuïsie gebaseer op belangrike sintetiese take soos selektiewe kopie- en induksiekoppe, ontwerp ons 'n eenvoudige seleksiemeganisme deur die SSM-parameters op grond van die insette te parameteriseer. Dit laat die model toe om irrelevante inligting uit te filter en relevante inligting onbepaald te onthou.

Hardeware-bewuste algoritme . Hierdie eenvoudige verandering stel 'n tegniese uitdaging vir die berekening van die model; om die waarheid te sê, alle vorige SSM-modelle moet tyd- en inset-onveranderlik wees om rekenkundig doeltreffend te wees. Ons oorkom dit met 'n hardeware-bewuste algoritme wat die model herhaaldelik met 'n skandering in plaas van konvolusie bereken, maar nie die uitgebreide toestand realiseer nie om IO-toegang tussen verskillende vlakke van die GPU-geheuehiërargie te vermy. Die gevolglike implementering is vinniger as vorige metodes, beide in teorie (skaal lineêr in volgorde lengte, in vergelyking met pseudo-lineêr vir alle konvolusie-gebaseerde SSM's) en op moderne hardeware (tot 3× vinniger op A100 GPU's).

Argitektuur. Ons vereenvoudig vorige diepvolgorde-modelargitekture deur die ontwerp van vorige SSM-argitekture (Dao, Fu, Saab, et al. 2023) met die MLP-blok van Transformers in 'n enkele blok te kombineer, wat lei tot 'n eenvoudige en homogene argitektuurontwerp (Mamba) wat selektiewe staatsruimtes insluit. Selektiewe SSM'e, en by uitbreiding die Mamba-argitektuur, is volledig herhalende modelle met sleuteleienskappe wat hulle geskik maak as die ruggraat van algemene fondamentmodelle wat op volgordes werk. (i) Hoë kwaliteit: selektiwiteit bring sterk prestasie op digte modaliteite soos taal en genomika. (ii) Vinnige opleiding en inferensie: berekening en geheue skale lineêr in volgorde lengte tydens opleiding, en die afrol van die model outoregressief tydens inferensie vereis slegs konstante tyd per stap aangesien dit nie 'n kas van vorige elemente vereis nie. (iii) Lang konteks: die kwaliteit en doeltreffendheid lewer saam prestasieverbeterings op werklike data tot op volgordelengte 1M.

Ons bevestig empiries Mamba se potensiaal as 'n algemene volgorde FM-ruggraat, in beide vooropleidingskwaliteit en domeinspesifieke taakverrigting, op verskeie tipes modaliteite en instellings:

• Sintetiese stowwe. Op belangrike sintetiese take soos kopiëring en induksiekoppe wat voorgestel is as die sleutel tot groot taalmodelle, los Mamba dit nie net maklik op nie, maar kan oplossings onbepaald lank ekstrapoleer (> 1M tokens).

• Oudio en Genomika. Mamba presteer beter as vorige moderne modelle soos SaShiMi, Hyena en Transformers met die modellering van oudiogolfvorms en DNS-volgordes, beide in vooropleidingskwaliteit en stroomaf-metrieke (bv. die vermindering van FID op 'n uitdagende spraakgenereringsdatastel met meer as die helfte). In beide instellings verbeter sy werkverrigting met langer konteks tot miljoen-lengte rye.

• Taalmodellering. Mamba is die eerste lineêre-tydvolgordemodel wat werklik transformatorgehalte-prestasie behaal, beide in voorafopleiding verwarring en stroomaf-evaluasies. Met skaalwette tot 1B-parameters wys ons dat Mamba die prestasie van 'n groot reeks basislyne oorskry, insluitend baie sterk moderne Transformer-opleidingsresepte gebaseer op LLaMa (Touvron et al. 2023). Ons Mamba-taalmodel het 5× generasie deurset in vergelyking met Transformers van soortgelyke grootte, en Mamba-3B se gehalte stem ooreen met dié van Transformers twee keer sy grootte (bv. 4 punte hoër gemiddelde op gesonde verstand redenering in vergelyking met Pythia-3B en selfs meer as Pythia-7B).

Modelkode en vooraf-opgeleide kontrolepunte is oopbron by https://github.com/state-spaces/mamba.

2 Staatsruimtemodelle

Algemene Staatsruimtemodelle. Ons let daarop dat die term staatsruimtemodel 'n baie wye betekenis het wat bloot die idee van enige herhalende proses met 'n latente toestand verteenwoordig. Dit is gebruik om na baie uiteenlopende konsepte in verskillende dissiplines te verwys, insluitend Markov-besluitprosesse (MDP) (versterkingsleer (Hafner et al. 2020)), dinamiese oorsaaklike modellering (DCM) (berekeningsneurologie (Friston, Harrison, en Penny 2003)), Kalman-filters (kontroles (kontroles (Kalman) en hidden-dinamiese lynmodelle (Kalman 19 Markovar-model) (HMM60). stelsels (LDS) (masjienleer), en herhalende (en soms konvolusionele) modelle in die algemeen (diep leer).

Dwarsdeur hierdie hele referaat gebruik ons die term "SSM" om uitsluitlik te verwys na die klas van gestruktureerde SSM'e of S4-modelle (Gu, Goel, en Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, en Berant 2022; Hasani et al. 2023; Ma et al. 2023, 2023, 2003, Linder Warrington en Linderman) uitruilbaar. Gerieflikheidshalwe kan ons ook afgeleides van sulke modelle insluit, soos dié wat fokus op óf die lineêre-herhaling óf globale-konvolusie-standpunte (Y. Li et al. 2023; Orvieto et al. 2023; Poli et al. 2023), en verduidelik nuanses wanneer nodig.

SSM Argitekture. SSM'e is selfstandige volgorde-transformasies wat by end-tot-end neurale netwerkargitekture geïnkorporeer kan word. (Ons noem ook soms SSM-argitekture SSNN'e, wat na SSM-lae is, soos CNN'e na lineêre konvolusie-lae is.) Ons bespreek sommige van die mees bekende SSM-argitekture, waarvan baie ook as ons primêre basislyne sal dien.

• Lineêre aandag (Katharopoulos et al. 2020) is 'n benadering van selfaandag wat 'n herhaling behels wat as 'n gedegenereerde lineêre SSM beskou kan word.

• H3 (Dao, Fu, Saab, et al. 2023) het hierdie herhaling veralgemeen om S4 te gebruik; dit kan gesien word as 'n argitektuur met 'n SSM wat deur twee omheinde verbindings vasgemaak is (Figuur 3). H3 voeg ook 'n standaard plaaslike konvolusie in, wat hulle raam as 'n skuif-SSM, voor die hoof SSM-laag.

• Hiëna (Poli et al. 2023) gebruik dieselfde argitektuur as H3 maar vervang die S4-laag met 'n MLP-geparameteriseerde globale konvolusie (Romero et al. 2021).

• RetNet (Y. Sun et al. 2023) voeg 'n bykomende hek by die argitektuur en gebruik 'n eenvoudiger SSM, wat 'n alternatiewe paralleliseerbare berekeningspad moontlik maak, deur 'n variant van multi-kop aandag (MHA) in plaas van konvolusies te gebruik.

• RWKV (B. Peng et al. 2023) is 'n onlangse RNN wat ontwerp is vir taalmodellering gebaseer op 'n ander lineêre aandagbenadering (aandagvrye Transformer (S. Zhai et al. 2021)). Sy hoof "WKV"-meganisme behels LTI-herhalings en kan gesien word as die verhouding van twee SSM'e.

Ander nou verwante SSM'e en argitekture word verder bespreek in 'n uitgebreide verwante werk (Bylae B). Ons beklemtoon veral S5 (Smith, Warrington en Linderman 2023), QRNN (Bradbury et al. 2016) en SRU (Lei et al. 2017), wat ons beskou as die metodes wat die naaste aan ons kern selektiewe SSM is.