GeekAlerts

جایی برای گیک‌ها

·

معماری Mixture of Recursions (MoR)، رویکردی نوین به ترنسفورمرها

معماری Mixture of Recursions (MoR)، رویکردی نوین به ترنسفورمرها
  • مدت زمان مطالعه: حدود ۸ دقیقه
  • پیش‌نیازها: آشنایی خیلی کلی با هوش مصنوعی و اینکه مدل‌های زبانی چی هستن.
  • اهداف درس: بعد از این درس، میفهمی معماری ترنسفورمر چه ایرادهایی داره، بازگشت یا Recursion توی مدل‌ها یعنی چی، و معماری جدید MoR چطور کار میکنه و چه تفاوت‌هایی با بقیه داره.

این روزها وقتشه که به جز معماری معروف ترنسفورمر، گزینه‌های دیگه‌ای هم برای ساخت مدل‌های هوش مصنوعی، مخصوصا مدل‌های زبانی بزرگ (LLM)، روی میز باشه. به تازگی، شرکت گوگل دیپ‌مایند مقاله‌ای منتشر کرده که به نظر میرسه میتونه خیلی انقلابی باشه و اسمش هست Mixture of Recursions یا به اختصار MoR. شاید با شنیدن اسمش یاد Mixture of Experts یا MoE بیفتید، اما این معماری بیشتر به خود ساختار ترنسفورمرها ربط داره.

فصل اول: چرا اصلا به یک معماری جدید نیاز داریم؟ (مشکلات ترنسفورمرها)

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

  1. پردازش یکسان برای همه کلمه‌ها (توکن‌ها): فرقی نداره یه کلمه چقدر ساده یا پیچیده باشه، مدل برای همه‌شون یه اندازه انرژی و پردازش میذاره و اونها رو از تمام لایه‌هاش رد میکنه. این یعنی کلی توان محاسباتی برای کلمه‌های ساده هدر میره و شاید برای کلمه‌های سخت هم به اندازه کافی پردازش انجام نشه.
  2. تعداد پارامترهای خیلی زیاد: ترنسفورمرها برای هر لایه، وزن‌های جداگانه‌ای در نظر میگیرن. این باعث میشه اندازه مدل به شدت بزرگ بشه و آموزش و استفاده از اون به منابع سخت‌افزاری خیلی زیادی نیاز داشته باشه.
  3. بهینه نبودن در زمان اجرا (Inference): وقتی مدل داره یک جواب رو تولید میکنه، هیچ کلمه‌ای نمیتونه زودتر از پردازش خارج بشه. همه باید تا آخر مسیر رو برن. این باعث محاسبات غیرضروری و تاخیر در پاسخ‌دهی میشه، حتی اگه کار بعضی از کلمه‌ها خیلی زودتر تموم شده باشه.
  4. مشکل حافظه پنهان KV Cache: برای هر کلمه در هر لایه، یک سری زوج داده به اسم «کلید-مقدار» یا Key-Value ذخیره میشه. این داده‌ها به سرعت حافظه کارت گرافیک (GPU) رو پر میکنن و باعث میشن نشه متن‌های خیلی طولانی رو به مدل داد.
  5. نبود قابلیت استدلال پنهان: هیچ راهی وجود نداره که مدل بتونه روی درک خودش از یک کلمه چند بار کار کنه و اون رو بهتر کنه. ترنسفورمرها فقط یک بار از اول تا آخر مسیر رو میرن و حلقه استدلال داخلی ندارن.
  6. بودجه محاسباتی ثابت: هیچ انعطافی در این مدل‌ها وجود نداره که بشه توان پردازشی رو در زمان آموزش یا اجرا، بسته به سختی ورودی یا اهمیت کلمه، کم و زیاد کرد.
  7. عدم تطبیق در زمان تست: یک ترنسفورمر استاندارد نمیتونه در زمان اجرا، بدون آموزش مجدد یا دستکاری‌های ساختاری، عمق یا میزان محاسباتش رو به صورت داینامیک تغییر بده.

