Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

Makalah · arXiv 2312.00752

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

Terbit 2023-12-01cs.LGcs.AI arXiv PDF

Baca dengan terjemahan dan AI

Ringkasan tiga baris

🎯 Transformer berskala kuadratik terhadap panjang sekuens dan membutuhkan KV cache yang besar selama inferensi autoregresif, sementara model subkuadratik sebelumnya seperti LTI state space model kesulitan pada tugas penalaran konten yang diskret seperti pemodelan bahasa.
🔧 Menjadikan parameter SSM bergantung pada masukan (selektif) untuk menyaring atau menyebarkan informasi secara dinamis, dan memadukannya dengan fusi parallel associative scan serta algoritma komputasi ulang yang sadar perangkat keras untuk menghindari materialisasi hidden state yang diperluas di HBM GPU.
📈 Mamba-3B menyamai atau melampaui Transformer berukuran dua kali lipatnya pada tugas downstream zero-shot (misalnya, rata-rata 63,3% vs. 59,1% untuk Pythia-2.8B dan 61,7% untuk Pythia-6.9B) sekaligus memberikan throughput inferensi hingga 5× lebih tinggi.

Abstrak

Model fondasi, yang kini mentenagai sebagian besar aplikasi menarik dalam deep learning, hampir secara universal didasarkan pada arsitektur Transformer dan modul perhatian intinya. Banyak arsitektur berwaktu subkuadratik seperti linear attention, konvolusi berpagar dan model rekuren, serta structured state space models (SSM) telah dikembangkan untuk mengatasi ketidakefisienan komputasi Transformer pada sekuens panjang, namun kinerjanya belum sebaik attention pada modalitas penting seperti bahasa. Kami mengidentifikasi bahwa kelemahan utama dari model-model tersebut adalah ketidakmampuannya dalam melakukan penalaran berbasis konten, dan kami membuat beberapa perbaikan. Pertama, membiarkan parameter SSM menjadi fungsi dari masukan secara sederhana dapat mengatasi kelemahannya pada modalitas diskret, memungkinkan model untuk secara selektif menyebarkan atau melupakan informasi di sepanjang dimensi panjang sekuens bergantung pada token saat ini. Kedua, meskipun perubahan ini mencegah penggunaan konvolusi yang efisien, kami merancang algoritma paralel sadar perangkat keras dalam mode rekuren. Kami mengintegrasikan selective SSM ini ke dalam arsitektur jaringan saraf ujung-ke-ujung (end-to-end) yang disederhanakan tanpa blok attention atau bahkan MLP (Mamba). Mamba menikmati inferensi yang cepat (throughput 5 $\times$ lebih tinggi daripada Transformer) dan penskalaan linear dalam panjang sekuens, serta kinerjanya meningkat pada data nyata hingga panjang sekuens jutaan. Sebagai tulang punggung model sekuens umum, Mamba mencapai performa mutakhir di berbagai modalitas seperti bahasa, audio, dan genomika. Pada pemodelan bahasa, model Mamba-3B kami mengungguli Transformer berukuran sama dan menyamai Transformer dua kali ukurannya, baik dalam pra-pelatihan maupun evaluasi hilir.

Abstrak asli (bahasa Inggris)

Foundation models, now powering most of the exciting applications in deep learning, are almost universally based on the Transformer architecture and its core attention module. Many subquadratic-time architectures such as linear attention, gated convolution and recurrent models, and structured state space models (SSMs) have been developed to address Transformers' computational inefficiency on long sequences, but they have not performed as well as attention on important modalities such as language. We identify that a key weakness of such models is their inability to perform content-based reasoning, and make several improvements. First, simply letting the SSM parameters be functions of the input addresses their weakness with discrete modalities, allowing the model to selectively propagate or forget information along the sequence length dimension depending on the current token. Second, even though this change prevents the use of efficient convolutions, we design a hardware-aware parallel algorithm in recurrent mode. We integrate these selective SSMs into a simplified end-to-end neural network architecture without attention or even MLP blocks (Mamba). Mamba enjoys fast inference (5 $\times$ higher throughput than Transformers) and linear scaling in sequence length, and its performance improves on real data up to million-length sequences. As a general sequence model backbone, Mamba achieves state-of-the-art performance across several modalities such as language, audio, and genomics. On language modeling, our Mamba-3B model outperforms Transformers of the same size and matches Transformers twice its size, both in pretraining and downstream evaluation.

