ترکیب متخصصان Mixture of Experts

با انتشار مدل Mixtral 8x7B ، یک مدل از ترانسفورمرها تبدیل به داغ‌ترین موضوع در جامعه متن‌باز هوش مصنوعی شده است: Mixture of Experts یا به اختصار MoEs. در این پست ، نگاهی به بلوک‌های سازنده MoEs، نحوه آموزش آن‌ها و مواردی که باید در هنگام استنتاج در نظر گرفت، می‌اندازیم.

ترکیبی متخصصان (MOE) چیست؟

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

ترکیب متخصصان یک معماری یادگیری ماشین است که از چندین مدل کوچکتر برای انجام یک کار استفاده می کند. MoE به مدل ها امکان می دهد با محاسبات بسیار کمتری از قبل آموزش داده شوند. این بدان معناست که می توان مدل یا اندازه مجموعه داده را با همان بودجه محاسباتی یک مدل متراکم(dense) به طرز چشمگیری افزایش داد.

به طور خاص، یک مدل MoE باید به همان کیفیتی که همتای متراکم خود دارد، در طول پیش‌آموزش سریع‌تر دست یابد. این به این دلیل است که MoE می تواند از موازی سازی محاسبات برای بهره برداری از چندین پردازنده استفاده کند.

بنابراین، MoE دقیقاً چیست؟ در زمینه مدل های ترانسفورماتور، یک MOE از دو عنصر اصلی تشکیل شده است:

• لایه‌های Sparse MoE یک نوع لایه شبکه عصبی هستند که از چندین متخصص برای انجام یک کار استفاده می‌کنند. لایه‌های MoE تعداد معینی «متخصص» دارند (به عنوان مثال 8)، که در آن هر متخصص یک شبکه عصبی است. در عمل، متخصصان، Feedforward Neural Network هستند، اما می توانند شبکه های پیچیده‌‌‌تر یا حتی خود، یک MOE باشند که منجر به MOE های سلسله مراتبی شوند!
• یک شبکه گیت یا مسیریاب(router)، که تعیین می‌کند کدام توکن‌ها برای کدام متخصص ارسال می‌شوند. به عنوان مثال، در تصویر زیر، توکن 'More' برای دومین متخصص و توکن 'Parameters' به شبکه اول ارسال می شود. همانطور که در ادامه بررسی خواهیم کرد، می توانیم یک توکن را به بیش از یک متخصص ارسال کنیم. نحوه مسیریابی یک توکن به یک متخصص یکی از تصمیمات بزرگ هنگام کار با MoEs است - مسیریاب از پارامترهای آموخته شده تشکیل شده است و همزمان با بقیه شبکه آموزش داده شده است.

بنابراین، به طور خلاصه، در MoEs ما هر لایه FFN مدل ترانسفورماتور را با یک لایه MoE جایگزین می کنیم که از یک شبکه گیت و تعدادی متخصص تشکیل شده است.

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

• آموزش: MOEها به طور قابل توجهی در آموزش کارآمدی محاسباتی دارند، اما در زمان تنظیم دقیق برای تعمیم تلاش می‌کنند که منجر به اورفیت می‌شود.
• استنتاج: اگرچه یک MOE ممکن است پارامترهای زیادی داشته باشد، اما تنها از برخی از آنها در طول استنتاج استفاده می شود. این مسئله، منجر به استنتاج بسیار سریعتر در مقایسه با یک مدل متراکم با همان تعداد پارامتر می شود. با این حال، تمام پارامترها باید در RAM بارگذاری شوند، بنابراین نیاز به حافظه زیادی دارند. به عنوان مثال، با توجه به MoE مانند Mixtral 8x7B، ما باید VRAM کافی برای نگهداری یک مدل پارامتر متراکم 47B داشته باشیم. حالا چرا پارامترهای 47B و نه 8x7B = 56B؟ به این دلیل که در مدل‌های MoE، تنها لایه‌های FFN به‌عنوان متخصصان در نظر گرفته می‌شوند و بقیه پارامترهای مدل به اشتراک گذاشته می‌شوند. در عین حال، با فرض استفاده از دو متخصص در هر توکن، سرعت استنتاج (FLOPs) مانند استفاده از یک مدل 12B (در مقابل مدل 14B) است، زیرا ضرب های ماتریس 2x7B را محاسبه می کند، اما با برخی از لایه های مشترک.

