Mga may-akda:
(1) Albert Gu, Machine Learning Department, Carnegie Mellon University na may Pantay na kontribusyon ([email protected]);
(2) Tri Dao, Department of Computer Science, Princeton University na may Pantay na kontribusyon ([email protected]).
Talaan ng mga Link
2 Mga Modelo ng Kalawakan ng Estado
3 Selective State Space Models at 3.1 Motivation: Selection as a Means of Compression
3.2 Pagpapabuti ng mga SSM na may Pinili
3.3 Mahusay na Pagpapatupad ng Selective SSMs
3.4 Isang Pinasimpleng Arkitektura ng SSM
3.5 Mga Katangian ng Mekanismo ng Pagpili
3.6 Karagdagang Detalye ng Modelo
4 Empirical Evaluation at 4.1 Sintetikong Gawain
4.4 Pagmomodelo at Pagbuo ng Audio
4.5 Mga Benchmark ng Bilis at Memorya
6 Konklusyon, Pagkilala at Mga Sanggunian
Isang Pagtalakay: Mekanismo ng Pagpili
B Kaugnay na Trabaho at B.1 S4 Variants at Derivatives
B.4 Linear Attention at B.5 Long Context Models
D Hardware-aware Algorithm Para sa Selective SSMs
E Mga Detalye ng Eksperimental at Karagdagang Resulta at E.1 Mga Sintetikong Gawain
Abstract
Ang mga modelo ng pundasyon, na ngayon ay nagpapagana sa karamihan ng mga kapana-panabik na aplikasyon sa malalim na pag-aaral, ay halos lahat ay nakabatay sa arkitektura ng Transformer at sa pangunahing module ng atensyon nito. Maraming mga subquadratic-time na arkitektura tulad ng linear na atensyon, gated convolution at paulit-ulit na mga modelo, at structured state space models (SSMs) ang binuo upang matugunan ang computational inefficiency ng Transformers sa mahabang pagkakasunud-sunod, ngunit hindi sila gumanap pati na rin ang pansin sa mahahalagang modalidad tulad ng wika. Tinutukoy namin na ang isang pangunahing kahinaan ng mga naturang modelo ay ang kanilang kawalan ng kakayahan na magsagawa ng pangangatwiran batay sa nilalaman, at gumawa ng ilang mga pagpapabuti. Una, hinahayaan lamang ang mga parameter ng SSM na maging function ng input na tumutugon sa kanilang kahinaan na may mga discrete modalities, na nagpapahintulot sa modelo na piliing magpalaganap o makalimot ng impormasyon kasama ang dimensyon ng haba ng sequence depende sa kasalukuyang token. Pangalawa, kahit na pinipigilan ng pagbabagong ito ang paggamit ng mga mahusay na convolution, nagdidisenyo kami ng parallel algorithm na may alam sa hardware sa paulit-ulit na mode. Isinasama namin ang mga piling SSM na ito sa isang pinasimple na end-to-end na arkitektura ng neural network nang walang pansin o kahit na mga MLP block (Mamba). Tinatangkilik ng Mamba ang mabilis na inference (5x na mas mataas na throughput kaysa sa mga Transformer) at linear scaling sa haba ng pagkakasunud-sunod, at ang pagganap nito ay bumubuti sa totoong data hanggang sa mga sequence na may milyong haba. Bilang backbone ng pangkalahatang sequence model, nakakamit ng Mamba ang makabagong pagganap sa ilang mga modalidad gaya ng wika, audio, at genomics. Sa pagmomodelo ng wika, ang aming Mamba-3B na modelo ay higit na gumaganap sa mga Transformer na may parehong laki at tumutugma sa mga Transformer nang dalawang beses sa laki nito, kapwa sa pretraining at downstream na pagsusuri.
1 Panimula
Ang mga foundation model (FM), o malalaking modelo na na-pretrained sa napakalaking data pagkatapos ay inangkop para sa mga downstream na gawain, ay lumitaw bilang isang epektibong paradigm sa modernong machine learning. Ang backbone ng mga FM na ito ay madalas na mga modelo ng pagkakasunud-sunod, na tumatakbo sa mga arbitrary na pagkakasunud-sunod ng mga input mula sa iba't ibang uri ng mga domain tulad ng wika, mga imahe, pagsasalita, audio, serye ng oras, at genomics (Brown et al. 2020; Dosovitskiy et al. 2020; Ismail Fawaz et al. 2019; Poli201t al. 2019; Poli201t al. 2019; Poli201t al. Sutskever, Vinyals, at Quoc V Le 2014). Bagama't agnostiko ang konseptong ito sa isang partikular na pagpipilian ng arkitektura ng modelo, ang mga modernong FM ay pangunahing nakabatay sa isang solong uri ng modelo ng pagkakasunud-sunod: ang Transformer (Vaswani et al. 2017) at ang pangunahing layer ng atensyon nito (Bahdanau, Cho, at Bengio 2015) Ang pagiging epektibo ng pansin sa sarili ay naiuugnay sa kakayahan nitong magruta ng impormasyon sa konteksto nang masinsinan. Gayunpaman, ang property na ito ay nagdudulot ng mga pangunahing disbentaha: isang kawalan ng kakayahang magmodelo ng anuman sa labas ng isang may hangganang window, at quadratic scaling na may kinalaman sa haba ng window. Ang isang napakalaking pangkat ng pananaliksik ay lumitaw sa mas mahusay na mga variant ng atensyon upang malampasan ang mga kakulangan na ito (Tay, Dehghani, Bahri, et al. 2022), ngunit kadalasan sa kapinsalaan ng mismong mga katangian na ginagawang epektibo. Sa ngayon, wala sa mga variant na ito ang naipakita na empirically epektibo sa laki sa mga domain. Kamakailan, ang structured state space sequence models (SSMs) (Gu, Goel, at Ré 2022; Gu, Johnson, Goel, et al. 2021) ay lumitaw bilang isang magandang klase ng mga arkitektura para sa sequence modelling. Ang mga modelong ito ay maaaring bigyang-kahulugan bilang isang kumbinasyon ng mga paulit-ulit na neural network (RNNs) at convolutional neural network (CNNs), na may inspirasyon mula sa mga klasikal na modelo ng espasyo ng estado (Kalman 1960). Ang klase ng mga modelong ito ay maaaring kalkulahin nang napakahusay bilang alinman sa pag-ulit o convolution, na may linear o malapit-linear na scaling sa haba ng pagkakasunud-sunod. Bukod pa rito, mayroon silang mga prinsipyong mekanismo para sa pagmomodelo ng mga long-range na dependency (Gu, Dao, et al. 2020) sa ilang partikular na data modalities, at nangingibabaw ang mga benchmark gaya ng Long Range Arena (Tay, Dehghani, Abnar, et al. 2021). Maraming lasa ng mga SSM (Gu, Goel, at Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, at Berant 2022; Y. Li et al. 2023; Ma et al. 2023; Orvieto et al. 2023; Smith, Warrington, at Linderman) ang naging matagumpay sa mga naturang domain ng signal. bilang audio at vision (Goel et al. 2022; Nguyen, Goel, et al. 2022; Saon, Gupta, and Cui 2023). Gayunpaman, hindi gaanong epektibo ang mga ito sa pagmomodelo ng discrete at siksik na impormasyon na data gaya ng text.
Iminumungkahi namin ang isang bagong klase ng mga pumipili na modelo ng espasyo ng estado, na nagpapabuti sa naunang gawain sa ilang mga axes upang makamit ang kapangyarihan ng pagmomodelo ng mga Transformer habang lumi-scale nang linear sa haba ng pagkakasunud-sunod.
Mekanismo ng Pagpili. Una, tinutukoy namin ang isang pangunahing limitasyon ng mga naunang modelo: ang kakayahang pumili ng data nang mahusay sa paraang nakadepende sa input (ibig sabihin, tumuon sa o huwag pansinin ang mga partikular na input). Pagbuo sa intuwisyon batay sa mahahalagang gawaing gawa ng tao tulad ng selective copy at induction head, nagdidisenyo kami ng simpleng mekanismo ng pagpili sa pamamagitan ng pag-parameter sa mga parameter ng SSM batay sa input. Nagbibigay-daan ito sa modelo na i-filter ang walang kaugnayang impormasyon at tandaan ang nauugnay na impormasyon nang walang katapusan.
Hardware-aware Algorithm . Ang simpleng pagbabagong ito ay nagdudulot ng teknikal na hamon para sa pagkalkula ng modelo; sa katunayan, ang lahat ng mga naunang modelo ng SSM ay dapat na time- at input-invariant upang maging mahusay sa computation. Napagtatagumpayan namin ito gamit ang isang algorithm na may kamalayan sa hardware na kino-compute ang modelo nang paulit-ulit gamit ang isang pag-scan sa halip na convolution, ngunit hindi natutupad ang pinalawak na estado upang maiwasan ang pag-access sa IO sa pagitan ng iba't ibang antas ng hierarchy ng memorya ng GPU. Ang resultang pagpapatupad ay mas mabilis kaysa sa mga nakaraang pamamaraan pareho sa teorya (pag-scale nang linear sa haba ng sequence, kumpara sa pseudo-linear para sa lahat ng convolution-based na SSM) at sa modernong hardware (hanggang sa 3x na mas mabilis sa A100 GPUs).
Arkitektura. Pinapasimple namin ang mga naunang arkitektura ng modelo ng deep sequence sa pamamagitan ng pagsasama-sama ng disenyo ng mga naunang arkitektura ng SSM (Dao, Fu, Saab, et al. 2023) kasama ang MLP block ng Transformers sa iisang bloke, na humahantong sa isang simple at homogenous na disenyo ng arkitektura (Mamba) na nagsasama ng mga piling puwang ng estado. Ang mga piling SSM, at sa pamamagitan ng extension ng arkitektura ng Mamba, ay ganap na paulit-ulit na mga modelo na may mga pangunahing katangian na ginagawang angkop ang mga ito bilang backbone ng mga pangkalahatang modelo ng pundasyon na tumatakbo sa mga pagkakasunud-sunod. (i) Mataas na kalidad: ang pagpili ay nagdudulot ng malakas na pagganap sa mga siksik na modalidad tulad ng wika at genomics. (ii) Mabilis na pagsasanay at hinuha: ang computation at memory scale ay linearly sa sequence length habang nagsasanay, at ang pag-unroll ng modelo nang autoregressive sa panahon ng inference ay nangangailangan lamang ng pare-parehong oras bawat hakbang dahil hindi ito nangangailangan ng cache ng mga nakaraang elemento. (iii) Mahabang konteksto: ang kalidad at kahusayan na magkasama ay nagbubunga ng mga pagpapahusay sa pagganap sa totoong data hanggang sa haba ng pagkakasunud-sunod na 1M.
Empirikal naming pinapatunayan ang potensyal ng Mamba bilang isang pangkalahatang sequence na backbone ng FM, sa parehong kalidad ng pretraining at pagganap ng gawain na partikular sa domain, sa ilang uri ng mga modalidad at setting:
• Synthetics. Sa mahahalagang gawaing gawa ng tao gaya ng pagkopya at mga induction head na iminungkahi bilang susi sa malalaking modelo ng wika, hindi lang madaling niresolba ng Mamba ang mga ito ngunit nagagawa nitong i-extrapolate ang mga solusyon nang walang katapusan (>1M token).
• Audio at Genomics. Mamba out-perform ang mga naunang makabagong modelo gaya ng SaShiMi, Hyena, at Transformers sa pagmomodelo ng mga audio waveform at DNA sequence, kapwa sa kalidad ng pretraining at downstream na sukatan (hal. pagbabawas ng FID sa isang mapaghamong dataset ng pagbuo ng pagsasalita ng higit sa kalahati). Sa parehong mga setting, ang pagganap nito ay bumubuti nang may mas mahabang konteksto hanggang sa milyong-haba na mga pagkakasunud-sunod.
• Pagmomodelo ng Wika. Ang Mamba ay ang unang linear-time na sequence model na tunay na nakakamit ng Transformer-kalidad na pagganap, kapwa sa pretraining perplexity at downstream na mga pagsusuri. Sa pag-scale ng mga batas hanggang sa 1B na mga parameter, ipinapakita namin na ang Mamba ay lumampas sa pagganap ng isang malaking hanay ng mga baseline, kabilang ang napakalakas na modernong mga recipe ng pagsasanay ng Transformer batay sa LLaMa (Touvron et al. 2023). Ang aming modelo ng wikang Mamba ay may 5x na henerasyong throughput kumpara sa Mga Transformer na may katulad na laki, at ang kalidad ng Mamba-3B ay tumutugma sa mga Transformer na dalawang beses ang laki nito (hal. 4 na puntos na mas mataas sa avg. sa common sense na pangangatwiran kumpara sa Pythia-3B at higit pa sa Pythia-7B).
Ang code ng modelo at mga pre-trained na checkpoint ay open-sourced sa https://github.com/state-spaces/mamba.
2 Mga Modelo ng Kalawakan ng Estado
Mga Pangkalahatang Modelo ng Kalawakan ng Estado. Pansinin namin na ang terminong modelo ng espasyo ng estado ay may napakalawak na kahulugan na kumakatawan lamang sa paniwala ng anumang paulit-ulit na proseso na may nakatagong estado. Ito ay ginamit upang sumangguni sa maraming magkakaibang konsepto sa iba't ibang disiplina, kabilang ang mga proseso ng desisyon ng Markov (MDP) (reinforcement learning (Hafner et al. 2020)), dynamic causal modeling (DCM) (computational neuroscience (Friston, Harrison, and Penny 2003)), Kalman filters (controls (60K) linear dynamical system (LDS) (machine learning), at paulit-ulit (at minsan convolutional) na mga modelo sa pangkalahatan (deep learning).
Sa buong papel na ito ginagamit namin ang terminong "SSM" para eksklusibong sumangguni sa klase ng mga structured na SSM o S4 na modelo (Gu, Goel, at Ré 2022; Gu, Gupta, et al. 2022; Gupta, Gu, at Berant 2022; Hasani et al. 2023; Ma et al., Warringman 2023, at Linderman 2; at Linderman 2; at Smith ang mga terminong ito; salitan. Para sa kaginhawahan, maaari rin naming isama ang mga derivatives ng naturang mga modelo, tulad ng mga tumutuon sa alinman sa linear-recurrence o global-convolution viewpoints (Y. Li et al. 2023; Orvieto et al. 2023; Poli et al. 2023), at linawin ang mga nuances kung kinakailangan.
Mga Arkitektura ng SSM. Ang mga SSM ay mga standalone sequence transformation na maaaring isama sa mga end-to-end na neural network architecture. (Minsan, tinatawag din namin ang mga arkitektura ng SSM na SSNN, na kung saan ay sa mga layer ng SSM dahil ang mga CNN ay sa mga linear convolution layer.) Tinatalakay namin ang ilan sa mga pinakakilalang arkitektura ng SSM, na marami sa mga ito ay magsisilbi rin bilang aming mga pangunahing baseline.
• Ang linear na atensyon (Katharopoulos et al. 2020) ay isang pagtatantya ng self-attention na kinasasangkutan ng pag-ulit na maaaring tingnan bilang isang degenerate linear SSM.
• Ang H3 (Dao, Fu, Saab, et al. 2023) ay nag-generalize ng pag-ulit na ito upang gamitin ang S4; maaari itong tingnan bilang isang arkitektura na may isang SSM na na-sandwich ng dalawang gated na koneksyon (Larawan 3). Naglalagay din ang H3 ng karaniwang lokal na convolution, na binabalangkas nila bilang shift-SSM, bago ang pangunahing layer ng SSM.
• Ginagamit ng Hyena (Poli et al. 2023) ang parehong arkitektura gaya ng H3 ngunit pinapalitan ang S4 layer ng MLP-parameterized global convolution (Romero et al. 2021).
• Ang RetNet (Y. Sun et al. 2023) ay nagdaragdag ng karagdagang gate sa arkitektura at gumagamit ng mas simpleng SSM, na nagbibigay-daan sa alternatibong parallelizable computation path, gamit ang isang variant ng multi-head attention (MHA) sa halip na mga convolution.
• Ang RWKV (B. Peng et al. 2023) ay isang kamakailang RNN na idinisenyo para sa pagmomodelo ng wika batay sa isa pang linear na pagtatantya ng atensyon (attention-free Transformer (S. Zhai et al. 2021)). Ang pangunahing mekanismo ng "WKV" nito ay kinabibilangan ng mga pag-ulit ng LTI at maaaring tingnan bilang ratio ng dalawang SSM.
Ang iba pang malapit na nauugnay na mga SSM at arkitektura ay tinalakay pa sa isang pinalawig na kaugnay na gawain (Appendix B). Binibigyang-diin namin sa partikular ang S5 (Smith, Warrington, and Linderman 2023), QRNN (Bradbury et al. 2016), at SRU (Lei et al. 2017), na tinitingnan namin bilang ang pinaka malapit na nauugnay na mga pamamaraan sa aming core selective SSM.
Ang papel na ito ay makukuha sa arxiv sa ilalim ng CC BY 4.0 DEED na lisensya.