Panduan membaca

Masalah dan motivasi

Model fondasi sangat bergantung pada arsitektur Transformer dan self-attention. Meskipun dense attention merutekan informasi secara efektif di seluruh jendela konteks sekuens, metode ini memiliki dua keterbatasan mendasar: biaya komputasi kuadratik $O (L^{2})$ terhadap panjang sekuens $L$ , dan keharusan menyimpan key dan value masa lalu (KV cache), yang menyebabkan inferensi per langkah membutuhkan waktu linear serta menjadi hambatan memori.

Pendekatan subkuadratik sebelumnya—seperti linear attention, arsitektur rekuren, dan structured state space models (SSM) seperti S4—berskala linear atau pseudo-linear dalam waktu. Namun, SSM terstruktur secara historis gagal menyamai kualitas Transformer pada modalitas padat informasi yang diskret seperti teks. Para penulis mengidentifikasi akar masalahnya: SSM terdahulu bersifat linear time-invariant (LTI). Karena dinamika transisinya tidak berubah berdasarkan konten masukan, model tersebut tidak dapat mengompresi konteks secara selektif, menyaring derau yang tidak relevan, atau memodulasi retensi melintasi jarak spasi yang bervariasi.

Ide utama

Mekanisme Seleksi sebagai Kompresi Konteks: Pemodelan sekuens membutuhkan kompresi riwayat ke dalam sebuah state. SSM LTI terdahulu menggunakan parameter transisi tetap, memaksa mereka memperlakukan semua langkah waktu secara seragam. Dengan memparameterisasi matriks SSM berdasarkan token masukan saat ini, model memperoleh kesadaran konten yang diperlukan untuk menyaring atau mengingat informasi secara selektif.
Mengatasi Hambatan Non-Konvolusi: Menjadikan parameter bergantung pada masukan memutus kesetaraan matematis antara rekurensi linear dan konvolusi global. Tanpa konvolusi, rekurensi naif akan membutuhkan materialisasi latent state yang diperluas berukuran $(B, L, D, N)$ dalam high-bandwidth memory (HBM) GPU. Mamba memecahkan masalah ini dengan memadukan diskritisasi, parallel associative scan, dan proyeksi keluaran langsung ke dalam SRAM on-chip GPU yang cepat, meminimalkan lalu lintas HBM.
Arsitektur Terpadu: Mamba menggantikan blok Multi-Head Attention (MHA) dan MLP yang berselang-seling dengan satu blok homogen tunggal yang menggabungkan desain berpagar (gated) milik H3 dengan MLP berpagar bergaya SwiGLU.

Metode

Model continuous state space memetakan masukan $x (t) \in R$ ke $y (t) \in R$ melalui latent state implisit $h (t) \in R^{N}$ : $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

Menggunakan diskritisasi zero-order hold (ZOH) dengan ukuran langkah $Δ$ , sistem kontinu berubah menjadi parameter rekurensi diskret: $\overline{A} = exp (Δ A), \overline{B} = (Δ A)^{- 1} (exp (Δ A) - I) \cdot Δ B$ $h_{t} = \overline{A} h_{t - 1} + \overline{B} x_{t}, y_{t} = C h_{t}$

Mekanisme Seleksi (S6)

Alih-alih menjaga $(Δ, B, C)$ konstan sepanjang waktu, Mamba membuatnya bergantung pada masukan:

$B_{t} = s_{B} (x_{t}) = Linear_{N} (x_{t})$
$C_{t} = s_{C} (x_{t}) = Linear_{N} (x_{t})$
$Δ_{t} = τ_{Δ} (Parameter + s_{Δ} (x_{t}))$ , dengan $τ_{Δ} = softplus$ dan $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

