Medusa Heads

رؤوس Medusa هي رؤوس فك ترميز خفيفة الوزن تُضاف فوق آخر طبقة مخفية في نموذج لغة كبير، ويُدرّب كل منها على التنبؤ بـرمز أبعد مما تتنبأ به طبقة خرج النموذج نفسه. وفي تمريرة أمامية واحدة، ينتج النموذج تنبؤه المعتاد بالرمز التالي، إلى جانب تخمينات للمواضع القليلة التالية؛ ثم تتحقق خطوة انتباه شجري من تلك التخمينات في التمريرة التالية، وتحتفظ بأطول تسلسل يوافق عليه النموذج. Medusa شكل من فك الترميز التخميني لا يحتاج إلى نموذج مسودة منفصل، فلا توجد شبكة ثانية لتدريبها أو استضافتها أو إبقاؤها متزامنة. وقد قدّم تيانلي كاي وتري داو وزملاؤهما هذا الأسلوب في Medusa: إطار بسيط لتسريع استدلال LLM باستخدام رؤوس فك ترميز متعددة في يناير 2024.

كيف تتنبأ الرؤوس بالرموز اللاحقة#

تُدخل خطوة توليد النص القياسية الحالة المخفية الأخيرة عند الموضع t إلى رأس النموذج اللغوي، الذي يقيّم الرمز عند t + 1. ويقرأ الرأس k في Medusa الحالة المخفية نفسها ويقيّم الرمز عند t + k + 1. وكل رأس طبقة تغذية أمامية واحدة، مع تنشيط SiLU ووصلة متبقية، يتبعها إسقاط إلى المعجم. ويبدأ الإسقاط نسخةً من رأس النموذج اللغوي، وتبدأ طبقة التغذية الأمامية بأصفار؛ لذا يعيد كل رأس قبل التدريب إنتاج تنبؤ النموذج الأساسي بالرمز التالي. ويجد البحث أن خمسة رؤوس كافية. ويبني هذا المخطط التوضيحي في PyTorch ثلاثة رؤوس، ويقارن مخرجاتها بمخرجات رأس النموذج الأساسي على حالة مخفية واحدة:

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)  # ⁨طبقة خرج النموذج المجمّدة⁩

# ⁨كل رأس من رؤوس Medusa: طبقة SiLU متبقية واحدة (مهيأة بالأصفار) ونسخة من رأس LM⁩
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)

# ⁨تقرأ جميع الرؤوس الحالة المخفية الأخيرة نفسها؛ ويستهدف الرأس k الموضع t + k + 1⁩
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] أن الطبقة المتبقية المهيأة بالأصفار تجعل كل رأس جديد يبدأ نسخةً مطابقة تمامًا لتنبؤ النموذج الأساسي. ثم ينقل التدريب الرأس k نحو موضعه المستقبلي الخاص، بينما يظل lm_head والعمود الفقري دون تغيير.

الانتباه الشجري والقبول#

تتحد التخمينات الأعلى للرؤوس لتكوين امتدادات مرشحة. فإذا أخذنا أعلى رمزين من الرأس 1 وأعلى ثلاثة من الرأس 2، نتج عن الجداء الديكارتي ستة مرشحين يشتركون في بادئات ويشكّلون شجرة:

في التمريرة الأمامية التالية، يُضاف كل رمز مرشح إلى المُدخل دفعةً واحدة، ويسمح قناع الانتباه لكل رمز برؤية أسلافه فقط في الشجرة، بحيث يقيّم النموذج جميع الفروع بالتوازي. ويُحتفظ بأطول فرع يقبله النموذج، وتغذّي حالته المخفية النهائية الرؤوس في الجولة التالية. ولأن جميع الرؤوس تقرأ الحالة المخفية نفسها، يقترح الرأس 2 الرموز نفسها ضمن كل فرع من فروع الرأس 1. ووجد المؤلفون أن شجرة متناثرة من 64 عقدة، بُنيت من التركيبات الأعلى دقة، حققت تسريعًا أكبر من شجرة كثيفة من 256 عقدة.

في فك الترميز الجشع، يُقبل الرمز المرشح عندما يطابق اختيار النموذج الأعلى. أما في أخذ العينات، فتتيح Medusa القبول النموذجي، الذي يُبقي أي مرشح يتجاوز احتماله وفق النموذج الأساسي عتبةً مرتبطة بإنتروبيا التوزيع، فيقبل رموزًا أكثر مقابل التخلي عن المطابقة الدقيقة للتوزيع.

تدريب Medusa-1 وMedusa-2#

يدرّب Medusa-1 الرؤوس الجديدة فقط على عمود فقري مجمّد، فتحافظ مخرجات النموذج على جودتها، ويذكر البحث أن التوليد أسرع بأكثر من 2.2x. واستغرق تدريب رؤوس Vicuna-7B نحو خمس ساعات على وحدة GPU واحدة من A100 باستخدام 60,000 محادثة من ShareGPT. وهذا هو الفصل نفسه الذي يهيئه freeze=23 عند الضبط الدقيق لـUltralytics YOLO26: تظل الطبقات 0-22 ثابتة ولا تتعلم من البيانات الجديدة إلا رأس الاكتشاف في الطبقة 23.

يُجري Medusa-2 ضبطًا دقيقًا للرؤوس والعمود الفقري معًا، ما يجعل الرؤوس أدق ويرفع التسريع إلى 2.3-3.6x، باستخدام وصفة تدريب تحمي قدرات العمود الفقري الأصلية. وعندما لا تتوفر بيانات التدريب الأصلية، يدرّب التقطير الذاتي الرؤوس على نص يولّده النموذج بنفسه.

يأتي Medusa ضمن عدة أساليب لصياغة رموز مسودة من النموذج المستهدف نفسه. ويصوغ EAGLE المسودات من السمات الداخلية للنموذج باستخدام طبقة ذاتية الانحدار صغيرة، بينما يدرّب التنبؤ متعدد الرموز، المستخدم في DeepSeek-V3، وحدات تنبؤ إضافية أثناء التدريب المسبق، ثم تعيد محركات الخدمة استخدامها كمولدات مسودات. وتقلل الأساليب الثلاثة زمن استجابة الاستدلال بأكبر قدر عند أحجام الدفعات الصغيرة الخاصة بـروبوتات الدردشة التفاعلية، حيث تترك كل تمريرة أمامية وحدة GPU في انتظار الذاكرة معظم الوقت.