اکنون که ما یک ایده تقریبی از چیستی MOE داریم، بیایید نگاهی به تحولات تحقیقاتی که منجر به اختراع آنها شده است بیندازیم.

تاریخچه مختصر ترکیب متخصصان

ریشه‌های MoEs از مقاله 1991 Adaptive Mixture of Local Experts می‌آید. ایده، شبیه به روش‌های ensemble، داشتن رویه‌ای تحت نظارت برای سیستمی از شبکه‌های مجزا بود که هر کدام زیرمجموعه‌ای متفاوت از دیتای آموزش را مدیریت می‌کردند. هر شبکه مجزا (یا متخصص) در ناحیه متفاوتی از فضای ورودی است. متخصص چگونه انتخاب می شود؟ یک شبکه گیت ، وزن ها را برای هر متخصص تعیین می کند. در طول تمرین، هم متخصص و هم گیت، آموزش می بینند.

بین سال‌های 2010-2015، دو حوزه تحقیقاتی مختلف به پیشرفت‌ بعدی ترکیب متخصصان کمک کردند:

• متخصصان به عنوان اجزاء: در سیستم‌های سنتی MOE، کل سیستم از یک مدل گیت و چندین مدل کوچکتر تشکیل می‌شود. این مدل‌های کوچکتر می‌توانند از انواع مختلفی از الگوریتم‌های یادگیری ماشین، مانند SVMها، فرآیندهای گاوسی یا شبکه‌های عمیق استفاده کنند. این متد امکان استفاده از MOE به عنوان لایه در یک شبکه چند لایه را فراهم می‌کند و این امر باعث می‌شود که مدل‌های MOE هم بزرگ و هم کارآمد باشند.
• محاسبه شرطی: در روش محاسبه شرطی، Yoshua Bengio روش‌هایی را برای فعال یا غیرفعال کردن پویای مؤلفه‌ها بر اساس توکن ورودی پیشنهاد کرد. این کار باعث می‌شود که شبکه بتواند داده‌های ورودی را به طور کارآمدتری پردازش کند و عملکرد بهتری داشته باشد.

این کارها منجر به آزمایش ترکیب متخصصان در زمینه NLP شد. به طور مشخص گوگل این ایده را در LSTM 137B (معماری NLP در آن زمان) استفاده کرد. استنتاج در یک مقیاس بالا و بر ترجمه متمرکز بود اما با چالش‌های زیادی مانند هزینه‌های ارتباطی بالا و بی‌ثباتی در آموزش روبرو بود.

‏MoEها امکان آموزش مدل‌های چند تریلیون پارامتری را فراهم کرده‌اند، مانند مدل‌های سوئیچ ترانسفورر با 1.6 تریلیون پارامتر که به صورت کد متن‌باز منتشر شده است. همچنین MoEها در حوزه بینایی کامپیوتر مورد بررسی قرار گرفته‌اند، اما این پست بر روی حوزه پردازش زبان طبیعی (NLP) تمرکز خواهد کرد.

پراکندگی (Sparsity) چیست؟

‏Sparity از ایده محاسبات شرطی استفاده می کند. در حالی که در مدل‌های متراکم از تمام پارامترها برای همه ورودی‌ها استفاده می‌شود، پراکندگی به ما اجازه می‌دهد فقط برخی از بخش‌های کل سیستم را اجرا کنیم. ایده محاسبات مشروط به افراد اجازه می‌دهد تا اندازه مدل را بدون افزایش محاسبات مقیاس‌پذیر کند و از این رو، این امر منجر به استفاده از هزاران متخصص در هر لایه MoE شد. این تنظیم چند چالش را ایجاد می‌کند. به عنوان مثال، اگرچه اندازه‌های batch بزرگ معمولاً برای عملکرد بهتر هستند، اما اندازه‌های batch در MoE‌ها با جریان داده در متخصصان فعال به طور موثری کاهش می‌یابد. به عنوان مثال، اگر ورودی دسته‌ای ما شامل 10 توکن باشد، ممکن است پنج توکن در یک متخصص و پنج توکن دیگر در پنج متخصص مختلف قرار گیرند که منجر به اندازه‌های batch ناهموار و ناکارآمد شود. چگونه می توانیم این را حل کنیم؟ یک شبکه گیت آموزش داده شده (G) تصمیم می‌گیرد که به کدام متخصصان (E) بخشی از ورودی را ارسال کند:

در این تنظیم، همه متخصصان برای همه ورودی‌ها اجرا می‌شوند - این یک ضرب وزنی است. اما اگر G برابر با 0 باشد چه اتفاقی می افتد؟ اگر چنین باشد، نیازی به محاسبه عملیات مربوط به متخصصان نیست و از این رو در محاسبات صرفه جویی می کنیم. یک تابع گیتینگ (gating function) معمولی چیست؟ در سنتی‌ترین حالت، ما فقط از یک شبکه ساده با تابع softmax استفاده می‌کنیم. این شبکه یاد می‌گیرد که کدام متخصص را برای ارسال ورودی انتخاب کند.

مکانیسم های دیگر گیت مانند Noisy Top-K Gating هم می‌توان بررسی کرد. این رویکرد، مقداری نویز (قابل تنظیم) را معرفی می کند و سپس مقادیر بالای k را حفظ می کند.

• مقداری نویز اضافه می کنیم

• سپس kتا بالا را انتخاب می کنیم

• سافت مکس را اعمال می کنیم.

این روش برخی خواص جالب را به همراه دارد. با استفاده از k کم (مثلاً یک یا دو)، می‌توانیم آموزش و اجرای استنتاج را بسیار سریع‌تر از زمانی که بسیاری از متخصصان فعال هستند، انجام دهیم. چرا فقط متخصص برتر را انتخاب نکنیم؟ حدس اولیه این بود که مسیریابی به بیش از یک متخصص برای یادگیری نحوه مسیریابی به متخصصان مختلف لازم بود، بنابراین حداقل دو متخصص باید انتخاب می‌شدند. بخش «سوئیچ ترانسفورماتورها» این تصمیم را دوباره بررسی می‌کند.

توزیع بار توکن‌ها برای MoE

همانطور که قبلاً گفته شد، اگر همه توکن‌های ما فقط به چند متخصص محبوب ارسال شوند، آموزش ناکارآمد خواهد شد و در آموزش MoE ، شبکه گیتینگ برای فعال کردن همان چند متخصص همگرا می‌شود. این کار تقویت هم می‌شود زیرا متخصصان مورد علاقه سریع‌تر آموزش دیده و در نتیجه بیشتر انتخاب می‌شوند. برای کاهش این مشکل، یک loss کمکی اضافه می‌شود تا به دادن اهمیت برابر به همه متخصصان تشویق شود. این loss تضمین می‌کند که همه متخصصان، تقریباً تعداد مساوی از نمونه‌های آموزشی را دریافت می‌کنند. بخش‌های بعدی همچنین مفهوم ظرفیت متخصص را بررسی خواهیم کرد که آستانه‌ای برای تعداد توکن‌هایی که می‌توانند توسط یک متخصص پردازش شوند را معرفی می‌کند. در ترانسفورمرها، loss کمکی از طریق پارامتر aux_loss مورد استفاده قرار می‌گیرد.

‏MoE و ترانسفورمرها

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

‏GShard هر لایه FFN را با یک لایه MoE با استفاده از گیتینگ top-2 هم در رمزگذار (encoder) و هم در رمزگشا (decoder) جایگزین می‌کند. تصویر بعدی نحوه نمایش آن را برای رمزگذار نشان می‌دهد. این تنظیم برای محاسبات در مقیاس زیاد بسیار مفید است: زمانی که ما به چندین دستگاه مقیاس می‌شویم، لایه MoE در بین دستگاه‌ها به اشتراک گذاشته می‌شود، در حالی که بقیه لایه‌ها تکرار می‌شوند.

برای حفظ بار متعادل و کارایی در مقیاس بالا، نویسندگان GShard علاوه بر loss کمکی مشابه آنچه در بخش قبلی مورد بحث قرار گرفت، چند تغییر دبگر هم ایجاد کردند:

مسیریابی تصادفی: در یک تنظیم top-2، ما همیشه متخصص برتر را انتخاب می‌کنیم، اما متخصص دوم با احتمالی متناسب با وزن آن انتخاب می‌شود.

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

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

این نکته مهم است زیرا به ما کمک می کند تا اندازه مدل و محاسبات مورد نیاز برای استنتاج را درک کنیم. به عنوان مثال، اگر یک مدل 47B با 8 متخصص داشته باشیم، می توانیم با محاسبات یک مدل متراکم 12B اجرا کنیم. این به این دلیل است که تنها 12B پارامتر در محاسبات استفاده می شوند و 2B پارامتر اضافی مربوط به متخصصان غیرفعال هستند.فرض کنید یک مدل با 4 متخصص داریم. هر متخصص 10 پارامتر دارد. اگر یک نمونه داده داشته باشیم، ممکن است شبکه گیتینگ تصمیم بگیرد که متخصص 1 و 2 فعال شوند. در این حالت، تنها 20 پارامتر (10 پارامتر برای متخصص 1 و 10 پارامتر برای متخصص 2) در محاسبات استفاده می شوند. 20 پارامتر اضافی مربوط به متخصصان غیرفعال هستند و هرگز استفاده نمی شوند.

سوئیچ بین ترانسفورمرها Switch Transformers

‏MoE ها پتانسیل زیادی دارند، اما آموزش و تنظیم دقیق آنها دشوار است. سوئیچ ترانسفورماتورها یک کار جدید و هیجان‌انگیز است که به این مشکلات می‌پردازد. نویسندگان این کار، حتی یک مدل MoE با 1.6 تریلیون پارامتر و 2048 متخصص را در Hugging Face منتشر کردند که با کتابخانه transformers قابل اجرا است. سوئیچ ترانسفورمرها، سرعت پیش‌آموزش را 4 برابر T5-XXL افزایش می‌دهد.

همانطور که در GShard، نویسندگان لایه‌های FFN را با یک لایه MoE جایگزین کردند. مقاله سوئیچ ترانسفورمرها، لایه Switch Transformer را پیشنهاد می‌کند که دو ورودی (دو توکن مختلف) را دریافت می‌کند و دارای چهار متخصص است.

برخلاف ایده اولیه استفاده از حداقل دو متخصص، سوئیچ ترانسفورمرها از استراتژی متخصص واحد ساده شده استفاده می‌کنند. اثرات این رویکرد عبارتند از:

• کاهش محاسبات روتر: با نیاز به مسیریابی برای تنها یک متخصص، بار محاسباتی روتر کاهش می یابد.
• کاهش حداقل اندازه دسته هر متخصص: با ارسال نیمی از توکن ها به هر متخصص، حداقل اندازه batch مورد نیاز برای هر متخصص نصف می شود.
• کاهش هزینه‌های ارتباطی: با ارسال توکن‌های کمتر بین متخصصان، هزینه‌های ارتباطی کاهش می‌یابد.
• حفظ کیفیت: با وجود استفاده از یک متخصص واحد، مدل همچنان عملکرد با کیفیت بالا را حفظ می‌کند.

سوئیچ ترانسفورمرها همچنین مفهوم ظرفیت متخصص را بررسی می‌کنند که آستانه‌ای برای تعداد توکن‌هایی که می‌توانند توسط یک متخصص پردازش شوند را تعیین می‌کند.

قدم به قدم

1. ورودی‌ها بلوک MoE دو توکن(آبی، نارنجی) دریافت می‌کند.
2. شبکه گیتینگ پردازش X1 (توکن آبی) را انجام می‌دهد و تعیین می‌کند Expert 2 باید فعال شود.
3. Expert 2 پردازش X1 (توکن آبی) را انجام می‌دهد.
4. شبکه گیتینگ پردازش X2 (توکن نارنجی) را انجام می‌دهد و تعیین می‌کند Expert 1 باید فعال شود.
5. Expert 1 پردازش X2 (توکن نارنجی) را انجام می‌دهد.
6. تابع فعال سازی ReLU خروجی های متخصصان را پردازش کرده و خروجی نهایی را تولید می کند.