Ini menggeneralisasi gating RNN klasik: ketika $N = 1, A = - 1, B = 1$ , rekurensi langsung disederhanakan menjadi pembaruan berpagar $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ , di mana $g_{t} = σ (Linear (x_{t}))$ .

Ekspansi State yang Sadar Perangkat Keras

Untuk mempertahankan throughput tinggi tanpa mematerialisasi state $(B, L, D, N)$ dalam HBM:

Parameter $(Δ, A, B, C)$ dimuat dari HBM langsung ke SRAM on-chip.
Diskritisasi menjadi $(\overline{A}, \overline{B})$ dan parallel associative scan yang efisien secara komputasi dihitung seluruhnya di dalam SRAM.
Latent state dikontraksikan dengan $C$ menjadi bentuk $(B, L, D)$ sebelum ditulis kembali ke HBM.
State perantara dikomputasi ulang di SRAM selama backward pass alih-alih disimpan di HBM.

Hasil

Ekstrapolasi Sintetis: Pada tugas Induction Heads, Mamba yang dilatih pada panjang sekuens 256 dapat berekstrapolasi secara sempurna (akurasi 100%) hingga panjang $2^{20} = 1, 048, 576$ ( $4000 \times$ panjang sekuens), sedangkan baseline attention gagal atau mengalami OOM di atas $2^{14}$ . Pada Selective Copying, Mamba mencapai akurasi 99,8% dibandingkan dengan 57,0% untuk H3 dan 18,3% untuk S4.
Pemodelan Bahasa: Dievaluasi di bawah penskalaan Chinchilla hingga 1,3B parameter pada the Pile, Mamba menyamai atau mengungguli Transformer++ (rotary embeddings, SwiGLU, RMSNorm). Pada tugas downstream zero-shot (misalnya LAMBADA, PIQA, HellaSwag, WinoGrande, ARC), Mamba-2.8B memperoleh akurasi rata-rata 63,3%, mengungguli Pythia-2.8B (59,1%), Pythia-6.9B (61,7%), serta menyamai atau melampaui model berparameter 6B–7B.
Genomika dan Audio: Pada klasifikasi DNA Kera Besar dengan panjang sekuens 1M ( $2^{20}$ ), Mamba mencapai akurasi 71,67% (1,4M parameter) dan 81,31% (7M parameter) dibandingkan dengan 54,87% untuk HyenaDNA. Pada pembangkitan wicara SC09, Mamba memangkas FID dari 1,99 (SaShiMi) menjadi 0,94 (6,1M parameter) dan 0,67 (24,3M parameter).
Efisiensi: Fused scan kernel dieksekusi 20–40× lebih cepat daripada implementasi scan bawaan PyTorch dan berjalan hingga 7× lebih cepat daripada FlashAttention-2 pada konteks 32K. Throughput inferensi mencapai hingga 5× lebih tinggi daripada Transformer yang sebanding.

Keterbatasan dan pertanyaan terbuka

Kinerja Sinyal Kontinu: Meskipun seleksi yang bergantung pada masukan meningkatkan penalaran diskret (teks, DNA), studi ablasi pada pra-pelatihan bentuk gelombang audio mentah menunjukkan bahwa model yang murni linear time-invariant (LTI) (seperti S4 standar) berkinerja lebih baik pada data kontinu yang disampel secara seragam.
Skala Model: Model yang dievaluasi dibatasi sekitar 3B parameter dan 300B token pra-pelatihan; kinerja empiris pada skala 7B–70B+ dan seterusnya memerlukan penyelidikan lebih lanjut.
Ekosistem Hilir: Transformer memiliki perkakas yang luas untuk prompt engineering, manipulasi in-context, kuantisasi KV-cache, dan adapter tuning. Memahami bagaimana SSM beradaptasi di seluruh mode interaksi ini masih menjadi eksplorasi yang aktif.

Istilah kunci

