Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu; Tri Dao

Makaleler · arXiv 2312.00752

Mamba: Seçici Durum Uzayları ile Doğrusal Zamanlı Dizi Modelleme

Mamba: Linear-Time Sequence Modeling with Selective State Spaces

Albert Gu, Tri Dao

Yayımlanma 2023-12-01cs.LGcs.AI arXiv PDF

Çeviri ve yapay zekâ ile oku

Üç satırda özet

🎯 Transformer'lar dizi uzunluğuyla kuadratik olarak ölçeklenir ve otoregresif çıkarım sırasında büyük KV önbellekleri gerektirir; oysa LTI durum uzayı modelleri gibi önceki alt-ikinci-dereceden modeller, dil modelleme gibi ayrık ve içerik akıl yürütmesi gerektiren görevlerde zorlanır.
🔧 Bilgiyi dinamik olarak filtrelemek veya iletmek için SSM parametrelerini girdiye bağımlı (seçici) hale getirir; genişletilmiş gizli durumun GPU HBM'sinde somutlaştırılmasını önlemek için bunu donanım duyarlı kaynaştırılmış paralel ilişkisel tarama ve yeniden hesaplama algoritmasıyla eşleştirir.
📈 Mamba-3B, sıfır-atışlı (zero-shot) alt görevlerde iki katı büyüklüğündeki Transformer'larla eşleşir veya onları geçer (örneğin, Pythia-2.8B için %59,1 ve Pythia-6.9B için %61,7'ye kıyasla ortalama %63,3) ve 5 kata kadar daha yüksek çıkarım verimi sağlar.

Özet

Derin öğrenmedeki en heyecan verici uygulamaların çoğuna güç veren temel modeller, neredeyse evrensel olarak Transformer mimarisine ve onun çekirdek dikkat (attention) modülüne dayanmaktadır. Doğrusal dikkat, kapılı evrişim ve yinelemeli modeller ile yapılandırılmış durum uzayı modelleri (SSM'ler) gibi birçok alt-ikinci-dereceden (subquadratic) zamanlı mimari, Transformer'ların uzun dizilerdeki hesaplama verimsizliğini gidermek için geliştirilmiştir; ancak dil gibi önemli modalitelerde dikkat mekanizması kadar başarılı olamamışlardır. Bu tür modellerin temel zayıflığının içerik tabanlı akıl yürütme yapamamaları olduğunu tespit ediyor ve çeşitli iyileştirmeler sunuyoruz. İlk olarak, SSM parametrelerinin doğrudan girdinin fonksiyonları olmasına izin vermek, bu modellerin ayrık modalitelerdeki zayıflıklarını giderir ve modelin mevcut belirtece (token) bağlı olarak dizi uzunluğu boyutu boyunca bilgiyi seçici bir şekilde iletmesini veya unutmasını sağlar. İkinci olarak, bu değişiklik verimli evrişimlerin kullanımını engellese de, yinelemeli modda donanım duyarlı paralel bir algoritma tasarlıyoruz. Bu seçici SSM'leri, dikkat ve hatta MLP blokları içermeyen basitleştirilmiş uçtan uca bir sinir ağı mimarisine (Mamba) entegre ediyoruz. Mamba, hızlı çıkarım (Transformer'lara kıyasla 5 kat daha yüksek verim) ve dizi uzunluğunda doğrusal ölçekleme avantajına sahiptir; gerçek veriler üzerindeki performansı milyon uzunluğundaki dizilere kadar artış gösterir. Genel bir dizi modeli omurgası olarak Mamba; dil, ses ve genomik gibi çeşitli modalitelerde son teknoloji (state-of-the-art) performans elde eder. Dil modellemesinde, Mamba-3B modelimiz hem ön eğitimde hem de alt görev değerlendirmelerinde kendi boyutundaki Transformer'ları geride bırakır ve iki katı büyüklüğündeki Transformer'larla eşleşir.

Orijinal özet (İngilizce)

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.

Okuma rehberi

Problem ve Motivasyon

