Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

शोध पत्र · arXiv 2312.00752

Mamba: सेलेक्टिव स्टेट स्पेसेज के साथ लीनियर-टाइम सीक्वेंस मॉडलिंग

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

प्रकाशित 2023-12-01cs.LGcs.AI arXiv PDF

अनुवाद और AI के साथ पढ़ें

तीन पंक्तियों में सारांश

🎯 Transformers सीक्वेंस की लंबाई के साथ क्वाड्रैटिक रूप से स्केल करते हैं और ऑटोरेग्रेसिव इन्फरेंस के दौरान बड़े KV कैशे की आवश्यकता होती है, जबकि पुराने सब-क्वाड्रैटिक मॉडल्स जैसे LTI स्टेट स्पेस मॉडल्स भाषा मॉडलिंग जैसे डिस्क्रीट, कंटेंट-रीजनिंग कार्यों पर संघर्ष करते हैं।
🔧 जानकारी को गतिशील रूप से फ़िल्टर या प्रोपेगेट करने के लिए SSM पैरामीटर्स को इनपुट-निर्भर (सेलेक्टिव) बनाता है, और GPU HBM में विस्तारित हिडन स्टेट को मैटेरियलाइज करने से बचने के लिए इसे एक हार्डवेयर-अवेयर फ्यूज्ड पैरेलल एसोसिएटिव स्कैन और रीकंप्यूटेशन एल्गोरिदम के साथ जोड़ता है।
📈 Mamba-3B डाउनस्ट्रीम ज़ीरो-शॉट कार्यों पर अपने से दोगुने आकार के Transformers की बराबरी करता है या उनसे आगे निकल जाता है (उदा. 63.3% औसत बनाम Pythia-2.8B के लिए 59.1% और Pythia-6.9B के लिए 61.7%) और 5× तक अधिक इन्फरेंस थ्रूपुट प्रदान करता है।

सार

फाउंडेशन मॉडल्स, जो अब डीप लर्निंग के अधिकांश रोमांचक अनुप्रयोगों को संचालित कर रहे हैं, लगभग पूरी तरह से Transformer आर्किटेक्चर और इसके कोर अटेंशन मॉड्यूल पर आधारित हैं। लंबी अनुक्रमों पर Transformers की कम्प्यूटेशनल अक्षमता को दूर करने के लिए लीनियर अटेंशन, गेटेड कनवल्शन और रिकरेंट मॉडल्स, और स्ट्रक्चर्ड स्टेट स्पेस मॉडल्स (SSMs) जैसे कई सब-क्वाड्रैटिक-टाइम आर्किटेक्चर विकसित किए गए हैं, लेकिन वे भाषा जैसी महत्वपूर्ण मोडैलिटीज पर अटेंशन जितना अच्छा प्रदर्शन नहीं कर पाए हैं। हमने पहचाना कि ऐसे मॉडल्स की एक प्रमुख कमजोरी कंटेंट-आधारित रीजनिंग करने में उनकी अक्षमता है, और हमने कई सुधार किए हैं। पहला, केवल SSM पैरामीटर्स को इनपुट का फंक्शन बनाने से असतत (डिस्क्रीट) मोडैलिटीज के साथ उनकी कमजोरी दूर हो जाती है, जिससे मॉडल मौजूदा टोकन के आधार पर सीक्वेंस लंबाई के आयाम में जानकारी को चुनिंदा रूप से प्रोपेगेट या रीसेट (भूलने) की अनुमति देता है। दूसरा, भले ही यह बदलाव कुशल कनवल्शन के उपयोग को रोकता है, हम रिकरेंट मोड में एक हार्डवेयर-अवेयर पैरेलल एल्गोरिदम डिजाइन करते हैं। हम इन सेलेक्टिव SSMs को अटेंशन या MLP ब्लॉक्स (Mamba) के बिना एक सरलीकृत एंड-टू-एंड न्यूरल नेटवर्क आर्किटेक्चर में एकीकृत करते हैं। Mamba तेज़ इन्फरेंस (Transformers की तुलना में 5 $\times$ अधिक थ्रूपुट) और सीक्वेंस की लंबाई में लीनियर स्केलिंग का लाभ उठाता है, और वास्तविक डेटा पर इसका प्रदर्शन 10 लाख तक की लंबाई वाले सीक्वेंस तक बेहतर होता जाता है। एक सामान्य सीक्वेंस मॉडल बैकबोन के रूप में, Mamba भाषा, ऑडियो और जीनोमिक्स जैसी कई मोडैलिटीज में स्टेट-ऑफ-द-आर्ट प्रदर्शन प्राप्त करता है। भाषा मॉडलिंग पर, हमारा Mamba-3B मॉडल समान आकार के Transformers से बेहतर प्रदर्शन करता है और प्री-ट्रेनिंग और डाउनस्ट्रीम मूल्यांकन दोनों में अपने से दोगुने आकार के Transformers की बराबरी करता है।

