کاهش حافظه پنهان KV با مکانیسم توجه جدید گام

مقدمه

استفاده روزافزون از مدل‌های زبان بزرگ (LLM) و ظهور پارادایم‌های جدید در استنتاج، چالش استنتاج کارآمد در مقیاس بزرگ را برجسته کرده است. یک گلوگاه مهم، حافظه پنهان Key-Value (KV) در مکانیسم‌های توجه سنتی است که به‌صورت خطی با اندازه دسته‌ای و طول دنباله افزایش می‌یابد و به یک «مصرف‌کننده حافظه» تبدیل می‌شود که مانع مقیاس‌پذیری و گسترش LLMها می‌شود.

در حالی که گونه‌هایی مانند MQA، GQA و MLA برای رفع این مشکل ظهور کرده‌اند، اغلب برای حفظ عملکرد در محدودیت‌های شدید حافظه با مشکل مواجه می‌شوند یا پیچیدگی‌هایی را معرفی می‌کنند که چالش‌های مهندسی و مشکلات سازگاری ایجاد می‌کنند.

توجه چند ماتریسی فاکتوریزاسیون (MFA)

مقاله جدیدی از محققان استپس، دانشگاه تسینگ‌هوا و سایر موسسات، معماری مکانیسم توجه جدیدی را معرفی می‌کند: توجه چند ماتریسی فاکتوریزاسیون (MFA) و نوع آن MFA-Key-Reuse (MFA-KR). این مکانیسم به طور قابل توجهی هزینه استنتاج مدل زبان را کاهش می دهد و در عین حال عملکرد را بهبود می بخشد.

MFA و MFA-KR نه تنها در عملکرد از MLA پیشی می گیرند، بلکه با عملکرد MHA سنتی نیز مطابقت دارند و در عین حال استفاده از حافظه پنهان KV را تا 93.7 درصد کاهش می دهند.

MFA برای سادگی، سهولت بازتولید، حساسیت کم به ابرپارامترها و سازگاری با روش های مختلف تعبیه موقعیت طراحی شده است.

رویکرد و تجزیه و تحلیل MFA

تیم تحقیقاتی طراحی کلی و ظرفیت مکانیسم های توجه را تجزیه و تحلیل کردند و دو بعد مهم مرتبط با ظرفیت را شناسایی کردند. این تجزیه و تحلیل منجر به توسعه روش های تحلیلی و اصول طراحی جدید شد.

آنها مفهوم توجه چند سر تعمیم یافته (GMHA) را به عنوان چارچوبی متحد برای درک انواع مختلف MHA معرفی کردند.

این تیم همچنین محاسبات و ذخیره سازی کلید-مقادیر را از منظر استنتاج بررسی کرد و ظرفیت مدل را از منظر تجزیه بررسی کرد.

توجه دوخطی کاملاً پارامتری شده (FPBA) به عنوان حد بالایی نظری عملکرد ایجاد شد. آنها دریافتند که MHA و انواع آن تجزیه رتبه پایین FPBA هستند.

مقایسه با MQA و MLA

تجزیه و تحلیل بر روی دو طرح بهبود نماینده متمرکز شد: توجه چند پرس و جو (MQA) و توجه نهفته چند سر (MLA).

MQA از یک استراتژی اشتراک گذاری پارامتر تهاجمی تر استفاده می کند، جایی که همه سرهای توجه از یک مجموعه پارامتر کلید-مقدار استفاده می کنند. این امر استفاده از حافظه را کاهش می دهد اما ممکن است بر بیان مدل تأثیر بگذارد.

MLA یک فضای نهفته مشترک برای فشرده سازی پارامتر معرفی می کند، اما قدرت بیان واقعی با کوچکترین بعد محدود می شود، به این معنی که افزایش ابعاد میانی به طور قابل توجهی عملکرد را بهبود نمی بخشد.

نوآوری های کلیدی MFA

توسعه MFA با هدف ایجاد یک مکانیسم توجه که مصرف منابع را به حداقل می رساند و در عین حال به محدودیت های عملکرد نظری نزدیک می شود، هدایت شد. طراحی MFA شامل سه نوآوری کلیدی است:

افزایش قابل توجه تعداد و ابعاد سرهای توجه برای به حداکثر رساندن ظرفیت مدل.
استفاده از یک استراتژی تجزیه رتبه پایین تهاجمی برای حفظ راندمان پارامتر در حین گسترش تعداد و ابعاد سر توجه.
استفاده از طراحی یک سر کلید-مقدار برای به حداقل رساندن مصرف حافظه، حتی با افزایش پیچیدگی مدل.

اندازه گیری ظرفیت و مقایسه

برای تجزیه و تحلیل بیشتر MFA و سایر مکانیسم های توجه، این تیم دو معیار کلیدی را معرفی کرد:

رتبه موثر کل (TER): حاصل ضرب تعداد سرهای توجه و رتبه فاکتوریزاسیون در هر سر (FRH).
بعد زیرفضای نهفته مشترک (SLSD): بعد فضای پنهان که توسط همه سرهای توجه به اشتراک گذاشته می شود.

MFA در مقایسه با MQA به SLSD و TER بالاتری دست می یابد.

در مقایسه با MLA، MFA با بودجه پارامتری مشابه، به اندازه حافظه پنهان KV کوچکتر و TER بالاتری دست می یابد، در حالی که SLSD قابل مقایسه ای را حفظ می کند.

در مقایسه با MHA سنتی، MFA دارای TER بالاتری است، حتی اگر SLSD آن کوچکتر باشد.

نتایج تجربی

آزمایش‌های گسترده‌ای برای ارزیابی عملکرد معماری جدید در مقیاس‌های بزرگ‌تر انجام شد و مدل‌هایی از 1B تا 7B پارامتر و داده‌های آموزشی از 10B تا 1T آزمایش شدند.

MFA قابلیت‌های مقیاس‌پذیری قابل مقایسه با MHA سنتی را نشان داد و عملکرد عالی را حتی در مقیاس‌های بزرگ‌تر حفظ کرد.

در حالی که MFA-KR عملکرد کمی پایین‌تری را نشان داد، روند مقیاس‌پذیری آن با MHA همسو بود. مزایای صرفه‌جویی در حافظه MFA و MFA-KR با اندازه مدل به گسترش خود ادامه داد، به‌طوری که MFA به 87.5 درصد صرفه‌جویی در حافظه و MFA-KR به کاهش استفاده از حافظه تا 6.25 درصد در بزرگترین مقیاس دست یافت.

مطالعات ابلیشن

مطالعات ابلیشن اثربخشی MFA و MFA-KR را تأیید کرد. مزایای عملکرد آنها نیز در روش های مختلف کدگذاری موقعیتی جریان اصلی تایید شد.

چشم انداز

MFA با طراحی ساده، بهبودهای قابل توجهی را ارائه می دهد و به طور موثر به گلوگاه حافظه در استنتاج LLM بدون افزودن پیچیدگی مهندسی اضافی می پردازد. این مکانیسم به طور یکپارچه در اکوسیستم ترانسفورمر موجود ادغام می شود و کاربرد LLM ها را در سناریوهای مختلف تسریع می کند.