معماری Mixture of Recursions اومده تا این مشکلات رو یکی یکی حل کنه. اما قبل از اینکه بریم سراغش، باید یه مفهوم پایه رو با هم یاد بگیریم.

فصل دوم: قبل از MoR، بیاید با مفهوم «ترنسفورمر بازگشتی» آشنا بشیم

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

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

اینجوری بهش فکر کن:

  • ترنسفورمر معمولی:
    لایه ۱ ← لایه ۲ ← لایه ۳ ← … ← لایه ۲۴ (هر کدوم متفاوت)
  • ترنسفورمر بازگشتی:
    بلوک A ← بلوک A ← بلوک A ← … (همون بلوک A، به تعداد N بار تکرار میشه)

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

چرا کسی باید این کار رو بکنه؟

  • هزینه حافظه و محاسبات با زیاد شدن عمق مدل به سرعت بالا میره.
  • خیلی از لایه‌ها در ترنسفورمرها کارهای مشابهی انجام میدن.
  • استفاده مجدد از لایه‌ها (یعنی بازگشت) باعث میشه مدل کوچیک‌تر، ارزون‌تر و بهینه‌تر بشه، البته اگه بتونی کاری کنی که درست کار کنه.

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

فصل سوم: خب، حالا Mixture of Recursions یا MoR چیه؟

Mixture-of-Recursions یا MoR یک نوع ترنسفورمر بازگشتیه. اما به جای اینکه حلقه پردازش رو برای همه کلمه‌ها (توکن‌ها) به یک تعداد ثابت اجرا کنه، برای هر توکن بسته به اینکه چقدر به «فکر کردن» نیاز داره، یک عمق بازگشت متفاوت اختصاص میده. یک شبکه عصبی کوچیک به اسم «مسیریاب» یا Router این تصمیم رو به صورت توکن به توکن، هم در زمان آموزش و هم در زمان اجرا، میگیره.

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

همه چیز خودش رو با این ساختار تطبیق میده: لایه‌ها، مکانیزم توجه (attention) و حتی حافظه پنهان کلید-مقدار (KV Cache). حافظه KV فقط برای توکن‌هایی ذخیره میشه که هنوز فعال هستن، نه برای همه توکن‌ها.

یه مثال ساده

فرض کن داری یه قلعه بزرگ با لگو میسازی. بعضی از قسمت‌هاش خیلی ساده‌ان و سریع به هم وصل میشن. اما بعضی قسمت‌هاش خیلی قلق دارن و باید چند بار سعی و خطا کنی تا درست از آب دربیان.

معماری Mixture-of-Recursions مثل یه دستیار هوشمنده که به هر تیکه لگو نگاه میکنه و میگه: «خب، این یکی ساده است، یه بار انجامش بدیم کافیه.» یا «این یکی سخته، بیا چند بار دیگه روش کار کنیم.»

پس به جای اینکه با همه تیکه‌ها یکسان رفتار کنه، برای قسمت‌های سخت وقت بیشتری میذاره و برای قسمت‌های آسون وقت کمتر. اینجوری کارش رو سریع‌تر و هوشمندانه‌تر تموم میکنه.

فصل چهارم: چرا این معماری مهمه؟

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

MoR این مشکل رو به این شکل حل میکنه:

  • لایه‌ها رو به اشتراک میذاره (استفاده مجدد از پارامترها) و در نتیجه وزن‌های منحصر به فرد کمتری داره.
  • توکن‌ها رو به صورت متفاوت مسیردهی میکنه (محاسبات تطبیقی) و در نتیجه فلاپس (FLOPs) کمتری هدر میره.
  • فقط چیزهایی که لازمه رو کش میکنه (کش انتخابی KV) و در نتیجه حافظه کمتری مصرف میکنه.

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

فصل پنجم: عملکردش در عمل چطوره؟ (بنچمارک‌ها)