Temel modeller ezici bir çoğunlukla Transformer mimarisine ve öz-dikkat (self-attention) mekanizmasına dayanmaktadır. Yoğun dikkat, bir dizi bağlam penceresi boyunca bilgiyi etkili bir şekilde yönlendirse de iki temel kısıtlamaya yol açar: dizi uzunluğu $L$ 'ye göre kuadratik hesaplama maliyeti $O (L^{2})$ ve geçmiş anahtarlar ile değerlerin (KV önbelleği) saklanması gereksinimi; bu da adım başına doğrusal zamanlı çıkarıma ve bellek darboğazlarına neden olur.

Doğrusal dikkat, yinelemeli mimariler ve S4 gibi yapılandırılmış durum uzayı modelleri (SSM'ler) dahil önceki alt-ikinci-dereceden yaklaşımlar, zamanda doğrusal veya sözde doğrusal ölçeklenir. Ancak yapılandırılmış SSM'ler, metin gibi ayrık ve bilgi yoğun modalitelerde tarihsel olarak Transformer kalitesine ulaşamamıştır. Yazarlar bunun temel nedenini tespit etmiştir: önceki SSM'ler doğrusal zamanla değişmezdir (LTI - linear time-invariant). Geçiş dinamikleri girdi içeriğine bağlı olarak değişmediğinden, bağlamı seçici olarak sıkıştıramaz, ilgisiz gürültüyü filtreleyemez veya değişken aralıklar boyunca bilgiyi akılda tutmayı modüle edemezler.

Temel Fikirler

Bağlam Sıkıştırması Olarak Seçim Mekanizması: Dizi modelleme, geçmişin bir duruma sıkıştırılmasını gerektirir. Önceki LTI SSM'ler sabit geçiş parametreleri kullanır ve bu da onları tüm zaman adımlarına tekdüze davranmaya zorlar. SSM matrisleri mevcut girdi belirtecine göre parametrelendirilerek, model bilgiyi seçici olarak filtrelemek veya hatırlamak için gereken içerik farkındalığını kazanır.
Evrişim Dışı Kalma Darboğazını Aşmak: Parametrelerin girdiye bağımlı hale getirilmesi, doğrusal yinelemeler ile genel evrişimler arasındaki matematiksel eşdeğerliği bozar. Evrişimler olmadan, naif yineleme GPU yüksek bant genişlikli belleğinde (HBM) $(B, L, D, N)$ boyutunda genişletilmiş bir örtük durumun somutlaştırılmasını gerektirir. Mamba, ayrıklaştırmayı, paralel ilişkisel taramayı ve çıktı projeksiyonunu hızlı GPU yonga üstü SRAM'inde kaynaştırarak (fusing) HBM trafiğini minimumda tutar ve bu sorunu çözer.
Birleşik Mimari: Mamba, ardışık Çok Başlı Dikkat (MHA) ve MLP bloklarını, H3'ün kapılı tasarımı ile SwiGLU tarzı kapılı MLP'yi birleştiren tek bir homojen blokla değiştirir.

Yöntem

Sürekli durum uzayı modelleri, bir $x (t) \in R$ girdisini örtük bir $h (t) \in R^{N}$ gizli durumu aracılığıyla $y (t) \in R$ çıktısına eşler: $h^{'} (t) = A h (t) + B x (t), y (t) = C h (t)$

$Δ$ adım boyutu ile sıfırıncı dereceden tutma (ZOH - zero-order hold) ayrıklaştırması kullanılarak, sürekli sistem ayrık yineleme parametrelerine dönüştürülür: $\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}$

Seçim Mekanizması (S6)

$(Δ, B, C)$ değerlerini zaman içinde sabit tutmak yerine, Mamba bunları girdiye bağımlı hale getirir:

$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}))$ , burada $τ_{Δ} = softplus$ ve $s_{Δ} (x) = Broadcast_{D} (Linear_{1} (x))$

Bu, klasik RNN kapılamasını genelleştirir: $N = 1, A = - 1, B = 1$ olduğunda, yineleme doğrudan kapılı bir güncellemeye sadeleşir: $h_{t} = (1 - g_{t}) h_{t - 1} + g_{t} x_{t}$ , burada $g_{t} = σ (Linear (x_{t}))$ .

Donanım Duyarlı Durum Genişletme

$(B, L, D, N)$ durumunu HBM'de somutlaştırmadan yüksek verimi korumak için:

$(Δ, A, B, C)$ parametreleri HBM'den doğrudan yonga üstü SRAM'e yüklenir.
$(\overline{A}, \overline{B})$ parametrelerine ayrıklaştırma ve iş verimliliği yüksek paralel ilişkisel tarama tamamen SRAM içinde hesaplanır.
Gizli durum, HBM'ye geri yazılmadan önce $C$ ile $(B, L, D)$ şekline daraltılır.
Ara durumlar HBM'ye kaydedilmek yerine geriye doğru geçiş (backward pass) sırasında SRAM içinde yeniden hesaplanır.

Sonuçlar

Sentetik Ekstrapolasyon: Induction Heads görevinde, 256 dizi uzunluğunda eğitilen Mamba, $2^{20} = 1, 048, 576$ uzunluğuna kadar ( $4000 \times$ dizi uzunluğu) mükemmel bir şekilde (%100 doğruluk) ekstrapolasyon yaparken, dikkat taban çizgileri $2^{14}$ sonrasında başarısız olur veya bellek yetersizliği (OOM) verir. Seçici Kopyalama (Selective Copying) görevinde Mamba, H3 için %57,0 ve S4 için %18,3'e kıyasla %99,8 doğruluğa ulaşır.
Dil Modellemesi: The Pile üzerinde 1.3B parametreye kadar Chinchilla ölçeklemesi altında değerlendirildiğinde Mamba, Transformer++ (rotary embeddings, SwiGLU, RMSNorm) ile eşleşir veya onu geride bırakır. Sıfır-atışlı alt görevlerde (örneğin LAMBADA, PIQA, HellaSwag, WinoGrande, ARC), Mamba-2.8B %63,3 ortalama doğruluk elde ederek Pythia-2.8B (%59,1) ve Pythia-6.9B'yi (%61,7) geride bırakır; 6B–7B parametreli modellerle eşleşir veya onları aşar.
Genomik ve Ses: $1 M$ ( $2^{20}$ ) dizi uzunluğundaki Büyük İnsansı Maymunlar DNA sınıflandırmasında Mamba, HyenaDNA'nın %54,87'lik başarımına kıyasla %71,67 (1.4M parametre) ve %81,31 (7M parametre) doğruluk elde eder. SC09 konuşma üretiminde Mamba, FID'yi 1,99'dan (SaShiMi), 0,94'e (6.1M parametre) ve 0,67'ye (24.3M parametre) düşürür.
Verimlilik: Kaynaştırılmış tarama çekirdeği (fused scan kernel), standart PyTorch tarama uygulamalarından 20–40 kat daha hızlı çalışır ve 32K bağlamda FlashAttention-2'den 7 kata kadar daha hızlıdır. Çıkarım verimi, karşılaştırılabilir Transformer'lara göre 5 kata kadar daha yüksektir.

Sınırlılıklar ve Açık Sorular

Sürekli Sinyal Performansı: Girdiye bağımlı seçim, ayrık akıl yürütmeyi (metin, DNA) geliştirirken, ham ses dalga biçimi ön eğitimi üzerindeki ablasyonlar, tamamen doğrusal zamanla değişmeyen (LTI) modellerin (standart S4 gibi) homojen örneklenmiş sürekli verilerde daha iyi performans gösterdiğini ortaya koymaktadır.
Model Ölçeği: Değerlendirilen modeller yaklaşık 3B parametre ve 300B ön eğitim belirteci ile sınırlıdır; 7B–70B+ ölçeklerinde ve ötesinde ampirik performans daha fazla araştırmayı gerektirir.
Aşağı Akış Ekosistemi: Transformer'lar; prompt mühendisliği, bağlam içi manipülasyon, KV önbellek kuantizasyonu ve adaptör ince ayarı (adapter tuning) için kapsamlı araçlara sahiptir. SSM'lerin bu etkileşim modlarına nasıl uyum sağladığını anlamak aktif bir araştırma konusudur.

Temel terimler