मूल सार (अंग्रेज़ी)

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.

पढ़ने की गाइड

समस्या और प्रेरणा

फाउंडेशन मॉडल्स मुख्य रूप से Transformer आर्किटेक्चर और सेल्फ-अटेंशन पर निर्भर करते हैं। यद्यपि डेंस अटेंशन एक सीक्वेंस संदर्भ विंडो में जानकारी को प्रभावी ढंग से रूट करता है, यह दो मूलभूत सीमाओं को जन्म देता है: सीक्वेंस की लंबाई $L$ के संबंध में क्वाड्रैटिक कम्प्यूटेशनल लागत $O (L^{2})$ , और पिछले कीज और वैल्यूज (KV cache) को स्टोर करने की आवश्यकता, जिससे प्रति स्टेप लीनियर-टाइम इन्फरेंस और मेमोरी अड़चनें पैदा होती हैं।

पूर्व सब-क्वाड्रैटिक दृष्टिकोण—जैसे लीनियर अटेंशन, रिकरेंट आर्किटेक्चर, और S4 जैसे स्ट्रक्चर्ड स्टेट स्पेस मॉडल्स (SSMs)—समय में लीनियर या स्यूडो-लीनियर रूप से स्केल करते हैं। हालाँकि, स्ट्रक्चर्ड SSMs ऐतिहासिक रूप से टेक्स्ट जैसी असतत, जानकारी से भरपूर मोडैलिटीज पर Transformer की गुणवत्ता से मेल खाने में विफल रहे हैं। लेखकों ने इसके मूल कारण की पहचान की है: पूर्व SSMs लीनियर टाइम-इनवेरिएंट (LTI) हैं। चूँकि उनकी ट्रांज़िशन डायनेमिक्स इनपुट कंटेंट के आधार पर नहीं बदलती हैं, वे संदर्भ को चुनिंदा रूप से कंप्रेस नहीं कर सकते, अप्रासंगिक शोर को फ़िल्टर नहीं कर सकते, या विभिन्न अंतरालों में अवधारण (रिटेंशन) को नियंत्रित नहीं कर सकते।

मुख्य विचार

कंटेक्स्ट कम्प्रेशन के रूप में सेलेक्शन मैकेनिज्म: सीक्वेंस मॉडलिंग के लिए इतिहास को एक स्टेट में कंप्रेस करने की आवश्यकता होती है। पूर्व LTI SSMs फिक्स्ड ट्रांज़िशन पैरामीटर्स का उपयोग करते हैं, जिससे वे सभी टाइमस्टेप्स को समान रूप से ट्रीट करने के लिए मजबूर होते हैं। वर्तमान इनपुट टोकन के आधार पर SSM मैट्रिसेस को पैरामीटर करके, मॉडल को चुनिंदा रूप से जानकारी को फ़िल्टर करने या याद रखने के लिए आवश्यक कंटेंट-अवेयरनेस प्राप्त होती है।
नॉन-कन्वॉल्विंग बॉटलनेक पर काबू पाना: पैरामीटर्स को इनपुट-निर्भर बनाने से लीनियर रिकरेंस और ग्लोबल कनवल्शन के बीच गणितीय समानता समाप्त हो जाती है। कनवल्शन के बिना, साधारण रिकरेंस के लिए GPU की हाई-बैंडविड्थ मेमोरी (HBM) में $(B, L, D, N)$ आकार की एक विस्तारित लेटेंट स्टेट को मैटेरियलाइज करने की आवश्यकता होगी। Mamba डिस्क्रेटाइजेशन, पैरेलल एसोसिएटिव स्कैन और आउटपुट प्रोजेक्शन को तेज़ GPU ऑन-चिप SRAM में फ्यूज करके इसे हल करता है, जिससे HBM ट्रैफिक न्यूनतम रहता है।
एकीकृत आर्किटेक्चर: Mamba इंटरलीव्ड मल्टी-हेड अटेंशन (MHA) और MLP ब्लॉक्स को एक सिंगल होमोजेनियस ब्लॉक से बदल देता है जो H3 के गेटेड डिज़ाइन को SwiGLU-शैली के गेटेड MLP के साथ जोड़ता है।

