Plastik ağırlıklı, Rust ile yazılmış PSSA modelu, eşdeğer bir transformer'ı geride bırakıyor

Attention yok, ML framework yok

PSSA, plastik state-space mimarisinin kısaltması, hem transformer tasarımını hem de geleneksel derin öğrenme yığınını terk eden küçük bir dil modeli. Hacker News ana sayfasına çıkan GitHub'daki proje README'sine göre uygulama Rust ile yazılmış ve altında PyTorch, TensorFlow ya da başka bir framework yok: lineer cebir elle yazılmış, eğitim bir CUDA cihazı kullanabiliyor ve her gradient, skaler bir CPU referansıyla doğrulanıyor; bildirilen en büyük sapma 2.98e-8.

Mimarinin dikkat çeken dört parçası var. Tekrarlayan bir state-space çekirdeği metni bir seferde bir token işliyor ve sabit boyutlu bir durumu bir bağlam penceresine dikkat etmek yerine öğrenilmiş sürekli matrisler üzerinden ileri taşıyor. 512 slotlu bir episodik bellek bankası, Poincaré tarzı hiperbolik erişim ve sınırlı top-4 arama ile model çalışırken yazılıp okunuyor. Plastik ağırlıklar anlık olarak güncelleniyor: hızlı değişiklikler işe yarayanı pekiştiriyor, yenilik büyümeyi tetikliyor ve bir refrakter kapı üzerine yazmayı hız sınırına tabi tutuyor, böylece tekrarlanan çelişkili girdi daha az zarar veriyor. Son olarak, kapalı formda bir ridge regression adımı bu hızlı güncellemeleri düzenli olarak temel geçiş matrisine katıyor; README bunu uykunun öğrenmeyi pekiştirmesine benzetiyor.

Durum sabit boyutlu olduğu için token başına maliyet, dikkat mekanizmasının bağlam üzerindeki kuadratik ölçeklenmesinin aksine, önceki metnin uzunluğuyla büyümüyor.

Karşılaştırma sayıları

Karşılaştırma, PSSA'yı yaklaşık eşit boyutta standart bir transformer'la (taban model 1.541.120 parametre olarak listeleniyor) yarıştırıyor; her ikisi de aynı temizlenmiş WikiText-103 geçişi, aynı tokenizer, optimizer planı ve seed ile eğitilmiş. Her iki çalışma da 200.000 kodlanmış tokenlik 64 bağlantı halinde düzenlendi; her biri önceki checkpoint'ten devam ediyordu, böylece kosinüs planı ve optimizer durumu tüm çalışma boyunca sürüyordu.

12,7M eğitim tokenı üzerinde PSSA 3,98 cross-entropy ile bitirdi, transformer ise 4,43; fark 0,45 nat, yani perplexity 53,7'ye karşı 83,7. Transformer, tüm token bütçesini harcayarak PSSA'nın yaklaşık 2M token işaretinde çoktan geçtiği bir loss'a ulaştı.

İki çalışmanın da dokunmadığı 198.939 tokenlik ayrılmış bir dilimde final checkpoint'leri 3,997'ye karşı 4,429 cross-entropy ve yüzde 24,1'e karşı yüzde 18,0 next-token doğruluğu aldı. O dilimin sınırlı pencerelerinde puanlanan her iki çalışmanın tüm checkpoint'lerinde transformer'ın eğrisi hiçbir zaman PSSA'nınkini geçmedi ve ayrılmış veri setindeki 0,43 natlık fark eğitim farkıyla özdeş; yazar bunu daha iyi genelleme olarak okuyor, daha sert ezberleme olarak değil. Bir pürüz de açıkça belirtilmiş: taban modelin ilk oturumu bir notebook süre sınırı nedeniyle 43. bağlantıda kesilmiş ve loss CSV'si kaybolmuş; bu bağlantılar, zincir devam edip bitmeden önce oturumun çalışma kaydından yeniden oluşturulmuş.

CPU'da yaklaşık on iki kat daha hızlı

Aynı CPU, prompt ve sampler ile 200 token üretmek PSSA için 226 ms, transformer için 2.735 ms sürdü. README, bu CPU'dan CPU'ya karşılaştırmanın adil olanı olduğunu vurgularken, eğitim iş hacmi rakamlarının donanımları eşit olmadığını belirtiyor: PSSA saniyede yaklaşık 900 token ile bir Kaggle T4'te eğitilirken taban model yalnızca CPU'da çalıştı, dolayısıyla bu sayılar mimari bir sonuç olarak okunmamalı.

Belirtilen sınırlar

Yazar ölçek konusunda açık sözlü. Bunlar 12,7M token üzerinde eğitilmiş 1,5M parametrelik modeller; herhangi bir üretim sisteminin rakibi değil, bir araştırma prototipi. Bu boyutta her ikisinin de metin kalitesi zayıf; README, PSSA'dan 'a barget of the Prian Academy' ve transformer'dan 'a material circulation of the United States' örnekleri veriyor. İki deney hâlâ ölçülmemiş: corpus değişikliğinde önceki becerilerin kalıp kalmadığı ve bellek bankasının kaldırılmasının loss'u değiştirip değiştirmeyeceği.

Çalıştırma ve projenin ihtiyaçları

Proje, Edition 2024 destekli bir Rust araç zinciriyle cargo build --release komutuyla derleniyor; iki doğrudan runtime bağımlılığı var (veri seti indirmeleri için ureq ve byte seviyesinde BPE için tokenizers); CUDA isteğe bağlı ve CPU yolu her zaman çalışıyor. Ana talep işlem gücü: şimdiye kadarki her şey birkaç saat sonra kesilen oturumlarla, giriş seviyesi bir GPU'ya sahip ücretsiz bir barındırılmış notebook'ta çalıştı; farkın 10 kat veya 100 kat parametrede sürüp sürmediği ya da bellek bankasının ölçekte önemli olup olmadığı gibi sorular gerçek VRAM ve günlerle ölçülen tahsisler gerektiriyor. En çok istenen kod katkıları kernel performansı, karşılaştırılacak modern bir tekrarlayan taban model ve next-token loss'un ötesinde değerlendirme.

Neden önemli

Baskın dil modelleri transformer'lardır ve maliyetleri bağlam uzunluğuyla büyür. PSSA, farklı bir bahsin kompakt ve framework'süz bir sınaması: sabit boyutlu tekrarlayan durum artı erişim belleği artı eğitimden sonra değişmeye devam eden ağırlıklar. Oyuncak ölçeğinde, olağandüstü açık uyarılarla denetlenebilir bir kod tabanının içinde, hem öğrenme verimliliğinde hem de çıkarım maliyetinde eşdeğer bir transformer'ı yeniyor. Üstünlüklerin daha büyük ölçekte, modern tekrarlayan taban modellere karşı ve bellek ablasyonlarında sürüp sürmediği, tam da projenin henüz test edemediği şey; GPU istemesinin nedeni bu.