مدل MobileLLM

چند وقت پیش با مدل MobileLLM آشنا شدم. یک مدل LLM سبک که شرکت متا برای دستگاه‌های با قدرت سخت افزاری کم مثل موبایل پیشنهاد کرده است. چون نکات آموزنده این مقاله زیاد بود، تصمیم گرفتم درباره آن بنویسم. با من همراه باشید تا از این مدل جالب برای شما بگویم…

آشنایی با مدل MobileLLM

مدل MobileLLM در مقاله‌ای با عنوان زیر توسط شرکت متا در سال 2024 منتشر شد:

MobileLLM: Optimizing Sub-billion Parameter Language Models for On-Device Use Cases [Link]

بیایید نگاهی به عنوان مقاله بندازیم:

• MobileLLM: دو کلمه Mobile و LLM نشان می‌دهد که با یک مدل زبانی طرف هستیم که مناسب دستگاه‌های سبک مانند موبایل است.
• Sub-billion Parameter Language Models: اگرچه، این مدل در عنوان خود LLM را یدک می‌کشد، اما به بزرگی مدل‌های LLM بیلیونری مانند Llama GPT-4 و غیره نیست. در واقع، باید مدل MobileLLM را در دسته شبکه‌های مگس‌وزن قرار داد! 😁 دو مدل اصلی MobileLLM شامل 125 و 350 میلیون پارامتر است.
• On-Device Use Cases: به برنامه‌ها و سناریوهایی گفته می‌شود که محاسبات، پردازش داده‌ها یا عملیات دیگر به‌صورت مستقیم روی خود دستگاه انجام می‌شود. یعنی، نیازی به استفاده از سرویس ابری یا سرورهای خارجی نیست. این یعنی دستگاه به تنهایی مجهز به سخت‌افزار و نرم‌افزار لازم برای انجام این وظایف هست؛ مثل، موبایل و تبلت، ساعت هوشمند، دستگاه‌های هوشمند خانگی، خودروها و دستگاه‌های اینترنت اشیا (IoT) و غیره.
• Optimizing: قرار هست مدل‌های سبک وزن را بهینه کنند. بهینه کردن به این معناست که ایده‌ها و فاکتورهای مختلف را طی آزمایش‌های متعدد بررسی می‌کنند تا بهترین نتیجه حاصل شود.

بیشتر بدانید تا همین چند روز پیش، نسخه‌های مختلف مدل MobileLLM بهترین عملکرد را در بین رقیبان مگس‌وزن داشتند. اما اخیرا، مدل SmolLM از هاگینگ‌فیس رو شده که طبق گزارش‌ها از مدل MobileLLM بهتر هست.

تذکر مدل MobileLLM یک مدل زبانی بزرگ یا Large Language Model هست. بنابراین، پیش‌نیاز این پست، آشنایی با شبکه ترنسفورمر، مدل زبانی و LLM هست. اگر آشنایی ندارید، پیشنهاد می‌کنم مطالب زیر را از هوسم مطالعه کنید.

شبکه ترنسفورمر مدل زبانی چیست؟ LLM چیست؟

چالش اجرای LLM در موبایل

در بخش Introduction مقاله، به یکسری از چالش‌های اساسی اجرای مدل LLM روی موبایل‌ها اشاره شده است. در این بخش، می‌خواهم درباره این چالش‌ها بنویسم. این سه چالش مهم عبارت است از:

• مصرف RAM موبایل در LLM-ها
• مصرف باتری
• سرعت تولید متن در ثانیه

قبل از اینکه درباره چالش‌های بالا بنویسم، می‌خواهم کمی درباره ساختار و سلسله مراتب حافظه (Memory) در موبایل‌های امروزی توضیح دهم.

ساختار حافظه در موبایل

شکل روبرو، چهار مولفه مهم در موبایل‌های امروزی را نشان می‌دهد:

• CPU/NPU/GPU: پردازنده‌ها هستند. وظیفه انجام عملیات پردازشی مختلف در موبایل را برعهده دارند.
• Flash Memory: میزان حافظه دائمی موبایل هست. همان حافظه‌ای که فایل‌های مختلف مانند عکس و ویدئو در آن ذخیره می‌شود.
• DRAM: میزان حافظه موقت هست که معمولا با نام RAM شناخته می‌شود.
• SRAM Cache: معمولا به نام کش شناخته می‌شود. البته، کمتر در دنیای موبایل فروش‌ها درباره آن صحبت می‌شود. ظرفیت بسیار کمی دارد!

ظرفیت Flash Memory قابل توجه هست؛ اما، در عمل این DRAM هست که برای اجرای یک اپلیکیشن موبایلی با سرعت بالا اهمیت دارد. درواقع، حین مکالمه با یک Chatbot در موبایل، مدل LLM در DRAM هست. این درحالی است که ظرفیت DRAM چندان زیاد نیست. طبق جدول زیر، برای موبایل‌های امروزی حداکثر 12 گیگابایت هست. حالا نگید نه داش آخرین مدل سامسونگ 24 گیگه! این اعداد رو مقاله گفته! 😁 آیا مدلی مانند Llama-7B (مدل 7 میلیارد پارامتری) در این DRAM جا می‌شود؟! به این سوال، در چالش اول جواب می‌دهم. اما خوب هست که خودتان به این سوال فکر کنید…

اما کارکرد SRAM و تفاوت آن با DRAM چیست؟ SRAM یک حافظه کم حجم هست که سرعت بالاتری نسبت به DRAM دارد و گران‌تر است. یعنی، مدت زمان دسترسی پردازنده (CPU/NPU/GPU) به SRAM کمتر از DRAM هست. فرض کنید، یک شبکه ترنسفورمر 12 لایه داریم که در DRAM بارگذاری (Load) شده است. حالا می‌خواهیم یک ورودی به این مدل ترنسفورمر بدهیم و خروجی ترنسفورمر (خروجی لایه دوازدهم) را دریافت کنیم. ورودی هم در DRAM قرار دارد. پردازش/محاسبات هم بر عهده پردازنده (CPU/NPU/GPU) است. طبیعتا، محاسبات هم لایه به لایه انجام می‌شود. یعنی، ورودی به لایه اول ترنسفورمر و دریافت خروجی، سپس خروجی لایه اول به عنوان ورودی به لایه دوم و دریافت خروجی و … بالاخره نقش SRAM چیست؟ قبل از انجام پردازش/محاسبات یک لایه ترنسفورمر، پارامترهای آن لایه در SRAM ریخته می‌شود. با توجه به سرعت دسترسی CPU به SRAM، خروجی لایه مدنظر سریع‌تر بدست می‌آید. جالب بود نه؟ نه! 😜

پس به‌صورت خلاصه این شد که:

1. یک مدل ترنسفورمر با L لایه در DRAM بارگذاری می‌شود.
2. داده ورودی در DRAM بارگذاری می‌شود.
3. یک حلقه به ازای تعداد لایه‌ها:
   1. پارامترهای لایه i ترنسفورمر از DRAM به SRAM منتقل می‌شود.
   2. پردازنده (CPU/NPU/GPU) خروجی لایه i را حساب می‌کند.
   3. بازگشت به ابتدای حلقه

حالا می‌توانیم سه چالش مطرح شده درباره اجرای LLM روی موبایل را بررسی کنیم…

چالش اول: مصرف RAM موبایل (DRAM) در LLM

درحال حاضر، مقدار DRAM برای موبایل‌های سطح بالا بین 6 گیگابایت (آیفون 15) تا 12 گیگابایت (گوگل پیکسل 8 پرو) هست. شاید بگویید که خب این ظرفیت 6 و 12 گیگابایت که زیاد هست. بله، اما نکته مهم اینجاست که ما نمی‌توانیم مدلی در موبایل قرار دهیم که کل این ظرفیت 6 تا 12 گیگ را پر کند! این ظرفیت DRAM سهم بخش‌های مختلف موبایل از جمله سیستم عامل و سایر اپلیکیشن‌های موبایل هم هست! در مقاله ذکر شده که یک اَپ موبایل نباید بیشتر از 10% DRAM را اشغال کند. یعنی، برای DRAM با 12 گیگ ظرفیت نهایتا 1.2 گیگابایت برای ما دردسترس هست! حالا، فرض کنید ما بخواهیم مدل Llama-7B (مدل 7 میلیارد پارامتری) را حتی با وزن‌های 8 بیتی در DRAM ذخیره کنیم. به نظر شما چقدر از ظرفیت رم را اشغال می‌کند؟حدود 7 گیگابایت! 🤯 پس این چالش انگیزه‌ای هست که مدل‌های LLM کوچکتر (پارامترهای کمتر) ساخته شود. به همین خاطر در عنوان مقاله گفته شده:

Sub-bilion parameter LLMs

چالش دوم: مصرف باتری در LLM

مدل LLaMA v2 با 7 بیلیون پارامتر، به ازای هر توکن 0.7 ژول انرژی مصرف می‌کند. به نظر شما زیاد هست یا کم؟! زیاده بابا جان زیااااد! ببینید، یک آیفون با شارژ کامل حدودا 50 هزار ژول انرژی دارد. مقاله می‌گوید، هر 64 توکن حدود 0.2% از باتری را مصرف می‌کند. یعنی، کل باتری در کمتر از دو ساعت مکالمه با مدل تَه می‌کشد. البته، من حساب و کتاب که کردم، هر 64 توکن حدود 0.08% باتری مصرف می‌کند. اما شاید این محاسبه روی کاغذ هست و بخش‌های جانبی هم باتری می‌خورند. پس، نتیجه اینکه مدل‌های هیولای LLM هم DRAM را اشغال می‌کنند و هم قاتل باتری هستند!

اما، نسخه 350 میلیون پارامتری مدل MobileLLM، حدود 0.035 ژول برای هر توکن مصرف می‌کند. مصرف انرژی لاما چقدر بود؟ 0.7 ژول! یعنی، نسخه 350 میلیونی 20 برابر مصرف انرژی کمتری دارد. به عبارتی، باتری فول شارژ برای یک روز مکالمه کافی است. اوضاع خیلی فرق کرد!

چالش سوم: سرعت مدل MobileLLM در موبایل

سرعت تولید متن مدل لامای 7 بیلیونی روی آیفون، 3~6 توکن بر ثانیه هست. کم هست یا زیاد؟ کم! یعنی، چت‌بات در هر ثانیه 4 5 کلمه برای ما می‌نویسد. اگر برای شما ملموس نیست، بروید یک سوال از ChatGPT بپرسید، ببینید در هر ثانیه چند توکن (همان کلمه) برای شما می‌نویسد. اما مدل 125 میلیونی این مقاله حدود 50 توکن بر ثانیه سرعت دارد. یعنی، در عرض 10 ثانیه، روی موبایل 500 توکن برای ما می‌نویسد. خیلی هم خوب!

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

معماری مدل MobileLLM

معماری مدل MobileLLM مشابه با معماری مدل‌های GPT هست. مدل GPT یک شبکه ترنسفورمر هست که از تعدادی لایه متوالی دیکدر تشکیل شده است. در شکل زیر، نسخه‌های مختلف مدل GPT-2 را مشاهده می‌کنید. این مدل در سال 2019 توسط OpenAI معرفی شد. جزئیات فنی، کدها و وزن‌های این مدل کامل در دسترس هست. اون موقعا OpenAI هنوز ColsedAI نشده بود! 😜 

بیایید کمی مدل GPT-2 را بررسی کنیم؛ نسخه‌های مختلف مدل GPT-2 در حجم پارامترها با هم تفاوت دارند. مثلا، سبک‌ترین مورد 117 میلیون پارامتر دارد و بزرگ‌ترین نسخه یک مدل 1.5 بیلیونی است. این تفاوت پارامترها از ]چهار فاکتور مهم نشات می‌گیرد:

• تعداد لایه‌ها (مدل سبک با 12 لایه و مدل سنگین با 48 لایه)
• ابعاد ویژگی مخفی مدل یا Model Dimensionality (مدل سبک 768 الی مدل سنگین با 1600)
• تعداد سَرهای اتنشن یا Attention Heads (مدل 117 میلیونی 12 اتنشن هِد دارد)
• سایز امبدینگ یا Vocab Size (همان لایه ورودی که به هر توکن یک بردار تخصیص می‌دهد)

تعداد لایه‌ها که از روی شکل مشخص هست؛ اما ممکن هست، ابعاد ویژگی مدل (Model Dimensionality) برای شما مبهم باشد؛ این فاکتور معادل با ابعاد ویژگی ورودی انکدر/دیکدر در ترنسفورمر هست. به شکل زیر نگاه کنید؛ لایه انکدر/دیکدر به‌گونه‌ای طراحی شده است که فارغ از محاسبات داخل انکدر/دیکدر، ابعاد ویژگی خروجی به اندازه ابعاد ویژگی ورودی باشد. در تعریف لایه انکدر/دیکدر در پایتورچ، آرگومانی بنام d_model داریم که معادل با همین Model Dimensionality هست. راستی، این جمله را هم به خاطر داشته باشید:

d_model تعیین کننده پهنای مدل هست. d_model کوچک/بزرگ معادل است با مدل ترنسفورمر لاغر/چاق!

البته، لازم به ذکر هست که مدل MobileLLM در جزئیات تشابه زیادی با مدل Llama دارد. خب به نظرم این تشابه قابل انتظار هست؛ چون هر دو مدل MobileLLM و Llama را متا ارائه کرده است. ضمن اینکه، هردو بروز هستند و در همین یک سال گذشته سروصدا کرده‌اند. اگر به سورس کدهای MobileLLM نگاه کنید، اثر Llama را به وضوح خواهید دید. محققان مقاله MobileLLM همان کدهای Llama در کتابخانه HuggingFace را مودیفای کرده‌اند. حالا نمیدونم بجای مودیفای میگفتم توسعه خوب بود یا نه! کلا، منظورم اینه که همون کد Llama رو برداشتن و یکسری جاها رو تغییر دادن و MobileLLM رو ساختن.

اما ایده‌های مقاله MobileLLM چیست؟ درادامه، در بخش‌های مجزا درباره ایده‌ها صحبت می‌کنم. اما، حواستان باشد که برای درک بعضی ایده‌ها باید با معماری شبکه ترنسفورمر آشنا باشید. ایده‌هایی که در ادامه توضیح می‌دهم عبارت است از:

• عمق مدل MobileLLM
• اشتراک پارامترهای Embedding
• اشتراک گذاری لایه‌ها
• تعداد Head و KV-Head

عمق مدل MobileLLM

محققان در این مقاله به این نتیجه رسیده‌اند که بهتر هست عمق مدل‌های سبک زیاد باشد. این نتیجه با باور رایج درباره ترنسفورمر تفاوت دارد؛ یک باور رایج در این حوزه می‌گوید که عملکرد مدل‌های ترنسفورمر در درجه اول توسط سه مولفه تعداد پارامترها، اندازه مجموعه داده آموزشی و تعداد تکرارهای آموزش تعیین می‌شود. همچنین، این باور می‌گوید که مسائل طراحی مدل تاثیر ناچیزی در عملکرد مدل ترنسفورمر دارند. اما محققان مقاله MobileLLM ادعا می‌کنند که تحقیقات آنها نشان می‌دهد که ممکن هست این باورها برای مدل‌های کوچک صدق نکند. محقق‌ها 19 مدل (9 مدل 125 میلیونی و 10 مدل 350 میلیونی) آموزش داده‌اند و به این نتیجه رسیده‌اند که:

مدل‌های عمیق و لاغر عملکرد بهتری نسبت به مدل‌های کم‌عمق و پهن دارند.