विधि

कंटीन्यूअस स्टेट स्पेस मॉडल्स एक इम्प्लिसिट लेटेंट स्टेट $h (t) \in R^{N}$ के माध्यम से एक इनपुट $x (t) \in R$ को $y (t) \in R$ में मैप करते हैं: $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

स्टेप साइज $Δ$ के साथ ज़ीरो-ऑर्डर होल्ड (ZOH) डिस्क्रेटाइजेशन का उपयोग करके, कंटीन्यूअस सिस्टम डिस्क्रीट रिकरेंस पैरामीटर्स में बदल जाता है: $\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}$

द सेलेक्शन मैकेनिज्म (S6)

समय के साथ $(Δ, B, C)$ को स्थिर रखने के बजाय, Mamba उन्हें इनपुट-निर्भर बनाता है:

$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}))$ , जहाँ $τ_{Δ} = softplus$ और $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

यह क्लासिकल RNN गेटिंग को सामान्यीकृत करता है: जब $N = 1, A = - 1, B = 1$ , तो रिकरेंस सीधे गेटेड अपडेट $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ में सरल हो जाता है, जहाँ $g_{t} = σ (Linear (x_{t}))$ ।

हार्डवेयर-अवेयर स्टेट एक्सपेंशन

HBM में $(B, L, D, N)$ स्टेट को मैटेरियलाइज किए बिना उच्च थ्रूपुट बनाए रखने के लिए:

पैरामीटर्स $(Δ, A, B, C)$ को सीधे HBM से ऑन-चिप SRAM में लोड किया जाता है।
$(\overline{A}, \overline{B})$ में डिस्क्रेटाइजेशन और एक कुशल पैरेलल एसोसिएटिव स्कैन पूरी तरह से SRAM के अंदर ही कंप्यूट किया जाता है।
HBM में वापस लिखने से पहले लेटेंट स्टेट को $C$ के साथ $(B, L, D)$ आकार में अनुबंधित (कॉन्ट्रैक्ट) किया जाता है।
बैकवर्ड पास के दौरान इंटरमीडिएट स्टेट्स को HBM में सेव करने के बजाय SRAM में फिर से कंप्यूट किया जाता है।

परिणाम

सिंथेटिक एक्स्ट्रापोलेशन: Induction Heads टास्क पर, 256 सीक्वेंस लंबाई पर प्रशिक्षित Mamba $2^{20} = 1, 048, 576$ लंबाई ( $4000 \times$ सीक्वेंस लंबाई) तक पूरी तरह से एक्स्ट्रापोलेट (100% सटीकता) करता है, जबकि अटेंशन बेसलाइन $2^{14}$ के बाद विफल हो जाते हैं या OOM हो जाते हैं। Selective Copying पर, Mamba H3 के 57.0% और S4 के 18.3% की तुलना में 99.8% सटीकता तक पहुँचता है।
लैंग्वेज मॉडलिंग: Pile पर 1.3B पैरामीटर्स तक Chinchilla स्केलिंग के तहत मूल्यांकित, Mamba Transformer++ (rotary embeddings, SwiGLU, RMSNorm) से मेल खाता है या उससे बेहतर प्रदर्शन करता है। डाउनस्ट्रीम ज़ीरो-शॉट कार्यों (उदा., LAMBADA, PIQA, HellaSwag, WinoGrande, ARC) पर, Mamba-2.8B 63.3% की औसत सटीकता प्राप्त करता है, जो Pythia-2.8B (59.1%), Pythia-6.9B (61.7%) से बेहतर है, और 6B–7B पैरामीटर मॉडल के बराबर या उससे बेहतर है।
जीनोमिक्स और ऑडियो: 1M ( $2^{20}$ ) की सीक्वेंस लंबाई पर ग्रेट एप्स DNA क्लासिफिकेशन में, Mamba HyenaDNA के 54.87% की तुलना में 71.67% (1.4M पैरामीटर्स) और 81.31% (7M पैरामीटर्स) सटीकता प्राप्त करता है। SC09 स्पीच जेनरेशन पर, Mamba FID को 1.99 (SaShiMi) से घटाकर 0.94 (6.1M पैरामीटर्स) और 0.67 (24.3M पैरामीटर्स) कर देता है।
दक्षता: फ्यूज्ड स्कैन कर्नेल स्टैंडर्ड PyTorch स्कैन कार्यान्वयन की तुलना में 20–40× तेज़ी से निष्पादित होता है और 32K संदर्भ पर FlashAttention-2 की तुलना में 7× तक तेज़ी से चलता है। इन्फरेंस थ्रूपुट तुलनीय Transformers की तुलना में 5× तक अधिक है।

