Apple, Pythia-1B Transformer'ı iki aşamalı bir yöntemle saf Mamba'ya damıtıyor

Transformer'ları hiçbir attention bloğu tutmadan dönüştürmek

Apple'ın makine öğrenmesi araştırma ekibi, attention tabanlı bir dil modelini Mamba tarzı bir state-space modeline damıtma yöntemi yayımladı; yöntem, hibrit tasarımlara geri başvurmak yerine attention'ı tamamen ortadan kaldırıyor. Hacker News ana sayfasına kadar yükselen yazı, Abhinav Moudgil, Ningyuan Huang, Eeshan Gunesh Dhekne, Pau Rodríguez Lopez, Luca Zappella ve Federico Danieli imzasını taşıyor; ilk iki yazar bu çalışmayı Apple'dayken yapmış.

Motivasyon mimari olduğu kadar ekonomik de. Yazarlara göre Mamba sınıfı state-space modelleri, benzer attention tabanlı modellere kıyasla daha küçük bir bellek ayak izi ve daha hızlı token üretimi sunuyor; bu da onları yoğun çıkarım gerektiren dağıtımlar için cazip kılıyor. Ne var ki ekosistemin eğitim bilgisi ve önceden eğitilmiş checkpoint envanteri ezici ölçüde Transformer tabanlı. Damıtma bir köprü vadediyor: Bir state-space modelini sıfırdan eğitmek yerine mevcut bir Transformer'ın bilgisini yeniden kullanmak.

Sorun şu ki en bariz yaklaşım başarısız oluyor. Yazarların atıf yaptığı, mimariler arası damıtma üzerine önceki çalışmalar, bir Transformer öğretmeni doğrudan bir Mamba öğrencisine aktarmanın öğretmenin performansının önemli bir kısmını kaybettirdiğini gösteriyor. Bilinen çözüm yolu, SSM bloklarının yanında bazı attention bloklarını koruyan hibrit bir model kullanmak; ama hibritler, bu geçişe zaten neden olan verimlilik kazanımlarını sulandırıyor.

İki aşamalı yöntem

Apple'ın önerisi, bilgiyi mimariler arasında doğrudan atlamak yerine ara bir temsil üzerinden aktarıyor.

İlk aşamada, özgün Transformer, kernel hilesinin bir uyarlamasıyla oluşturulmuş attention'ın lineerleştirilmiş bir varyantına damıtılıyor. Sonuç, sequence-mixing işlemi biçim olarak bir state-space modelinin ifade edebileceğine çok daha yakın bir ara öğretmen.

İkinci aşamada, bu lineerleştirilmiş model, hiçbir attention bloğu içermeyen uyarlanmış bir Mamba'ya damıtılıyor. Yazarların merkezî argümanı başlatmayla (initialization) ilgili: Mamba öğrencisini matematiksel olarak gerekçelendirilmiş bir başlangıç noktasıyla donatmak — aktarımı sıradan bir eğitim koşusu gibi ele almak yerine — yöntemin bütün olarak işe yaramasını sağlayan şey.

Ne kadar iyi çalışıyor

Ana deney öğretmen olarak Pythia-1B'yi kullanıyor. Makaleye göre damıtılmış Mamba 14.01 perplexity değerine ulaşıyor; öğretmenin 13.86 sonucuna yakın olan bu fark, öğrencinin özgün modelin yeteneğinin çoğunu koruduğu kadar küçük.

Aşağı yönlü görevlerdeki sonuçlar da büyük ölçüde öğretmeni takip ediyor; tek istisna tutarlı biçimde şu: bağlam içi hatırlamaya (in-context recall) ağırlık veren görevlerde hâlâ bir açık var. Yazarlar bunu, yöntemin çözdüğü bir şey olarak değil, kalan zayıflık olarak sunuyor.

Çalışma, 10B damıtma token'ı kullanılarak 1B ölçeğinde yapılan ablation'larla destekleniyor. Bunlar sequence-mixer mimarisi seçimini, model boyutları ve toplam damıtma token bütçeleri arası bir ölçeklendirme analizini ve token'ların iki aşama arasında nasıl bölünmesi gerektiğine dair bir duyarlılık çalışmasını kapsıyor — son nokta önemli, çünkü yöntemi izleyen bir uygulayıcı, sabit bir hesaplama bütçesini nereye harcayacağına karar vermek zorunda.

Neden önemli

Bellek veya iş hacmi kısıtları altında dil modeli dağıtan herkes için bu, somut bir aktarım yolu sunuyor: Yetenek kaynağı olarak önceden eğitilmiş Transformer'ları koruyun, çıkarım zamanında saf bir Mamba dağıtın ve verimlilik kazanımlarını sulandıran hibrit attention bloklarından kaçının. Bu, SSM tabanlı çıkarımı benimseme eşiğini düşürüyor; çünkü topluluğun zaten elinde olan checkpoint'ler ve eğitim pratikleri üzerine inşa ediliyor.

Sonucu çerçeveleyen iki uyarı var. Yöntem 1B parametre ölçeğinde gösterildi; yani öne çıkan rakamlar tek başına yaklaşımın en üst ölçekte de geçerli olduğunu kanıtlamıyor — ne kadar ki dahil edilen ölçeklendirme analizleri, performansın model boyutu ve token bütçesiyle nasıl değiştiğini ele alıyor. Ve bağlam içi hatırlama açığı, attentionsız modellerin her yerde birebir yerine geçebilir olmadığının bir hatırlatıcısı — retrieval ağırlıklı iş yükleri, damıtılmış öğrencinin öğretmeninin gerisinde kalma ihtimalinin en yüksek olduğu yerler.