رمزگشایی پیشگویانه: ترفند استنتاجی که LLMها را بدون کاهش کیفیت 3 برابر سریع‌تر می‌کند

رمزگشایی پیشگویانه تأخیر استنتاج LLM را 2 تا 3 برابر بدون تغییر وزن‌های مدل، کیفیت خروجی یا توزیع احتمال زیربنایی کاهش می‌دهد. این کار با بهره‌برداری از یک عدم تقارن ساختاری انجام می‌شود که بیشتر خطوط لوله استنتاج نادیده می‌گیرند: تأیید توکن‌ها بسیار ارزان‌تر از تولید آن‌ها از ابتدا است.

این تکنیک در مقاله سال ۲۰۲۳ DeepMind رسمی شد و از آن زمان در سیستم‌های تولیدی در Google، Meta و تعدادی رو به رشد از فریم‌ورک‌های بهینه‌سازی استنتاج از جمله vLLM و TensorRT-LLM به کار گرفته شده است. درک دلیل کارکرد آن — و جایی که با شکست مواجه می‌شود — برای هر تیمی که LLMها را در مقیاس اجرا می‌کند ضروری است.

نحوه عملکرد واقعی رمزگشایی پیشگویانه

تولید خودرگرسیون استاندارد به طور ذاتی ترتیبی است: مدل یک توکن را در هر بار تولید می‌کند، هر کدام مشروط به تمام توکن‌های قبلی. این بدان معناست که یک مدل با 70 میلیارد پارامتر روی یک A100 حدود 20 تا 30 توکن در ثانیه تولید می‌کند و بخش عمده محاسبات خود را به جای عملیات حسابی صرف پهنای باند حافظه می‌کند. GPU بیشتر منتظر بارگذاری وزن‌ها از HBM است، نه اینکه اعداد را محاسبه کند.

رمزگشایی پیشگویانه یک مدل "پیش‌نویس" دوم و بسیار کوچک‌تر معرفی می‌کند — معمولاً 7 میلیارد پارامتر یا کمتر — که دنباله‌ای از توکن‌های نامزد (معمولاً 4 تا 8 توکن در یک زمان) را پیش از مدل "هدف" بزرگ‌تر پیشنهاد می‌دهد. سپس مدل هدف تمام توکن‌های پیش‌نویس را در یک عبور رو به جلو ارزیابی می‌کند. از آنجایی که توجه روی یک دنباله کوتاه تقریباً به اندازه یک عبور تک‌توکن روی یک مدل بزرگ سریع است، مدل هدف می‌تواند هر توکن پیش‌نویس را به صورت موازی تأیید یا رد کند.

زمانی که توکن‌های پیش‌نویس با آنچه مدل هدف تولید می‌کرد مطابقت دارند (یعنی در محدوده احتمال پذیرش آن قرار می‌گیرند)، همه آن‌ها پذیرفته می‌شوند و فرآیند چندین موقعیت را هم‌زمان پیش می‌برد. هنگامی که یک توکن پیش‌نویس رد می‌شود، خط لوله به نمونه‌گیری از توزیع هدف در آن موقعیت بازمی‌گردد و پیش‌نویسی را از آنجا از سر می‌گیرد. بینش کلیدی که در مقاله اصلی به طور دقیق اثبات شده است این است که این طرح نمونه‌گیری رد، توزیع خروجی یکسان با آنچه از مدل هدف به تنهایی به دست می‌آید تولید می‌کند — بنابراین کیفیت از نظر ریاضی حفظ می‌شود.

نرخ پذیرش و جایی که کیفیت پیش‌نویس اهمیت دارد

ضریب افزایش سرعت مستقیماً با نرخ متوسط پذیرش توکن‌های پیش‌نویس متناسب است. اگر مدل پیش‌نویس 80٪ از توکن‌های پیشنهادی را بپذیرد و شما 4 توکن را پیش‌بینی کنید، عملاً به جای 1 توکن، حدود 3.2 توکن را در هر عبور رو به جلو مدل هدف پردازش می‌کنید. نرخ پذیرش در دنیای واقعی بسته به وظیفه به طور قابل توجهی متفاوت است:

• خروجی‌های تکراری یا قالبی (کدهای استاندارد، داده‌های ساختاریافته): پذیرش 85–95٪، افزایش سرعت 3 تا 4 برابر معمول
• چت عمومی و پیروی از دستورالعمل: پذیرش 60–80٪، افزایش سرعت 1.5 تا 2.5 برابر معمول
• تولید خلاقانه و زنجیره‌های استدلالی: پذیرش 40–65٪، بهبودهای نسبتاً کمتر

این بدان معناست که رمزگشایی پیشگویانه زمانی بهترین عملکرد را دارد که وظیفه ساختار قابل پیش‌بینی داشته باشد — که بخش بزرگی از بارهای کاری تولیدی LLM را توصیف می‌کند: تولید کد، خلاصه‌سازی اسناد با قالب‌بندی یکسان، و پاسخ‌های چت‌بات که از الگوهای prompt پیروی می‌کنند. این روش در وظایف خلاقانه باز که مدل پیش‌نویس به طور سیستماتیک از مدل هدف منحرف می‌شود، عملکرد ضعیف‌تری دارد.

رمزگشایی پیشگویانه خودکار و سرهای Medusa

اجرای یک مدل پیش‌نویس مجزا پیچیدگی عملیاتی اضافه می‌کند: باید دو مدل را هم‌مکان کنید، حافظه هر دو را مدیریت کرده و آن‌ها را همگام نگه دارید. چندین رویکرد برای رسیدگی به این موضوع پدیدار شده است.

رمزگشایی پیشگویانه خودکار از خود مدل هدف به عنوان تولیدکننده پیش‌نویس استفاده می‌کند اما با لایه‌های خروج زودهنگام. برای مثال، مدل 70 میلیاردی Llama 3.1 را می‌توان برای پیش‌نویسی پیکربندی کرد که از لایه 32 خارج شود (به جای 80 لایه کامل)، سپس برای تأیید دوباره وارد پشته کامل شود. این کار نیاز به یک مدل مجزا را برطرف می‌کند اما نیازمند پشتیبانی معماری برای پرش از لایه است.

Medusa که در دانشگاه پرینستون توسعه یافته و در سال 2024 به صورت Open Source منتشر شده است، رویکرد متفاوتی دارد: چندین سر پیش‌بینی اضافی به حالت پنهان نهایی مدل هدف اضافه می‌کند که هر کدام توکن‌هایی را در موقعیت‌های +1، +2، +3 و غیره پیش‌بینی می‌کنند. این سرها به صورت ارزان از طریق Fine-tuning نظارت‌شده بر روی خروجی‌های موجود مدل آموزش داده می‌شوند. Medusa-2 با ترکیب تأیید درختی، بدون یک مدل پیش‌نویس مجزا و با کاهش حداقل دقت، به افزایش سرعت 2.8 برابر در Vicuna-13B دست می‌یابد.

SpecInfer از دانشگاه Carnegie Mellon این ایده را با استفاده از یک درخت از دنباله‌های پیش‌نویس به جای یک دنباله واحد که در یک عبور رو به جلو دسته‌ای تأیید می‌شود، گسترش می‌دهد. این احتمال پذیرش را افزایش می‌دهد زیرا مدل هدف فقط باید یک مسیر معتبر در درخت پیدا کند.

ملاحظات تولیدی

استقرار رمزگشایی پیشگویانه در تولید نیازمند توجه به چندین عامل است که اغلب در Benchmarkها کمتر برآورد می‌شوند:

• تعامل اندازه دسته: رمزگشایی پیشگویانه بیشترین مزیت را در اندازه دسته 1 (موارد استفاده تعاملی) فراهم می‌کند. با افزایش اندازه دسته، تنگنای پهنای باند حافظه نسبت به عملیات حسابی کاهش می‌یابد و هزینه اجرای مدل پیش‌نویس کمتر توجیه‌پذیر می‌شود. در اندازه دسته بالای 32–64، تولید استاندارد اغلب از رمزگشایی پیشگویانه بهتر عمل می‌کند.
• مدیریت KV Cache: مدل پیش‌نویس و مدل هدف هر دو KV Cache را نگهداری می‌کنند. تخصیص حافظه باید هر دو را در نظر بگیرد و این امر فریم‌ورک‌های سرویس‌دهی را که استفاده از Cache را بهینه می‌کنند (مانند PagedAttention در vLLM) پیچیده می‌کند.
• انتخاب مدل پیش‌نویس: مدل پیش‌نویس باید از خانواده مشابه مدل هدف باشد. استفاده از Llama-3-8B برای پیش‌نویسی برای Llama-3-70B نرخ پذیرش بسیار بالاتری نسبت به استفاده از مدلی که روی مجموعه داده متفاوت یا با توکن‌سازی متفاوت آموزش دیده است، به همراه دارد.
• مبادله تأخیر در برابر توان عملیاتی: رمزگشایی پیشگویانه بیشتر برای تأخیر (زمان تکمیل یک درخواست) بهینه‌سازی می‌کند تا توان عملیاتی (توکن در ثانیه در تمام درخواست‌ها). تیم‌هایی که استنتاج سرویس‌دهی با QPS بالا را اجرا می‌کنند باید ارزیابی کنند کدام معیار برای SLA آن‌ها مهم‌تر است.

مسیر آینده این تکنیک

استقرار رمزگشایی پیشگویانه در زیرساخت سرویس‌دهی Gemini در سال 2024 توسط Google گزارش شده است که هزینه استنتاج را برای بارهای کاری حساس به تأخیر 40٪ کاهش داده است. NVIDIA پشتیبانی از رمزگشایی پیشگویانه را در TensorRT-LLM 0.8 با سرهای پیش‌بینی چند توکن (MTP) یکپارچه کرد و این رویکرد اکنون یک ویژگی درجه یک به جای یک افزونه آزمایشی است.

مرز بعدی رمزگشایی پیشگویانه آبشاری است: استفاده از سه یا چند مدل در یک سلسله‌مراتب (مثلاً 1 میلیارد → 7 میلیارد → 70 میلیارد)، که در آن کوچک‌ترین مدل برای مدل متوسط و مدل متوسط برای مدل بزرگ پیش‌نویسی می‌کند. نتایج اولیه از تیم تحقیقاتی Meta نشان می‌دهد که کاهش تأخیر 4 تا 5 برابر در وظایف ساختاریافته با این رویکرد حاصل می‌شود، هرچند ردپای حافظه و پیچیدگی هماهنگی به همان نسبت افزایش می‌یابد.

همچنین کار فعالی روی رمزگشایی پیشگویانه برای مدل‌های مخلوط کارشناسان وجود دارد، جایی که الگوهای محاسباتی نامنظم مسیریابی کارشناسان پراکنده به روش‌های غیربدیهی با تأیید پیش‌نویس تعامل می‌کند. معماری Mixtral از Mistral به طور خاص برای این موضوع بررسی شده است؛ نتایج اولیه نشان می‌دهد دینامیک نرخ پذیرش به طور معناداری از مدل‌های متراکم متفاوت است.

نکات عملی قابل اقدام

• اگر بار کاری شما حساس به تأخیر است و اندازه دسته پایین (1–8) است، رمزگشایی پیشگویانه باید اولین بهینه‌سازی باشد که ارزیابی می‌کنید — فریم‌ورک‌هایی مانند vLLM با یک پرچم پیکربندی از آن پشتیبانی می‌کنند.
• مدل پیش‌نویس خود را با خانواده و داده‌های آموزشی مدل هدف مطابقت دهید. پیش‌نویس‌های نامتناسب نرخ پذیرش ضعیفی به شما می‌دهند و ممکن است در واقع کار را کندتر کنند.
• نرخ پذیرش را روی promptهای تولیدی واقعی خود اندازه‌گیری کنید، نه روی Benchmarkها. ساختار وظیفه برای پیش‌بینی افزایش سرعت مهم‌تر از اندازه مدل است.
• برای تیم‌هایی که نمی‌توانند یک مدل دوم اجرا کنند، سرهای Medusa را به عنوان یک راه میانی ارزیابی کنید — آن‌ها سربار حافظه حداقلی اضافه می‌کنند و می‌توانند در عرض چند ساعت روی یک GPU واحد Fine-tuning شوند.
• رمزگشایی پیشگویانه را برای استنتاج دسته‌ای با توان عملیاتی بالا به کار نگیرید مگر اینکه دینامیک اندازه دسته را در نقطه عملیاتی خاص خود تأیید کرده باشید.