Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

الأوراق البحثية · arXiv 2312.00752

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

تاريخ النشر 2023-12-01cs.LGcs.AI arXiv PDF

اقرأ مع الترجمة والذكاء الاصطناعي

الملخص في ثلاثة أسطر

🎯 تتوسع نماذج Transformers تربيعياً مع طول التسلسل وتتطلب مخابئ KV ضخمة أثناء الاستدلال الذاتي، بينما تواجه النماذج السابقة دون التربيعية، كنماذج فضاء الحالة LTI، صعوبات في مهام الاستدلال على المحتوى المنفصل مثل نمذجة اللغة.
🔧 يجعل معلمات SSM معتمدة على المدخلات (انتقائية) لترشيح المعلومات أو تمريرها ديناميكياً، ويدمج هذا مع مسح تجميعي متوازٍ مدمج وخوارزمية إعادة حساب تراعي العتاد لتجنب تخزين الحالة الخفية الموسعة في ذاكرة HBM الخاصة بـ GPU.
📈 يطابق Mamba-3B أو يتجاوز أداء نماذج Transformers الأكبر منه حجماً بمرتين في مهام الصفر المعرفي (zero-shot) اللاحقة (على سبيل المثال، متوسط 63.3% مقارنة بـ 59.1% لـ Pythia-2.8B و61.7% لـ Pythia-6.9B) مع توفير إنتاجية استدلال أعلى بما يصل إلى 5 أضعاف.

الملخص

تعتمد النماذج التأسيسية، التي تشغل حالياً معظم التطبيقات المثيرة في التعلم العميق، بشكل شبه شامل على بنية Transformer ووحدة الانتباه الأساسية الخاصة بها. طُوِّرت العديد من البنى ذات التعقيد الحسابي دون التربيعي مثل الانتباه الخطي (linear attention)، والالتفاف البوابي (gated convolution)، والنماذج المتكررة، ونماذج فضاء الحالة المهيكلة (SSMs) لمعالجة القصور الحسابي لـ Transformers في التسلسلات الطويلة، لكنها لم تقدم أداءً يضاهي الانتباه في وسائط مهمة مثل اللغة. حددنا أن نقطة الضعف الأساسية في هذه النماذج هي عدم قدرتها على الاستدلال المعتمد على المحتوى، وقدمنا العديد من التحسينات. أولاً، جعل معلمات SSM دوالاً معتمدة على المدخلات يعالج ببساطة نقطة ضعفها مع الوسائط المنفصلة، مما يسمح للنموذج بنقل المعلومات أو نسيانها انتقائياً عبر بعد طول التسلسل بناءً على الرمز الحالي. ثانياً، ورغم أن هذا التغيير يحول دون استخدام الالتفافات الفعالة، فقد صممنا خوارزمية تفرعية تراعي طبيعة العتاد في الوضع المتكرر. دمجنا نماذج SSM الانتقائية هذه في بنية شبكة عصبية مبسطة من طرف إلى طرف دون طبقات انتباه أو حتى كتل MLP، وأطلقنا عليها اسم (Mamba). يتميز Mamba باستدلال سريع (إنتاجية أعلى بـ 5 $\times$ أضعاف من Transformers) وتوسع خطي مع طول التسلسل، ويتحسن أداؤه على البيانات الحقيقية حتى تسلسلات بطول مليون رمز. وبصفته نموذجاً أساسياً عاماً للتسلسلات، يحقق 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 والانتباه الذاتي. ورغم أن الانتباه الكثيف ينقل المعلومات بكفاءة عبر نافذة سياق التسلسل، إلا أنه يعاني من قيدين أساسيين: تكلفة حسابية تربيعية $O (L^{2})$ بالنسبة لطول التسلسل $L$ ، وضرورة تخزين المفاتيح والقيم السابقة (مخبأ KV cache)، مما يؤدي إلى استدلال بزمن خطي لكل خطوة وحدوث اختناقات في الذاكرة.

