Princeton ve CMU, Mamba Dizisi Modeli ile Yapay Zeka Sınırlarını Zorluyor

Bağlantılar Tablosu

Özet ve 1 Giriş

2 Durum Uzay Modeli

3 Seçici Durum Uzay Modelleri ve 3.1 Motivasyon: Sıkıştırma Aracı Olarak Seçim

3.2 Seçimle SSM'leri İyileştirme

3.3 Seçici SSM'lerin Etkin Uygulanması

3.4 Basitleştirilmiş Bir SSM Mimarisi

3.5 Seçim Mekanizmalarının Özellikleri

3.6 Ek Model Ayrıntıları

4 Ampirik Değerlendirme ve 4.1 Sentetik Görevler

4.2 Dil Modelleme

4.3 DNA Modellemesi

4.4 Ses Modelleme ve Üretimi

4.5 Hız ve Bellek Karşılaştırmaları

4.6 Model Ablasyonları

5 Tartışma

6 Sonuç ve Referanslar

Bir Tartışma: Seçim Mekanizması

B İlgili Çalışma

Seçici SSM'lerin C Mekaniği

Seçici SSM'ler için Donanım Farkında Algoritma

E Deneysel Ayrıntılar ve Ek Sonuçlar

Soyut

Derin öğrenmedeki heyecan verici uygulamaların çoğunu destekleyen temel modeller, neredeyse evrensel olarak Transformer mimarisine ve onun temel dikkat modülüne dayanmaktadır. Doğrusal dikkat, kapılı evrişim ve yinelemeli modeller ve yapılandırılmış durum alanı modelleri (SSM'ler) gibi birçok alt kuadratik zamanlı mimari, Transformer'ların uzun dizilerdeki hesaplama yetersizliğini gidermek için geliştirilmiştir, ancak dil gibi önemli kiplerde dikkat kadar iyi performans göstermemişlerdir. Bu tür modellerin temel zayıflıklarından birinin içerik tabanlı akıl yürütme yapamamaları olduğunu tespit ediyoruz ve birkaç iyileştirme yapıyoruz. İlk olarak, SSM parametrelerinin girdinin fonksiyonları olmasına izin vermek, ayrık kiplerdeki zayıflıklarını giderir ve modelin geçerli belirtece bağlı olarak dizi uzunluğu boyutu boyunca bilgiyi seçici olarak yaymasına veya unutmasına olanak tanır. İkinci olarak, bu değişiklik verimli evrişimlerin kullanımını engellese de, yinelemeli modda donanım farkında paralel bir algoritma tasarlıyoruz. Bu seçici SSM'leri, dikkat veya hatta MLP blokları (Mamba) olmadan basitleştirilmiş bir uçtan uca sinir ağı mimarisine entegre ediyoruz. Mamba, hızlı çıkarım (Transformatörlerden 5 kat daha yüksek verim) ve dizi uzunluğunda doğrusal ölçekleme sunar ve performansı milyon uzunluktaki dizilere kadar gerçek verilerde iyileşir. Genel bir dizi modeli omurgası olarak Mamba, dil, ses ve genomik gibi çeşitli modalitelerde en son teknoloji performansına ulaşır. Dil modellemesinde, Mamba-3B modelimiz aynı boyuttaki Transformatörlerden daha iyi performans gösterir ve hem ön eğitimde hem de aşağı akış değerlendirmesinde boyutunun iki katı olan Transformatörlerle eşleşir.

1 Giriş

Temel modeller (FM'ler) veya büyük veriler üzerinde önceden eğitilmiş ve daha sonra alt akış görevleri için uyarlanmış büyük modeller, modern makine öğreniminde etkili bir paradigma olarak ortaya çıkmıştır. Bu FM'lerin omurgası genellikle dil, görüntü, konuşma, ses, zaman serisi ve genomik gibi çok çeşitli alanlardan gelen girdilerin keyfi dizileri üzerinde çalışan dizi modelleridir (Brown ve diğerleri 2020; Dosovitskiy ve diğerleri 2020; Ismail Fawaz ve diğerleri 2019; Oord ve diğerleri 2016; Poli ve diğerleri 2023; Sutskever, Vinyals ve Quoc V Le 2014). Bu kavram belirli bir model mimarisi seçimine bağlı olmasa da, modern FM'ler ağırlıklı olarak tek bir dizi model türüne dayanmaktadır: Transformatör (Vaswani vd. 2017) ve onun çekirdek dikkat katmanı (Bahdanau, Cho ve Bengio 2015). Öz-dikkatin etkinliği, karmaşık verileri modellemesine olanak tanıyan bir bağlam penceresi içinde bilgileri yoğun bir şekilde yönlendirme yeteneğine atfedilir. Ancak, bu özellik temel dezavantajlar getirir: sonlu bir pencerenin dışındaki hiçbir şeyi modelleyememe ve pencere uzunluğuna göre ikinci dereceden ölçekleme. Bu dezavantajların üstesinden gelmek için daha verimli dikkat varyantları üzerinde muazzam bir araştırma gövdesi ortaya çıktı (Tay, Dehghani, Bahri, vd. 2022), ancak çoğu zaman onu etkili kılan özelliklerin pahasına. Şimdiye kadar, bu varyantların hiçbirinin alanlar arasında ölçekte deneysel olarak etkili olduğu gösterilmemiştir.

Son zamanlarda, yapılandırılmış durum uzayı dizi modelleri (SSM'ler) (Gu, Goel ve Ré 2022; Gu, Johnson, Goel ve diğerleri 2021), dizi modelleme için umut verici bir mimari sınıfı olarak ortaya çıkmıştır. Bu modeller, klasik durum uzayı modellerinden (Kalman 1960) esinlenerek, tekrarlayan sinir ağları (RNN'ler) ve evrişimli sinir ağlarının (CNN'ler) bir kombinasyonu olarak yorumlanabilir. Bu model sınıfı, dizi uzunluğunda doğrusal veya doğrusala yakın ölçekleme ile tekrarlama veya evrişim olarak çok verimli bir şekilde hesaplanabilir. Ek olarak, belirli veri modalitelerinde uzun menzilli bağımlılıkları modellemek için ilkeli mekanizmalara sahiptirler (Gu, Dao ve diğerleri 2020) ve Uzun Menzilli Arena (Tay, Dehghani, Abnar ve diğerleri 2021) gibi kıyaslama ölçütlerine hakim olmuşlardır. SSM'lerin birçok çeşidi (Gu, Goel ve Ré 2022; Gu, Gupta ve ark. 2022; Gupta, Gu ve Berant 2022; Y. Li ve ark. 2023; Ma ve ark. 2023; Orvieto ve ark. 2023; Smith, Warrington ve Linderman 2023) ses ve görüntü gibi sürekli sinyal verilerini içeren alanlarda başarılı olmuştur (Goel ve ark. 2022; Nguyen, Goel ve ark. 2022; Saon, Gupta ve Cui 2023). Ancak, metin gibi ayrı ve bilgi yoğun verileri modellemede daha az etkili olmuşlardır.

Transformatörlerin modelleme gücünü elde etmek ve dizi uzunluğunda doğrusal ölçekleme yapmak için birkaç eksende daha önce yapılmış çalışmaları geliştiren yeni bir seçici durum uzayı modelleri sınıfı öneriyoruz.

Seçim Mekanizması. Öncelikle, önceki modellerin temel bir sınırlamasını tanımlıyoruz: girdiye bağlı bir şekilde verileri verimli bir şekilde seçme yeteneği (yani belirli girdilere odaklanma veya onları görmezden gelme). Seçici kopyalama ve tümevarım kafaları gibi önemli sentetik görevlere dayalı sezgiye dayanarak, SSM parametrelerini girdiye göre parametrelendirerek basit bir seçim mekanizması tasarlıyoruz. Bu, modelin alakasız bilgileri filtrelemesine ve alakalı bilgileri sonsuza kadar hatırlamasına olanak tanır.

Donanım Farkında Algoritma. Bu basit değişiklik, modelin hesaplanması için teknik bir zorluk teşkil eder; aslında, tüm önceki SSM modelleri hesaplama açısından verimli olmak için zaman ve girdi değişmez olmalıdır. Bunu, modeli evrişim yerine bir tarama ile yinelemeli olarak hesaplayan, ancak GPU bellek hiyerarşisinin farklı seviyeleri arasında G/Ç erişimini önlemek için genişletilmiş durumu somutlaştırmayan donanım farkında bir algoritma ile aşıyoruz. Ortaya çıkan uygulama, hem teoride (tüm evrişim tabanlı SSM'ler için sözde doğrusala kıyasla sıra uzunluğunda doğrusal olarak ölçekleme) hem de modern donanımda (A100 GPU'larda 3 kata kadar daha hızlı) önceki yöntemlerden daha hızlıdır.

Mimarlık . Önceki derin dizi modeli mimarilerini, önceki SSM mimarilerinin (Dao, Fu, Saab, vd. 2023) tasarımını Transformatörlerin MLP bloğuyla tek bir blokta birleştirerek basitleştiriyoruz ve bu da seçici durum alanlarını içeren basit ve homojen bir mimari tasarımına (Mamba) yol açıyor.

Seçici SSM'ler ve dolayısıyla Mamba mimarisi, diziler üzerinde çalışan genel temel modellerin omurgası olarak uygun hale getiren temel özelliklere sahip tamamen tekrarlayan modellerdir. (i) Yüksek kalite: Seçicilik, dil ve genomik gibi yoğun modalitelerde güçlü performans sağlar. (ii) Hızlı eğitim ve çıkarım: Eğitim sırasında hesaplama ve bellek, dizi uzunluğunda doğrusal olarak ölçeklenir ve çıkarım sırasında modeli otoregresif olarak açmak, önceki öğelerin önbelleğini gerektirmediğinden adım başına yalnızca sabit zaman gerektirir. (iii) Uzun bağlam: Kalite ve verimlilik birlikte, 1M dizi uzunluğuna kadar gerçek verilerde performans iyileştirmeleri sağlar.

Mamba'nın genel bir dizi FM omurgası olarak potansiyelini, hem ön eğitim kalitesi hem de alan-özgü görev performansı açısından, çeşitli türdeki modaliteler ve ayarlar üzerinde deneysel olarak doğruluyoruz:

• Sentetikler. Büyük dil modelleri için anahtar olarak önerilen kopyalama ve indüksiyon başlıkları gibi önemli sentetik görevlerde, Mamba bunları yalnızca kolayca çözmekle kalmaz, aynı zamanda sonsuza kadar uzun çözümler (>1M belirteç) çıkarabilir.

• Ses ve Genomik. Mamba, SaShiMi, Hyena ve Transformers gibi önceki son teknoloji modellerini, hem ön eğitim kalitesi hem de aşağı akış metriklerinde (örneğin, zorlu bir konuşma oluşturma veri kümesinde FID'yi yarıdan fazla azaltarak) ses dalga formlarını ve DNA dizilerini modellemede geride bırakıyor. Her iki ayarda da, performansı milyon uzunluktaki dizilere kadar daha uzun bağlamla iyileşiyor.

• Dil Modelleme. Mamba, hem ön eğitim karmaşasında hem de aşağı akış değerlendirmelerinde Transformatör kalitesinde performansa gerçekten ulaşan ilk doğrusal zamanlı dizi modelidir. 1B parametreye kadar ölçekleme yasalarıyla, Mamba'nın LLaMa'ya (Touvron ve diğerleri 2023) dayalı çok güçlü modern Transformatör eğitim tarifleri dahil olmak üzere geniş bir temel çizgi aralığının performansını aştığını gösteriyoruz. Mamba dil modelimiz, benzer boyuttaki Transformatörlere kıyasla 5 kat daha fazla üretim verimine sahiptir ve Mamba-3B'nin kalitesi, boyutunun iki katı olan Transformatörlerin kalitesiyle eşleşir (örneğin, Pythia-3B'ye kıyasla sağduyulu muhakemede ortalama 4 puan daha yüksek ve hatta Pythia-7B'yi bile aşar).