Ekstra Maliyet Olmadan Yapay Zekayı Hızlandıran Son Teknoloji Teknikler

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

3.3 Seçici SSM'lerin Etkin Uygulanması

Konvolüsyonlar (Krizhevsky, Sutskever ve Hinton 2012) ve Transformatörler (Vaswani ve ark. 2017) gibi donanım dostu mimariler yaygın uygulama alanına sahiptir. Burada seçici SSM'leri modern donanımlarda (GPU) da verimli hale getirmeyi amaçlıyoruz. Seçim mekanizması oldukça doğaldır ve daha önceki çalışmalar, yinelemeli SSM'lerde ∆'nin zamanla değişmesine izin vermek gibi özel seçim durumlarını dahil etmeye çalışmıştır (Gu, Dao ve ark. 2020). Ancak, daha önce belirtildiği gibi SSM'lerin kullanımındaki temel bir sınırlama, hesaplama verimliliğidir, bu nedenle S4 ve tüm türevleri çoğunlukla küresel konvolüsyonlar biçiminde LTI (seçici olmayan) modeller kullanmıştır.

3.3.1 Önceki Modellerin Motivasyonu

Öncelikle bu motivasyonu tekrar ele alıp, önceki yöntemlerin sınırlamalarını aşmak için yaklaşımımızı gözden geçireceğiz.

• Yüksek düzeyde, SSM'ler gibi tekrarlayan modeller her zaman ifade gücü ve hız arasında bir denge kurar: Bölüm 3.1'de tartışıldığı gibi, daha büyük gizli durum boyutuna sahip modeller daha etkili ancak daha yavaş olmalıdır. Bu nedenle, hız ve bellek maliyetleri ödemeden gizli durum boyutunu en üst düzeye çıkarmak istiyoruz.

• Tekrarlayan modun, ikincisinin (3) birincisinin (2) genişletilmesinden türetilmesi nedeniyle, evrişim modundan daha esnek olduğunu unutmayın (Gu, Goel ve Ré 2022; Gu, Johnson, Goel ve diğerleri 2021). Ancak, bu, (B, L, D, N) şekline sahip gizli durum ℎ'nin hesaplanmasını ve somutlaştırılmasını gerektirir; bu, (B, L, D) şeklinin giriş x ve çıkış y'sinden çok daha büyüktür (N faktörüyle, SSM durum boyutu). Bu nedenle, durum hesaplamasını atlatabilen ve yalnızca (B, L, D)'nin bir evrişim çekirdeğini (3a) somutlaştıran daha verimli evrişim modu tanıtıldı.

• Önceki LTI SSM'ler, geleneksel RNN'lerden çok daha büyük olan ve verimlilikten ödün vermeyen Nx (≈ 10 - 100) faktörüyle etkin durum boyutunu artırmak için ikili tekrarlayan-evrişimli formlardan yararlanıyordu.

3.3.2 Seçmeli Taramaya Genel Bakış: Donanım Farkında Durum Genişletmesi

Seçim mekanizması LTI modellerinin sınırlamalarını aşmak için tasarlanmıştır; bu nedenle aynı zamanda SSM'lerin hesaplama problemini yeniden ele almamız gerekir. Bunu üç klasik teknikle ele alıyoruz: çekirdek birleştirme, paralel tarama ve yeniden hesaplama. İki ana gözlem yapıyoruz:

• Naif tekrarlı hesaplama O(BLDN) FLOP kullanırken, evrişimli hesaplama O(BLD log(L)) FLOP kullanır ve ilkinin daha düşük bir sabit faktörü vardır. Bu nedenle uzun diziler ve çok büyük olmayan durum boyutu N için, tekrarlı mod aslında daha az FLOP kullanabilir.

• İki zorluk, yinelemenin ardışık doğası ve büyük bellek kullanımıdır. İkincisini ele almak için, tıpkı evrişimsel modda olduğu gibi, tam durum ℎ'yi gerçekte somutlaştırmamaya çalışabiliriz.

Ana fikir, modern hızlandırıcıların (GPU'lar) özelliklerini kullanarak ℎ durumunu yalnızca bellek hiyerarşisinin daha verimli seviyelerinde somutlaştırmaktır. Özellikle, çoğu işlem (matris çarpımı hariç) bellek bant genişliğiyle sınırlıdır (Dao, Fu, Ermon, vd. 2022; Ivanov vd. 2021; Williams, Waterman ve Patterson 2009). Bu, tarama işlemlerimizi içerir ve bellek G/Ç miktarını azaltmak için çekirdek birleştirmeyi kullanırız ve bu da standart bir uygulamaya kıyasla önemli bir hızlanmaya yol açar.

Sıralı tekrarı önlemek için, doğrusal olmasa da yine de iş açısından verimli bir paralel tarama algoritmasıyla paralel hale getirilebileceğini gözlemliyoruz (Blelloch 1990; Martin ve Cundy 2018; Smith, Warrington ve Linderman 2023).

Son olarak, geri yayılım için gerekli olan ara durumları kaydetmekten de kaçınmalıyız. Bellek gereksinimlerini azaltmak için klasik yeniden hesaplama tekniğini dikkatlice uygularız: ara durumlar saklanmaz, ancak girdiler HBM'den SRAM'a yüklendiğinde geri geçişte yeniden hesaplanır. Sonuç olarak, kaynaştırılmış seçici tarama katmanı, FlashAttention ile optimize edilmiş bir dönüştürücü uygulamasıyla aynı bellek gereksinimlerine sahiptir.

Birleştirilmiş çekirdeğin ve yeniden hesaplamanın ayrıntıları Ek D'de yer almaktadır. Tam Seçici SSM katmanı ve algoritması Şekil 1'de gösterilmiştir.