تتوسع المناهج السابقة دون التربيعية—مثل الانتباه الخطي، والبنى المتكررة، ونماذج فضاء الحالة المهيكلة (SSMs) مثل S4—بشكل خطي أو شبه خطي من حيث الزمن. ومع ذلك، فشلت نماذج SSM المهيكلة تاريخياً في مطابقة جودة Transformer في الوسائط المنفصلة الكثيفة بالمعلومات مثل النصوص. حدد المؤلفون السبب الجذري: نماذج SSM السابقة هي نماذج خطية ثابتة زمنياً (LTI). ونظراً لأن ديناميكيات الانتقال الخاصة بها لا تتغير وفقاً لمحتوى المدخلات، فإنها لا تستطيع ضغط السياق بشكل انتقائي، أو ترشيح الضوضاء غير المهمة، أو تعديل الاحتفاظ بالمعلومات عبر تباعدات متغيرة.

الأفكار الرئيسية

آلية الانتقاء بوصفها ضغطاً للسياق: تتطلب نمذجة التسلسلات ضغط السياق السابق في حالة معينة. تستخدم نماذج SSM ذات نمط LTI السابقة معلمات انتقال ثابتة، مما يجبرها على التعامل مع جميع الخطوات الزمنية بشكل متساوٍ. من خلال جعل مصفوفات SSM معتمدة على رمز الإدخال الحالي، يكتسب النموذج الوعي بالمحتوى المطلوب لترشيح المعلومات أو تذكرها بشكل انتقائي.
تجاوز معضلة غياب الالتفاف: يؤدي جعل المعلمات معتمدة على المدخلات إلى كسر التكافؤ الرياضي بين التكرارات الخطية والالتفافات العامة. بدون الالتفافات، سيتطلب التكرار البسيط حفظ حالة كامنة موسعة بحجم $(B, L, D, N)$ في ذاكرة النطاق الترددي العالي (HBM) لوحدة معالجة الرسوميات (GPU). يتغلب Mamba على ذلك عبر دمج التقطيع (discretization)، والمسح التجميعي المتوازي، وإسقاط المخرجات داخل ذاكرة SRAM السريعة المدمجة بالرقاقة، مما يقلل حركة البيانات عبر HBM إلى الحد الأدنى.
بنية موحدة: يستبدل Mamba كتل الانتباه متعدد الرؤوس (MHA) وMLP المتناوبة بكتلة متجانسة موحدة تجمع بين التصميم البوابي لـ H3 ونمط SwiGLU للـ MLP البوابي.

الطريقة

تقوم نماذج فضاء الحالة المستمرة برسم خريطة من المدخلات $x (t) \in R$ إلى $y (t) \in R$ عبر حالة كامنة ضمنية $h (t) \in R^{N}$ : $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}))$ .

توسيع الحالة بمراعاة العتاد

للحفاظ على إنتاجية عالية دون حفظ الحالة $(B, L, D, N)$ في ذاكرة HBM:

يتم تحميل المعلمات $(Δ, A, B, C)$ من HBM مباشرة إلى ذاكرة SRAM المدمجة بالرقاقة.
يتم حساب التقطيع إلى $(\overline{A}, \overline{B})$ والمسح التجميعي المتوازي الفعال حسابياً بالكامل داخل ذاكرة SRAM.
تُقلص الحالة الكامنة باستخدام $C$ إلى الشكل $(B, L, D)$ قبل كتابتها مجدداً في HBM.
يُعاد حساب الحالات الوسيطة في SRAM أثناء مسار التمرير الخلفي (backward pass) بدلاً من حفظها في HBM.

النتائج