در جمله بالا، معنی عبارت‌های عمیق و کم‌عمق که مشخص هست؛ به تعداد لایه‌های مدل ترنسفورمری اشاره دارد. درباره مدل لاغر و پهن هم بالاتر صحبت کردیم. به میزان طول ویژگی شبکه یا همان d_model اشاره دارد. نتایج آزمایش‌های این مقاله نشان می‌دهد که یک شبکه 125 میلیون پارامتری با 30 یا حتی 42 لایه بهتر از مدل‌های 12 لایه عمل می‌کند. درست هست که تعداد لایه‌ها فرق زیادی دارد، اما با تغییر طول ویژگی می‌توان به مدل‌های با سایز تقریبا برابر رسید.

اشتراک پارامترهای Embedding

در ساختار LLM دو لایه Embedding ورودی و خروجی داریم. طبق شکل زیر این دو لایه به‌صورت زیر عمل می‌کنند:

• لایه امبدینگ ورودی: اندیس توکن‌ها را دریافت و یک بردار ویژگی به طول مشخص (مثلا 512) برای هر توکن ارائه می‌کند.
• لایه امبدینگ خروجی: نقش کلاسیفایر را دارد. بردارهای خروجی ترنسفورمر (همان ابعاد 512) را دریافت می‌کند و یک عدد (همان شماره توکن) را در خروجی می‌دهد.

لایه امبدینگ پارامترهای زیادی دارد. به عنوان مثال، اگر امبدینگ شامل 32 هزار توکن و هر توکن برداری به طول 512 باشد، حدود 16 میلیون پارامتر ایجاد می‌شود. حالا، ما دو لایه امبدینگ ورودی و خروجی داریم و این یعنی 32 میلیون پارامتر به کل مدل اضافه می‌شود. وقتی ما مدل بیلیونی داریم، 32 میلیون در آن گم هست، اما زمانی که مدل ما سبک باشد (همین 125 و 350 میلیونی)، 32 میلیون سهم قابل توجهی از کل مدل را شامل می‌شود. مثلا، در مدل لامای 7 بیلیونی، 32 میلیون پارامتر حدود 3.7% و در مدل 125 میلیونی حدود 20% از کل پارامترهاست.

در اینجا، محققان مقاله از تکنیک به اشتراک گذاری پارامترهای لایه امبدینگ ورودی و خروجی استفاده کرده‌اند. باتوجه به ماهیت دو لایه امبدینگ ورودی و خروجی و همچنین سایز یکسان ماتریس پارامترهای این دو لایه، می‌توان پارامترهای مشترک برای دو لایه تعریف کرد. این باعث می‌شود میزان پارامترهای امبدینگ از 32 به 16 میلیون پارامتر کاهش یابد.

لازم به ذکر هست که این ایده، جدید نیست و پیش از این پیشنهاد شده بود. مثلا، ما در دوره دیپ کاتالیست در پروژه مدل‌سازی زبان با شبکه LSTM همین به اشتراک گذاری لایه امبدینگ ورودی و خروجی را با نام Weight Tying پیاده‌سازی کردیم. پس تکنیک جدیدی نیست، اما به گفته محققان فیسبوک، به دلیل حجم زیاد پارامتر مدل‌های امروزی، نادیده گرفته شده بود.

محقق‌های مقاله وقتی اشتراک گذاری پارامترها را روی یک مدل 30 لایه 125 میلیون پارامتری انجام دادند، میانگین دقت 0.2% افت کرد. اما، نکته مثبت ماجرا این هست که بجای 16 میلیون پارامترِ کم شده، دو لایه به مدل اضافه کردند و میانگین دقت این مدل 32 لایه 0.4% بهتر شد. 😇

اشتراک گذاری لایه‌ها

اشتراک گذاری لایه‌ها ایده جذابی است که مطمئن هستم از آن لذت خواهید برد! تا الان دو ایده مهم را دیدیم: (1) بهبود عملکرد مدل با افزایش عمق مدل (افزایش تعداد لایه‌ها) و (2) اشتراک گذاری در لایه‌های امبدینگ. ایده اشتراک گذاری لایه‌ها این هست که:

بدون افزایش تعداد پارامترها و هزینه ذخیره‌سازی، تعداد لایه‌های مدل را افزایش دهیم.

