‏Speculative Decoding: چطور LLMها با حدس زدن سریع‌تر می‌شن

نویسنده تو این پست از پایه نشون می‌ده چرا decoding (یعنی مرحله‌ای که مدل جواب رو توکن به توکن می‌سازه) کنده. مشکل محاسبه نیست، حافظه‌ست. مثال مقاله یه مدل 8B با bf16 روی H100 هست: فقط خوندن 16 گیگ وزن با پهنای باند 3.35 TB/s حدود 4.8 میلی‌ثانیه طول می‌کشه. یعنی سقف حدود 208 توکن در ثانیه با batch size 1. نکتهٔ اصلی اینه که اگه به جای یه توکن 8 تا بدی، باز همون 4.8 میلی‌ثانیه طول می‌کشه، چون وزن‌ها یه بار خونده می‌شن.

راه‌حل، speculative decoding هست. یه مدل کوچیک به اسم drafter (یعنی مدلی که ارزون و سریع توکن‌های بعدی رو حدس می‌زنه) چند توکن پیشنهاد می‌ده. بعد مدل اصلی یا target فقط یه بار روی همه‌شون اجرا می‌شه. بلندترین بخشی از حدس‌ها که target باهاش موافقه نگه داشته می‌شه و سر اولین اختلاف، توکن خود target جاش می‌شینه. پس هر دور حداقل یه توکن و حداکثر γ+1 توکن می‌ده و خروجی هیچ فرقی با decoding معمولی نداره.

وقتی با temperature نمونه‌برداری می‌کنی، قانون speculative sampling (از Leviathan و Chen در 2023) به کار میاد: توکن پیشنهادی با احتمال min(1, p/q) قبول می‌شه و اگه رد شد، از یه توزیع باقی‌مونده نمونه گرفته می‌شه. نویسنده ثابت می‌کنه خروجی دقیقاً از توزیع target پیروی می‌کنه. یعنی drafter بد فقط سرعت رو کم می‌کنه، نه کیفیت رو. نرخ قبول هم برابر یک منهای total variation distance بین دو توزیعه.

یه روش دیگه Gumbel coupling هست: drafter و target از یه نویز Gumbel مشترک استفاده می‌کنن (یعنی بخش تصادفی sampling از خود مدل جدا می‌شه) و چک کردن می‌شه یه مقایسهٔ سادهٔ argmax. مزیتش اینه که با یه seed ثابت، هر drafterی همون متن رو می‌ده و تست و دیباگ خیلی راحت‌تر می‌شه. هزینه‌ش یه نرخ قبول کمی پایین‌تره؛ تو مثال مقاله 0.679 در برابر 0.70.

برای حساب سرعت، مقاله فرمول speedup = τ / (γc + 1) رو می‌ده؛ τ میانگین توکن‌های هر دور و c هزینهٔ هر پاس drafter نسبت به target هست. با α=0.8، γ=5 و c=0.05 به τ حدود 3.69 و سرعتی نزدیک 2.95 برابر می‌رسیم. نویسنده بعدش سراغ token tree (چند شاخه حدس که با یه attention mask درختی یکجا چک می‌شن)، block verification که کل پیش‌نویس رو با هم قضاوت می‌کنه، و معماری‌هایی مثل EAGLE-3، multi-token prediction و DFlash می‌ره.

نکات کلیدی:

• یه مدل 8B روی H100 با batch size 1 به‌خاطر پهنای باند حافظه حداکثر حدود 208 توکن در ثانیه می‌ده.
• تا حدود 295 توکن در هر پاس، توکن اضافه تقریباً هزینه‌ای نداره چون پاس memory-bound هست.
• speculative sampling خروجی رو دقیقاً هم‌توزیع با مدل اصلی نگه می‌داره؛ drafter فقط روی سرعت اثر داره.
• Gumbel coupling خروجی رو مستقل از drafter می‌کنه، با نرخ قبول کمی پایین‌تر.
• با α=0.8، γ=5 و c=0.05 سرعت حدود 2.95 برابر می‌شه.