Structured State Space Models (SSMs): 1 boyutlu girdileri örtük bir gizli durum aracılığıyla eşleyen, sürekli doğrusal dinamik sistemlerden türetilmiş bir derin dizi modelleri sınıfı.
Linear Time Invariance (LTI): Sistem geçiş parametrelerinin tüm zaman adımlarında statik kaldığı, hızlı küresel evrişim temsillerine olanak tanıyan ancak içeriğe bağlı akıl yürütmeyi kısıtlayan matematiksel bir özellik.
Selection Mechanism: SSM geçiş dinamiklerini (Δ, B, C) girdi belirtecinin fonksiyonları olarak parametrelendirerek modelin dizi uzunluğu boyunca bağlamı dinamik olarak yaymasına, sıfırlamasına veya yok saymasına olanak tanıyan mekanizma.
S6: Evrişim yerine ilişkisel paralel tarama yoluyla hesaplanan seçici yapılandırılmış durum uzayı modelleri için kullanılan bir kısaltma.
Parallel Associative Scan: Hızlandırıcı donanımlarda dizi uzunlukları boyunca doğrusal yinelemeyi paralel olarak değerlendiren, iş verimliliği yüksek bir önek (prefix) hesaplama algoritması.
Hardware-Aware State Expansion: Büyük ara durumların yüksek bant genişlikli belleğe (HBM) gidip gelmesini önlemek için ayrıklaştırma ve taramayı hızlı yonga üstü SRAM'de gerçekleştiren bir GPU çekirdek kaynaştırma ve yeniden hesaplama stratejisi.
Induction Heads: Bir modelin önceki bağlam belirtecini takip eden daha önce gözlemlenmiş bir belirteci yeniden üretip üretemeyeceğini değerlendiren sentetik bir ilişkisel hatırlama mekanizması.
Mamba Block: Ayrı dikkat veya MLP katmanları olmaksızın girdi projeksiyonlarını, 1D evrişimi, seçici SSM'leri ve çarpımsal kapılamayı birleştiren basitleştirilmiş, homojen bir dizi modeli bloğu.

Sorular

Seçici SSM'ler neden hızlı FFT evrişimleri kullanılarak eğitilemez?

Hızlı genel FFT evrişimleri, sistem dinamiklerinin (A, B ve C matrisleri) zaman boyunca sabit kalmasını (doğrusal zamanla değişmezlik) gerektirir. Mamba, B, C ve Δ değerlerini her zaman adımında girdi belirtecinin fonksiyonları haline getirdiğinden, evrişim çekirdeği her konumda değişir; bu da evrişimsel eşdeğerliği bozar ve bunun yerine yinelemeli paralel taramalar gerektirir.

Mamba, eğitim sırasında yinelemeli durum genişlemesinin yol açtığı bellek patlamasını nasıl önler?

N durum genişletme boyutuna sahip naif yineleme, GPU HBM'sinde (B, L, D, N) şeklinde bir ara durumun somutlaştırılmasını gerektirir. Mamba, SRAM'e yalnızca düşük boyutlu girdileri yüklemek için çekirdek kaynaştırma (kernel fusion) kullanır; ayrıklaştırmayı ve paralel taramayı hızlı SRAM içinde hesaplar, HBM'ye yalnızca daraltılmış (B, L, D) çıktısını geri yazar ve geriye yayılım sırasında ara aktivasyonları yeniden hesaplar.

Mamba'nın Δ parametresi ile yinelemeli kapılama arasındaki ilişki nedir?

Adım boyutu parametresi Δ, genelleştirilmiş bir kapılama mekanizması gibi işlev görür. N=1 ve A=-1 ile sıfırıncı dereceden tutma (ZOH) altında ayrıklaştırıldığında, seçici yineleme matematiksel olarak standart bir RNN kapısına indirgenir: büyük bir Δ, modelin mevcut girdiye odaklanmasına ve durumu sıfırlamasına neden olurken, küçük bir Δ geçmiş durumu korur ve girdiyi yok sayar.

Mamba üretim sırasında bir KV önbelleği kullanır mı?

Hayır. Mamba üretim sırasında otoregresif bir durum uzayı yinelemesi olarak çalıştığı için, geçmiş tüm belirteçlerin anahtarlarını ve değerlerini önbelleğe almak yerine yalnızca sabit boyutlu bir gizli durumu (katman başına D × N boyutu) korur; bu da adım başına sabit zamanlı O(1) çıkarım sağlar.

Bu sayfadaki açıklama yapay zekâ tarafından makaleden yazılmıştır ve hata içerebilir. Ayrıntılar için makaleyi okuyun.

Çeviri ve yapay zekâ ile oku