الاستقراء الاصطناعي: في مهمة رؤوس الاستقراء (Induction Heads)، يستقرئ Mamba المدرب على طول تسلسل 256 بشكل مثالي (دقة 100%) حتى طول $2^{20} = 1, 048, 576$ ( $4000 \times$ من طول التسلسل)، بينما تفشل النماذج المرجعية المعتمدة على الانتباه أو تعاني من نفاد الذاكرة (OOM) بعد $2^{14}$ . وفي النسخ الانتقائي، يصل Mamba إلى دقة 99.8% مقارنة بـ 57.0% لـ H3 و18.3% لـ S4.
نمذجة اللغة: عند التقييم تحت مقاييس Chinchilla وصولاً إلى 1.3B معامل على The Pile، يطابق Mamba أو يتفوق على Transformer++ (الذي يستخدم تضمينات دوارة، وSwiGLU، وRMSNorm). وفي مهام الصفر المعرفي اللاحقة (مثل LAMBADA، وPIQA، وHellaSwag، وWinoGrande، وARC)، يحقق Mamba-2.8B متوسط دقة يبلغ 63.3%، متفوقاً على Pythia-2.8B (59.1%)، وPythia-6.9B (61.7%)، ومطابقاً أو متجاوزاً لنماذج تحتوي على 6B–7B معلمات.
علم الجينوم والصوتيات: في تصنيف الحمض النووي للقرود العليا عند طول تسلسل قدره مليون رمز ( $2^{20}$ )، حقق Mamba دقة 71.67% (مع 1.4M معامل) و81.31% (مع 7M معامل) مقارنة بـ 54.87% لـ HyenaDNA. وفي توليد الكلام على SC09، خفض Mamba مقياس FID من 1.99 (SaShiMi) إلى 0.94 (مع 6.1M معامل) و0.67 (مع 24.3M معامل).
الكفاءة: تعمل نواة المسح المدمجة بسرعة تفوق تطبيقات المسح القياسية في PyTorch بـ 20–40 ضعفاً، وتعمل بسرعة أعلى بما يصل إلى 7 أضعاف من FlashAttention-2 عند سياق 32K. كما أن إنتاجية الاستدلال أعلى بما يصل إلى 5 أضعاف مقارنة بنماذج Transformers المماثلة.

القيود والأسئلة المفتوحة

أداء الإشارات المستمرة: على الرغم من أن الانتقاء المعتمد على المدخلات يعزز الاستدلال المنفصل (النصوص، والحمض النووي)، إلا أن دراسات الاستئصال على التدريب المسبق للإشارات الصوتية الخام تظهر أن النماذج الخطية الثابتة زمنياً (LTI) البحتة (مثل S4 القياسي) تقدم أداءً أفضل على البيانات المستمرة المأخوذة بانتظام.
حجم النموذج: اقتصرت النماذج التي تم تقييمها على حجم يقارب 3B معامل و300B رمز تدريب مسبق؛ ويتطلب الأداء التجريبي على مقاييس 7B–70B+ وما بعدها مزيداً من البحث والتقصي.
المنظومة البيئية للتطبيقات اللاحقة: تمتلك نماذج Transformers أدوات متطورة وواسعة لهندسة التعليمات الفورية (prompt engineering)، والتعديل داخل السياق، وتكميم مخزن KV، وضبط المحولات (adapters). ولا يزال فهم كيفية تكيف نماذج SSM عبر أنماط التفاعل هذه مجالاً للاستكشاف النشط.

المصطلحات الرئيسية