بیاید ببینیم این معماری در عمل چه نتایجی داشته:

  • یک مدل MoR با ۱۱۸ میلیون پارامتر تونسته در دقت few-shot یک ترنسفورمر معمولی با ۳۱۵ میلیون پارامتر رو شکست بده.
  • مدل MoR با همون بودجه محاسباتی ترنسفورمر (۱۶.۵e۱۸ فلاپس) آموزش داده شده، اما ۲۵ درصد حافظه کمتری مصرف کرده.
  • سرعت اجرا (Inference) بسته به عمق بازگشتی که انتخاب میشه، میتونه تا ۲.۰۶ برابر سریع‌تر بشه.

فصل ششم: بیاید عمیق‌تر بشیم: MoR چطور کار میکنه؟

اینکه MoR کار میکنه دلایل مختلفی داره:

  1. اشتراک‌گذاری وزن (Weight Sharing): مدل MoR از یک پشته لایه به صورت تکراری استفاده میکنه، شبیه به مدل‌های Universal Transformer. اما تفاوتش اینجاست که برای هر توکن تصمیم میگیره چند بار این لایه‌ها رو تکرار کنه.
  2. مسیریابی پویا (Dynamic Routing): یک مسیریاب یا Router که خودش آموزش میبینه، برای هر توکن مشخص میکنه که عمق بازگشتش چقدر باشه. بر خلاف حلقه‌هایی با عمق ثابت یا خروج زودهنگام که به صورت دستی کدنویسی میشن، این قابلیت در حین آموزش یاد گرفته میشه.
  3. بهینه‌سازی حافظه پنهان KV Cache: فقط زوج‌های کلید-مقدار (KV) توکن‌هایی ذخیره میشن که هنوز در حلقه پردازش هستن. اگه یه توکن بعد از ۲ مرحله بازگشت از حلقه خارج بشه، دیگه حافظه‌ای براش اشغال نمیشه. این کار به شدت در مصرف حافظه صرفه‌جویی میکنه. یه گزینه دیگه هم وجود داره که در اون، حافظه پنهان مرحله اول در تمام مراحل بعدی بازاستفاده میشه. این کار حافظه رو حتی بیشتر هم کاهش میده، اما کمی روی دقت تاثیر منفی میذاره.

مسیریابی یا Routing چطوری انجام میشه؟

مدل MoR از یه مسیریاب عصبی کوچیک استفاده میکنه تا تصمیم بگیره: آیا این توکن باید یه دور دیگه پردازش بشه یا از حلقه خارج بشه؟ برای این کار دو روش وجود داره:

  • انتخاب توسط متخصص (Expert-choice): توی این روش، هر مرحله از بازگشت مثل یه نگهبان عمل میکنه. به توکن‌ها نگاه میکنه و چندتایی رو که فکر میکنه لیاقت پردازش بیشتر دارن، انتخاب میکنه. بقیه از حلقه خارج میشن. این روش پویاست و در حین آموزش یاد گرفته میشه. اما یه ریسکی داره و اونم اینه که ممکنه اطلاعات آینده رو ببینه (که بهش میگن نشت علیت). برای همین، یک تابع زیان کمکی بهش اضافه میکنن تا جلوی این مشکل رو بگیرن.
  • انتخاب توسط توکن (Token-choice): توی این روش، هر توکن یک بار در همان ابتدای کار عمق بازگشت خودش رو انتخاب میکنه و تمام. دیگه در هر مرحله تصمیمی گرفته نمیشه. این روش ساده‌تر و قابل پیش‌بینی‌تره، اما عملکردش کمی پایین‌تره چون نمیتونه به خوبی خودش رو تطبیق بده.

پس بله، MoR میتونه هوشمندانه حلقه بزنه. اما اینکه چقدر هوشمندانه این کار رو میکنه بستگی به این داره که کی کار مسیریابی رو انجام میده.

مشکل حافظه پنهان KV Cache

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