نتیجه آزمایش‌های انجام شده در این مطالعه این هست که صرفا با تکرار کردن لایه‌های شبکه ترنسفورمر می‌توان عملکرد مدل را بهبود داد. برای انجام این کار حتی نیازی به تغییر معماری شبکه ترنسفورمر هم نیست! همان‌طور که در شکل زیر نشان داده شده، در این مطالعه به سه شکل مختلف اشتراک گذاری پارامترهای لایه‌ها را مورد بررسی قرار گرفته است. اما توضیح هریک از شکل‌ها:

1. مدل اصلی ترنسفورمر بدون هرگونه اشتراک گذاری لایه‌ها. رنگ بلوک‌ها نشان‌دهنده وضعیت اشتراک پارامترهاست. در این شکل، هر بلوک یک رنگ اختصاصی دارد که نشان می‌دهد از اشتراک گذاری پارامترها خبری نیست.
2. این حالت Immediate block-wise sharing نام دارد؛ هر بلوک دوبار به‌صورت متوالی تکرار شده است. مشاهده می‌کنید، ابتدا دو بلوک بنفش، سپس دو بلوک زرد و بعد هم دو بلوک نارنجی داریم. همین روند تا خروجی مدل ادامه دارد. پس، پارامترهای دو لایه همسایه یکسان هست.
3. این حالت Repeat-all-over sharing نام دارد؛ یعنی، اگر ترنسفورمر اصلی 30 لایه دارد، ابتدا یک بار 30 لایه را قرار می‌دهیم و دوباره 30 لایه را تکرار می‌کنیم. به عبارت دیگر، مدل ترنسفورمر 30 لایه را دو بار پشت هم تکرار کرده‌ایم! رنگ بلوک‌ها نشان می‌دهد که کدام لایه‌ها مشترک هستند و پارامترهای یکسانی دارند.
4. این حالت Reverse sharing نام دارد؛ مشابه با حالت c هست. با این تفاوت که مدل دومی را برعکس (Reverse) می‌کنیم و سپس به مدل ترنسفورمر اول می‌چسبانیم.

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

قبلا درمورد DRAM SRAM صحبت کردم؛ گفتم که داده بین SRAM و DRAM مبادله می‌شود. دسترسی پردازنده (CPU) به SRAM سریع‌تر از DRAM هست. SRAM ظرفیت پایینی دارد (در حد 20 مگابایت). به همین خاطر نمی‌توان پارامتر همه لایه‌ها را در کش یا SRAM ریخت. این حجم معمولا برای صرفا یک لایه ترنسفورمر کفایت می‌کند. بنابراین، چه خوب که وقتی یک لایه در SRAM هست، دو بار در محاسبات استفاده شود.

تعداد Head و KV-Head

این بخش نسبت به قبلی‌ها مشکل‌تر هست. بهتر هست با تصویر زیر از مفهوم چندسَری (Multi-head) در ترنسفورمر شروع کنیم. می‌دانید که در ترنسفورمر مراحل زیر را روی بردار توکن‌ها در ماژول Multi-head Attention اجرا می‌کنیم:

1. تنسور شامل بردار توکن‌ها (سایز LxF) را به h تکه در راستای F تقسیم می‌کنیم.
2. سپس به‌صورت جداگانه روی هرکدام ماژول اتنشن را اعمال می‌کنیم.
3. نهایتا، خروجی‌ها را با هم در همان راستای F ترکیب می‌کنیم.

مثلا، در شکل زیر ورودی 9×4 با h=3 به 3 هِد با سایز 3×4 تقسیم شده است. حالا شما بگویید: یک ورودی به 512×10 با h=32 به چه شکلی می‌شود؟ اول، جواب بده بعد ادامه متن رو بخون! خب، وقتی h=32 هست، یعنی 32 هِد (مسیر موازی) داریم. سایز تنسور هر مسیر یا شاخه هم از تقسیم 512 بر 32 بدست می‌آید. پس، سایز تنسور هر شاخه برابر با 16×10 هست. به عدد 32 تعداد هِد و عدد 16 سایز هِد گفته می‌شود. 