सीमाएँ और खुले प्रश्न

कंटीन्यूअस सिग्नल प्रदर्शन: जबकि इनपुट-निर्भर सेलेक्शन असतत तर्क (टेक्स्ट, DNA) में सुधार करता है, रॉ ऑडियो वेवफॉर्म प्री-ट्रेनिंग पर किए गए एब्लेशन से पता चलता है कि शुद्ध लीनियर टाइम-इनवेरिएंट (LTI) मॉडल्स (जैसे मानक S4) समान रूप से सैंपल किए गए, कंटीन्यूअस डेटा पर बेहतर प्रदर्शन करते हैं।
मॉडल का पैमाना: मूल्यांकित मॉडल्स लगभग 3B पैरामीटर्स और 300B प्री-ट्रेनिंग टोकन तक सीमित हैं; 7B–70B+ और उससे अधिक के पैमाने पर अनुभवजन्य प्रदर्शन के लिए आगे की जांच की आवश्यकता है।
डाउनस्ट्रीम इकोसिस्टम: Transformers के पास प्रॉम्प्ट इंजीनियरिंग, इन-कॉन्टेक्स्ट मैनिपुलेशन, KV-कैशे क्वांटाइजेशन और एडॉप्टर ट्यूनिंग के लिए व्यापक टूल्स हैं। SSMs इन इंटरेक्शन मोड्स में कैसे अनुकूलित होते हैं, यह समझना अभी भी एक सक्रिय शोध का विषय है।

मुख्य शब्द

Structured State Space Models (SSMs): कंटीन्यूअस लीनियर डायनेमिकल सिस्टम्स से प्राप्त डीप सीक्वेंस मॉडल्स का एक वर्ग जो एक इम्प्लिसिट लेटेंट स्टेट के माध्यम से 1D इनपुट को मैप करता है।
Linear Time Invariance (LTI): एक गणितीय विशेषता जहाँ सिस्टम ट्रांज़िशन पैरामीटर्स सभी टाइमस्टेप्स में स्थिर रहते हैं, जो तेज़ ग्लोबल कनवल्शन अभ्यावेदन को सक्षम बनाता है लेकिन कंटेंट-निर्भर रीजनिंग को सीमित करता है।
Selection Mechanism: SSM ट्रांज़िशन डायनेमिक्स (Δ, B, C) को इनपुट टोकन के फंक्शन्स के रूप में पैरामीटर करना, जिससे मॉडल को सीक्वेंस लंबाई के साथ संदर्भ को गतिशील रूप से प्रोपेगेट करने, रीसेट करने या अनदेखा करने की अनुमति मिलती है।
S6: कनवल्शन के बजाय एक एसोसिएटिव पैरेलल स्कैन के माध्यम से कंप्यूट किए गए सेलेक्टिव स्ट्रक्चर्ड स्टेट स्पेस मॉडल्स का संक्षिप्त नाम।
Parallel Associative Scan: एक कार्य-कुशल प्रिफिक्स कंप्यूटेशन एल्गोरिदम जो एक्सेलेरेटर हार्डवेयर पर समानांतर में सीक्वेंस लंबाई में लीनियर रिकरेंस का मूल्यांकन करता है।
Hardware-Aware State Expansion: एक GPU कर्नेल फ्यूजन और रीकंप्यूटेशन रणनीति जो हाई-बैंडविड्थ मेमोरी (HBM) में बड़े इंटरमीडिएट स्टेट्स के राउंड-ट्रिपिंग से बचने के लिए तेज़ ऑन-चिप SRAM में डिस्क्रेटाइजेशन और स्कैन करती है।
Induction Heads: एक सिंथेटिक एसोसिएटिव रिकॉल मैकेनिज्म जो यह मूल्यांकन करता है कि क्या कोई मॉडल अपने पिछले संदर्भ टोकन के बाद पहले देखे गए टोकन को पुन: उत्पन्न कर सकता है।
Mamba Block: एक सरलीकृत, सजातीय सीक्वेंस मॉडल ब्लॉक जो अलग अटेंशन या MLP परतों के बिना इनपुट प्रोजेक्शन, 1D कनवल्शन, सेलेक्टिव SSMs और मल्टीप्लिकेटिव गेटिंग को एकीकृत करता है।