Structured State Space Models (SSMs): Kelas model sekuens mendalam yang diturunkan dari sistem dinamik linear kontinu yang memetakan masukan 1D melalui latent state implisit.
Linear Time Invariance (LTI): Sifat matematis di mana parameter transisi sistem tetap statis di semua langkah waktu, memungkinkan representasi konvolusi global yang cepat tetapi membatasi penalaran yang bergantung pada konten.
Selection Mechanism: Memparameterisasi dinamika transisi SSM (Δ, B, C) sebagai fungsi dari token masukan, memungkinkan model menyebarkan, mengatur ulang, atau mengabaikan konteks secara dinamis di sepanjang panjang sekuens.
S6: Singkatan untuk selective structured state space models yang dihitung melalui parallel associative scan alih-alih konvolusi.
Parallel Associative Scan: Algoritma komputasi prefiks yang efisien secara beban kerja yang mengevaluasi rekurensi linear di seluruh panjang sekuens secara paralel pada perangkat keras akselerator.
Hardware-Aware State Expansion: Strategi fusi kernel GPU dan komputasi ulang yang melakukan diskritisasi dan scan dalam SRAM on-chip yang cepat untuk menghindari transfer bolak-balik state perantara yang besar ke high-bandwidth memory (HBM).
Induction Heads: Mekanisme pemanggilan asosiatif sintetis yang mengevaluasi apakah suatu model dapat mereproduksi token yang diamati sebelumnya mengikuti token konteks pendahulunya.
Mamba Block: Blok model sekuens homogen yang disederhanakan yang mengintegrasikan proyeksi masukan, konvolusi 1D, selective SSM, dan multiplicative gating tanpa lapisan attention atau MLP terpisah.

Pertanyaan

Mengapa selective SSM tidak dapat dilatih menggunakan konvolusi FFT cepat?

Konvolusi FFT global yang cepat membutuhkan dinamika sistem (matriks A, B, dan C) yang konstan sepanjang waktu (linear time invariance). Karena Mamba membuat B, C, dan Δ menjadi fungsi dari token masukan pada setiap langkah waktu, kernel konvolusi bervariasi pada setiap posisi, merusak kesetaraan konvolusional dan mengharuskan penggunaan recurrent parallel scan sebagai gantinya.

Bagaimana cara Mamba menghindari lonjakan memori dari ekspansi state rekuren selama pelatihan?

Rekurensi naif dengan dimensi ekspansi state N memerlukan materialisasi state perantara berukuran (B, L, D, N) di HBM GPU. Mamba menggunakan fusi kernel untuk memuat masukan berdimensi rendah saja ke dalam SRAM, menghitung diskritisasi dan parallel scan di dalam SRAM yang cepat, hanya menuliskan keluaran yang dikontraksikan (B, L, D) kembali ke HBM, dan mengomputasi ulang aktivasi perantara selama perambatan balik (backpropagation).

Apa hubungan antara parameter Δ Mamba dan gating rekuren?

Parameter ukuran langkah Δ bertindak sebagai mekanisme gating yang digeneralisasi. Ketika didiskritisasi di bawah zero-order hold (ZOH) dengan N=1 dan A=-1, rekurensi selektif secara matematis tereduksi menjadi gerbang RNN standar: nilai Δ yang besar membuat model berfokus pada masukan saat ini dan mengatur ulang state, sedangkan nilai Δ yang kecil mempertahankan state riwayat dan mengabaikan masukan.

Apakah Mamba menggunakan KV cache selama pembuatan teks (generasi)?

Tidak. Karena Mamba beroperasi sebagai rekurensi state space autoregresif selama proses generasi, model ini hanya mempertahankan hidden state berukuran tetap (dimensi D × N per lapisan) alih-alih menyimpan cache key dan value untuk semua token masa lalu, menghasilkan inferensi dengan waktu konstan O(1) per langkah.

Ulasan di halaman ini ditulis oleh AI dari makalah dan mungkin mengandung kesalahan. Baca makalahnya untuk detail.

Baca dengan terjemahan dan AI