Structured State Space Models (SSMs): فئة من نماذج التسلسل العميقة المشتقة من أنظمة ديناميكية خطية مستمرة ترسم خريطة لمدخلات أحادية البعد عبر حالة كامنة ضمنية.
Linear Time Invariance (LTI): خاصية رياضية تظل فيها معلمات انتقال النظام ثابتة عبر جميع الخطوات الزمنية، مما يتيح تمثيلات الالتفاف العالمي السريعة ولكنه يحد من الاستدلال المعتمد على المحتوى.
Selection Mechanism: تحديد ديناميكيات انتقال SSM (المعلمات Δ وB وC) كدوال لرمز الإدخال، مما يتيح للنموذج تمرير السياق أو إعادة ضبطه أو تجاهله ديناميكياً على طول التسلسل.
S6: اختصار لنماذج فضاء الحالة المهيكلة الانتقائية التي تُحسب عبر مسح تجميعي متوازٍ بدلاً من الالتفاف.
Parallel Associative Scan: خوارزمية حساب بادئات فعالة من حيث الشغل تقيم التكرار الخطي عبر أطوال التسلسل بالتوازي على عتاد المسرعات.
Hardware-Aware State Expansion: استراتيجية دمج الأنوية البرمجية وإعادة الحساب في GPU تقوم بالتقطيع والمسح داخل ذاكرة SRAM السريعة المدمجة بالرقاقة لتجنب التبادل المتكرر للحالات الوسيطة الضخمة مع ذاكرة النطاق الترددي العالي (HBM).
Induction Heads: آلية تذكر تجميعي اصطناعية تختبر ما إذا كان النموذج قادراً على إعادة إنتاج رمز لوحظ سابقاً عقب الرمز الذي يسبقه في السياق.
Mamba Block: كتلة نموذج تسلسلي مبسطة ومتجانسة تدمج إسقاطات المدخلات، والالتفاف أحادي البعد، ونماذج SSM الانتقائية، والبوابات الضربّية دون الحاجة إلى طبقات انتباه أو MLP منفصلة.

أسئلة شائعة

لماذا لا يمكن تدريب نماذج SSM الانتقائية باستخدام التفافات FFT السريعة؟

تتطلب التفافات FFT العالمية السريعة أن تكون ديناميكيات النظام (المصفوفات A وB وC) ثابتة عبر الزمن (ثبات زمني خطي). ولأن Mamba يجعل B وC وΔ دوالاً تعتمد على رمز الإدخال في كل خطوة زمنية، فإن نواة الالتفاف تختلف عند كل موضع، مما يكسر التكافؤ الالتفافي ويتطلب استخدام مسوحات متوازية متكررة بدلاً منها.

كيف يتجنب Mamba تضخم الذاكرة الناتج عن توسيع الحالة المتكررة أثناء التدريب؟

يتطلب التكرار البسيط مع بُعد توسيع الحالة N تخزين حالة وسيطة بالشكل (B, L, D, N) في ذاكرة HBM الخاصة بـ GPU. يستخدم Mamba دمج الأنوية لتحميل المدخلات منخفضة الأبعاد فقط إلى ذاكرة SRAM، ويحسب التقطيع والمسح المتوازي داخل SRAM السريعة، ويكتب فقط المخرجات المقلصة (B, L, D) مرة أخرى إلى HBM، مع إعادة حساب التنشيطات الوسيطة أثناء التمرير العكسي.

ما هي العلاقة بين المعلمة Δ في Mamba والتحكم البوابي المتكرر؟

تعمل معلمة حجم الخطوة Δ كآلية بوابية معممة. عند تقطيعها تحت أسلوب الاستبقاء من الرتبة صفر (ZOH) مع N=1 وA=-1، يتبسط التكرار الانتقائي رياضياً إلى بوابة RNN قياسية: تجعل القيمة الكبيرة لـ Δ النموذج يركز على المدخلات الحالية ويعيد تعيين الحالة، بينما تجعل القيمة الصغيرة لـ Δ الحالة التاريخية مستمرة ويتجاهل المدخلات.

هل يستخدم Mamba مخبأ KV أثناء توليد المخرجات؟

لا. نظراً لأن Mamba يعمل كتكرار فضاء حالة انحداري تلقائي أثناء التوليد، فإنه يحتفظ فقط بحالة خفية ثابتة الحجم (بأبعاد D × N لكل طبقة) بدلاً من تخزين المفاتيح والقيم لجميع الرموز السابقة، مما ينتج عنه استدلال بزمن ثابت O(1) لكل خطوة.

كتب الذكاء الاصطناعي الشرح في هذه الصفحة استنادًا إلى الورقة، وقد يحتوي على أخطاء. راجع الورقة للتفاصيل.

اقرأ مع الترجمة والذكاء الاصطناعي