Medusa Heads

Medusa başlıkları, bir büyük dil modelinin son gizli katmanının üzerine eklenen hafif çözümleme başlıklarıdır; her biri modelin kendi çıktı katmanından daha ilerideki bir token'ı tahmin edecek şekilde eğitilir. Tek bir ileri geçişte model, olağan sonraki token tahminini ve sonraki birkaç konum için tahminleri üretir; bir sonraki geçişte ağaç dikkati bu tahminleri denetler ve modelin kabul ettiği en uzun diziyi tutar. Medusa, ayrı bir taslak model gerektirmeyen bir spekülatif çözümleme yöntemidir; dolayısıyla eğitilecek, barındırılacak veya eşzamanlı tutulacak ikinci bir ağ yoktur. Tianle Cai, Tri Dao ve çalışma arkadaşları, yöntemi Ocak 2024'te Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads başlıklı çalışmada tanıttı.

Başlıklar İlerideki Token'ları Nasıl Tahmin Eder#

Standart bir metin üretimi adımında, t konumundaki son gizli durum dil modeli başlığına beslenir ve başlık t + 1 konumundaki token'ı puanlar. Medusa k başlığı aynı gizli durumu okur ve t + k + 1 konumundaki token'ı puanlar. Her başlık, SiLU aktivasyonu ve artık bağlantı içeren tek bir ileri beslemeli katmandan, ardından söz dağarcığına yapılan bir izdüşümden oluşur. İzdüşüm, dil modeli başlığının kopyası olarak başlar; ileri beslemeli katman ise sıfırdan başlatılır. Böylece eğitimden önce her başlık, temel modelin sonraki token tahminini aynen üretir. Makaleye göre beş başlık yeterlidir. Bu PyTorch taslağı üç başlık oluşturur ve tek bir gizli durum üzerindeki çıktılarını temel modelin başlığıyla karşılaştırır:

import torch
import torch.nn.functional as F

hidden_size, vocab_size, num_heads = 256, 1000, 3
lm_head = torch.nn.Linear(hidden_size, vocab_size, bias=False)  # dondurulmuş modelin çıktı katmanı

# Her Medusa başlığı: sıfır başlatılmış tek bir artık SiLU katmanı ve LM başlığının kopyası
w1 = [torch.nn.Linear(hidden_size, hidden_size) for _ in range(num_heads)]
w2 = [torch.nn.Linear(hidden_size, vocab_size, bias=False) for _ in range(num_heads)]
for layer, out in zip(w1, w2):
    torch.nn.init.zeros_(layer.weight)
    torch.nn.init.zeros_(layer.bias)
    out.weight.data.copy_(lm_head.weight.data)

# Tüm başlıklar aynı son gizli durumu okur; k başlığı t + k + 1 konumunu hedefler
h = torch.randn(1, hidden_size)
head_logits = [out(h + F.silu(layer(h))) for layer, out in zip(w1, w2)]
print([torch.equal(logits, lm_head(h)) for logits in head_logits])

[True, True, True] çıktısı, sıfır başlatılmış artık katmanın her yeni başlığın temel model tahmininin birebir kopyası olarak başlamasını sağladığını gösterir. Eğitim daha sonra k başlığını kendi gelecek konumuna doğru kaydırırken, lm_head ve omurga değişmeden kalır.

Ağaç Dikkati ve Kabul#

Başlıkların en yüksek puanlı tahminleri aday devamları oluşturmak üzere birleştirilir. 1. başlıktan en yüksek puanlı iki token, 2. başlıktan ise en yüksek puanlı üç token alındığında Kartezyen çarpım, ortak önekleri paylaşan altı aday verir ve bir ağaç oluşturur:

Bir sonraki ileri geçişte her aday token girdiye aynı anda eklenir ve bir dikkat maskesi her token'ın yalnızca ağaçtaki atalarını görmesini sağlar; böylece model tüm dalları paralel olarak puanlar. Modelin kabul ettiği en uzun dal tutulur ve bu dalın son gizli durumu, sonraki tur için başlıkları besler. Tüm başlıklar aynı gizli durumu okuduğundan, 2. başlık her 1. başlık dalında aynı token'ları önerir. Yazarlar, en isabetli birleşimlerden oluşturulan 64 düğümlü seyrek ağacın, yoğun 256 düğümlü ağaçtan daha yüksek hızlanma sağladığını buldu.

Açgözlü çözümlemede aday token, modelin en yüksek seçimle eşleştiğinde kabul edilir. Örnekleme için Medusa, temel model altındaki olasılığı dağılımın entropisine bağlı bir eşiği aşan tüm adayları tutan tipik kabul yöntemini sunar; bu yöntem daha fazla token kabul eder, ancak dağılımla tam eşleşme pahasına çalışır.

Medusa-1 ve Medusa-2 Eğitimi#

Medusa-1 yalnızca dondurulmuş bir omurga üzerindeki yeni başlıkları eğitir; böylece modelin çıktıları kalitesini korur ve makale 2,2 kattan fazla hızlı üretim bildirir. Vicuna-7B başlıklarının eğitimi, 60.000 ShareGPT konuşmasıyla tek bir A100 GPU üzerinde yaklaşık beş saat sürdü. Bu, freeze=23 tarafından Ultralytics YOLO26'ya ince ayar yapılırken kurulan ayrımın aynısıdır: 0-22. katmanlar sabit kalır ve yalnızca 23. katmandaki algılama başlığı yeni verilerden öğrenir.

Medusa-2, başlıklarla birlikte omurgaya da ince ayar yapar; bu, başlıkların isabetini artırır ve omurganın özgün becerilerini koruyan bir eğitim tarifiyle hızlanmayı 2,3-3,6 kata çıkarır. Özgün eğitim verileri mevcut olmadığında öz damıtma, başlıkları modelin kendi ürettiği metinler üzerinde eğitir.

Medusa, hedef modelin kendisinden token taslağı oluşturmanın çeşitli yöntemleri arasında yer alır. EAGLE, modelin iç özelliklerinden küçük bir otoregresif katmanla taslak üretir; DeepSeek-V3'te kullanılan çoklu token tahmini ise ön eğitim sırasında ek tahmin modülleri eğitir ve sunum motorları bunları daha sonra taslak oluşturucu olarak yeniden kullanır. Bu üç yöntemin tümü, her ileri geçişin GPU belleğini bekleyerek büyük ölçüde boş kaldığı etkileşimli sohbet botlarının küçük toplu iş boyutlarında çıkarım gecikmesini en fazla azaltır.