حالا با این مقدمه، به این نکته کلیدی توجه کنید: نکته کلیدی در انتخاب سایز هِد، مصالحه (trade-off) بین ترکیب غیرخطی بیشتر بین هدهای مختلف و معنی (Semantic) بیشتر در هر هِد هست. یعنی، با افزایش h به سمت ترکیب غیرخطی بیشتر بین هدهای مختلف می‌رویم. با کاهش h معنی بیشتری در هر هِد خواهیم داشت. چون سایز هر هِد بیشتر شده است.

در این مطالعه، دو کار مهم در راستای تعداد هِد و سایز هِد انجام شده است:

• سایز بهینه هِد مورد آزمایش قرار گرفته است.
• تعداد هدِ کمتری برای Key و Value نسبت به Query در نظر گرفته شده است.

مقاله می‌گوید، اکثر کارهای گذشته در حوزه مدل‌های کوچک یک مقدار یکسان را برای تعداد هِد Key-Value و Query درنظر می‌گرفتند. اما، مولفان مقاله با استفاده از تکنیک Grouped Query Attention (GQA) تعداد هد key-value را کاهش دادند. احتمالا با ایده GQA آشنا نیستید. همان‌طور که از عنوان آن مشخص هست، این تکنیک یک ماژول اتنشن هست که بر پایه گروه‌بندی Query کار می‌کند. در این تکنیک، کوئری به چند گروه تقسیم و سپس برای هر گروه یک Key-Value یکسان استفاده می‌شود. شکل زیر، به خوبی تفاوت GQA را با حالت MHA نشان می‌دهد. در حالت MHA، تعداد Query با Key-Value برابر است. درحالی‌که در GQA، با گروه‌بندی Query، به ازای هر Query از یک Key-Value استفاده شده است. 

حالا این GQA چه خاصیتی دارد؟

• برای کاهش سایز کش (Cache) در Key-Value پیشنهاد شده است.
• می‌تواند افزونگی (Redundancy) را در هد Key-Value کاهش دهد.

به نوعی، می‌توان GQA را شکلی از اشتراک‌گذاری وزن‌ها در نظر گرفت. چرا؟ چون، تعداد هد کوئری n برابر تعداد هد Key-Value است. درواقع، تعداد هد Key-Value را n بار تکرار می‌کنیم تا محاسبات درون اتنشن همراه با کوئری انجام شود.

آزمایش‌های مختلفی برای این بخش انجام شده است. زمانی‌که تعداد هد کوئری 16 درنظر گرفته شد، بهترین نتیجه بدست آمده است. همچنین، تعداد هد Key-Value از 16 به 4 کاهش یافته و تنها 0.2% میانگین دقت در مدل 350 میلیونی افت کرده، درحالی‌که سایز مدل 10% کاهش داشته است. حال، اگر این کاهش سایز را با افزایش ابعاد امبدینگ جبران کنیم، نتیجه جالب می‌شود. مثلا، در مدل 125 میلیونی میانگین دقت 0.4 بهتر می‌شود. یک انتقادی که به این مقاله دارم این هست که دائم بین دو مدل 125 و 350 میلیونی جابجا میشه. مثلا، اول میگه مدل 350 میلیونی 0.2% افت دقت و 10% کاهش سایز داشته. بعدش میگه در مدل 125 میلیونی 0.4% دقت بهتر شده! بهتر نبود اینا رو روی یک مدل گزارش کنی؟! 🤔

هوراااا، تمااام شد! البته، مطالب ارزشمند دیگری هم در مقاله هست. اما، به نظرم این آموزش تا همین اندازه کافیست. پیشنهاد می‌کنم مقاله MobileLLM را هم بخوانید. خیلی خوشحالم که بالاخره این کار را به انتها رساندم. لابه‌لای کارهای مختلف، زمانی باز می‌کردم که این آموزش را بنویسم. من کار روی MobileLLM را ادامه می‌دهم. یک ایده‌هایی دارم که اگر همه چیز خوب پیش برود، به زودی از آن باخبر خواهید شد. از کیفیت نهایی کار تقریبا راضی هستم. اما دوست دارم نظر شما را هم بدانم. لطفا نظرتان را برای من کامنت کنید.