MoR این مشکل رو از دو راه حل میکنه:

  • کش کردن بر اساس بازگشت (Recursion-wise caching): فقط زوج‌های KV توکن‌هایی ذخیره میشن که هنوز «توی بازی» هستن. اگه یه توکن بعد از ۲ مرحله بازگشت خارج بشه، ما هم بعد از اون مرحله دیگه براش چیزی کش نمیکنیم. این کار کلی در حافظه صرفه‌جویی میکنه.
  • اشتراک‌گذاری بازگشتی KV (Recursive KV sharing): همه چیز یک بار در مرحله اول بازگشت کش میشه. بعد در مراحل بعدی از همون کش استفاده میشه. این روش ارزون‌تره، اما دقتش کمتره. با این حال، اگه با کمبود حافظه کارت گرافیک مواجه باشید، این روش مثل طلا میمونه.

در هر صورت، MoR در مورد حافظه پنهان خیلی بی‌رحمانه عمل میکنه و اجازه نمیده کسی بی‌دلیل حافظه اشغال کنه.

فصل هفتم: تفاوت MoR با MoE چیه؟

این دو تا رو با هم اشتباه نگیرید. خیلی با هم فرق دارن.

  • MoE (Mixture of Experts) متخصص‌های مختلفی رو انتخاب میکنه. یعنی عرض مدل رو زیاد میکنه.
  • MoR (Mixture of Recursions) عمق بازگشت رو انتخاب میکنه. یعنی عمق مدل رو زیاد میکنه.

MoE مثل انتخاب کردن از منوی غذاست. MoR مثل اینه که یه غذا رو برای مواد اولیه سخت‌ترش، مدت زمان بیشتری بپزی. MoE قطعات بیشتری اضافه میکنه. MoR از همون قطعه موجود، به شکل هوشمندانه‌تری استفاده میکنه.

فصل هشتم: نقاط ضعف MoR

البته نباید وانمود کنیم که این معماری بی‌نقص و عالیه. چندتا مشکل هم داره:

  • مسیریابی به روش «انتخاب توسط توکن» بیش از حد ساده و خشکه. کار میکنه، اما مثل اینه که یه تایمر رو تنظیم کنی و بری. هیچ انعطافی در حین اجرا نداره.
  • اشتراک‌گذاری حافظه KV در روش «انتخاب توسط متخصص» به عملکرد آسیب میزنه. درسته که در حافظه صرفه‌جویی میکنید، اما در ازای اون کمی از دقت رو از دست میدید.
  • ظرفیت مسیریاب بعد از آموزش قفل میشه. یعنی وقتی مسیریاب یاد گرفت که در هر مرحله چند تا توکن رو انتخاب کنه، دیگه به سختی میشه بعدا اون رو تغییر داد.
  • روی مدل‌های خیلی کوچیک (مثلا ۱۳۵ میلیون پارامتر یا کمتر) عملکرد درخشانی نداره.
  • به مهندسی نیاز داره. نمیتونید این معماری رو به راحتی توی کتابخانه‌هایی مثل HuggingFace پیدا و استفاده کنید.

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

Mixture-of-Recursions فقط یک کار تحقیقاتی آکادمیک نیست، بلکه یک تلاش واقعیه برای فکر کردن دوباره به اینکه مدل‌ها چطور باید زمان و انرژی خودشون رو مصرف کنن. به جای اینکه هر توکن رو با خشونت از ۲۴ لایه رد کنه، انگار که این یک قانون مقدسه، مکث میکنه و میپرسه: «آیا این کلمه اصلا به این همه فکر کردن نیاز داره؟». این یک تغییر نگرشه، نه فقط یک بهینه‌سازی کوچیک.

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

لازم نیست هر مدلی بزرگ‌تر باشه. بعضی‌ها فقط لازمه در مورد نحوه حلقه زدنشون هوشمندتر باشن.

منابع

دیدگاه‌ها

دیدگاهتان را بنویسید

نشانی ایمیل شما منتشر نخواهد شد. بخش‌های موردنیاز علامت‌گذاری شده‌اند *