प्रश्न

तेज़ FFT कनवल्शन का उपयोग करके सेलेक्टिव SSMs को प्रशिक्षित क्यों नहीं किया जा सकता है?

तेज़ ग्लोबल FFT कनवल्शन के लिए सिस्टम डायनेमिक्स (मैट्रिसेस A, B, और C) का समय के साथ स्थिर होना आवश्यक है (लीनियर टाइम इनवेरिएंस)। चूँकि Mamba प्रत्येक टाइमस्टेप पर B, C, और Δ को इनपुट टोकन का फंक्शन बनाता है, इसलिए कनवल्शन कर्नेल हर स्थिति में बदलता है, जिससे कनवल्शनल समानता टूट जाती है और इसके बजाय रिकरेंट पैरेलल स्कैन की आवश्यकता होती है।

प्रशिक्षण के दौरान Mamba रिकरेंट स्टेट एक्सपेंशन की अत्यधिक मेमोरी खपत से कैसे बचता है?

स्टेट एक्सपेंशन आयाम N के साथ नेइव रिकरेंस के लिए GPU HBM में (B, L, D, N) आकार की एक इंटरमीडिएट स्टेट को मैटेरियलाइज करने की आवश्यकता होती है। Mamba केवल लो-डायमेंशनल इनपुट्स को SRAM में लोड करने के लिए कर्नेल फ्यूजन का उपयोग करता है, तेज़ SRAM के भीतर डिस्क्रेटाइजेशन और पैरेलल स्कैन की गणना करता है, केवल कॉन्ट्रैक्टेड (B, L, D) आउटपुट को वापस HBM में लिखता है, और बैकप्रॉपैगैशन के दौरान इंटरमीडिएट एक्टिवेशन्स को फिर से कंप्यूट करता है।

Mamba के Δ पैरामीटर और रिकरेंट गेटिंग के बीच क्या संबंध है?

स्टेप साइज पैरामीटर Δ एक सामान्यीकृत गेटिंग तंत्र के रूप में कार्य करता है। जब N=1 और A=-1 के साथ ज़ीरो-ऑर्डर होल्ड (ZOH) के तहत डिस्क्रेटाइज़ किया जाता है, तो सेलेक्टिव रिकरेंस गणितीय रूप से एक मानक RNN गेट में सरल हो जाता है: बड़ा Δ मॉडल को वर्तमान इनपुट पर ध्यान केंद्रित करने और स्थिति को रीसेट करने के लिए प्रेरित करता है, जबकि छोटा Δ ऐतिहासिक स्थिति को बनाए रखता है और इनपुट को अनदेखा करता है।

क्या Mamba जनरेशन के दौरान KV कैशे का उपयोग करता है?

नहीं। चूँकि Mamba जेनरेशन के दौरान एक ऑटोरेग्रेसिव स्टेट स्पेस रिकरेंस के रूप में काम करता है, यह सभी पिछले टोकन के लिए कीज और वैल्यूज को कैश करने के बजाय केवल एक निश्चित आकार की हिडन स्टेट (प्रति लेयर D × N आयाम) को बनाए रखता है, जिसके परिणामस्वरूप प्रति स्टेप कॉन्स्टेंट-टाइम O(1) इन्फरेंस होता है।

इस पृष्ठ की व्याख्या AI ने शोध पत्र से लिखी है और इसमें गलतियाँ हो सकती हैं। विवरण के लिए मूल शोध पत्र पढ़ें।

अनुवाद और AI के साथ पढ़ें