سازوکار مدل‌های فضای حالت: چگونه مامبا معماری ترنسفورمر را به چالش می‌کشد

به مدت هفت سال، معماری ترنسفورمر موتور بلامنازع هوش مصنوعی بوده است. اما این معماری یک نقص ریاضی مهلک دارد: با افزایش حجم متنی که پردازش می‌کند، هزینه محاسباتی آن به صورت درجه دوم (نمایی) افزایش می‌یابد.

اکنون، کلاس جدیدی از معماری‌ها به نام مدل‌های فضای حالت (SSMs) – که توسط چارچوبی به نام مامبا هدایت می‌شود – این برتری را به چالش می‌کشد. مامبا با پردازش داده‌ها در زمان خطی به جای زمان درجه دوم، راهی برای ساخت مدل‌های هوش مصنوعی ارائه می‌دهد که می‌توانند میلیون‌ها کلمه را بدون فروپاشی تحت بار خود بخوانند.[1][6]

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

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

این مقیاس‌پذیری درجه دوم یک دیوار حافظه عظیم در طول استنتاج ایجاد می‌کند – مرحله‌ای که مدل واقعاً برای کاربر متن تولید می‌کند. برای جلوگیری از محاسبه مجدد کل توالی برای هر کلمه جدیدی که تولید می‌شود، ترنسفورمرها باید یک «حافظه پنهان کلید-مقدار (KV)» از تمام توکن‌های قبلی را در حافظه پردازنده گرافیکی (GPU) ذخیره کنند.

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

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

برخلاف ترنسفورمرها که کل توالی را به یکباره در یک عملیات موازی عظیم مشاهده می‌کنند، SSMها داده‌ها را به صورت متوالی پردازش می‌کنند. آنها یک توکن را می‌خوانند، حالت پنهان داخلی خود را برای بازتاب اطلاعات جدید به‌روزرسانی می‌کنند و سپس به سراغ توکن بعدی می‌روند، بسیار شبیه به نحوه خواندن یک جمله کلمه به کلمه توسط انسان.[3]

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

این تفاوت معماری به آنها اجازه می‌دهد تا به صورت خطی مقیاس‌پذیر باشند: دو برابر کردن طول ورودی تنها محاسبات را دو برابر می‌کند، نه چهار برابر. در تئوری، این امر پردازش پنجره‌های متنی عظیم را ساده می‌سازد و به یک مدل هوش مصنوعی اجازه می‌دهد تا یک کتابخانه کامل اطلاعات را بدون نیاز به یک ابرکامپیوتر برای نگهداری زمینه در حافظه فعال، جذب کند.[1]

با این حال، مدل‌های فضای حالت اولیه یک ضعف حیاتی داشتند که مانع از برکناری ترنسفورمر شد. از آنجایی که دینامیک ریاضی آنها در طول زمان ثابت بود – خاصیتی که در مهندسی به عنوان «ثبات زمانی خطی» شناخته می‌شود – در پردازش داده‌های گسسته و متراکم اطلاعاتی مانند زبان انسان مشکل داشتند.

آنها نمی‌توانستند «استدلال مبتنی بر محتوا» انجام دهند. به عبارت عملی، این بدان معنا بود که آنها نمی‌توانستند به طور انتخابی یک کلمه مهم (مانند نام یک شخصیت در رمان) را به خاطر بسپارند یا یک کلمه نامربوط (مانند یک کلمه پرکننده) را فراموش کنند و با همه داده‌ها با وزن ریاضی یکسان رفتار می‌کردند.[1]

در اواخر سال ۲۰۲۳، محققان آلبرت گو و تری دائو این محدودیت اساسی را با معرفی معماری مامبا حل کردند. مامبا مکانیزم «فضای حالت انتخابی» را معرفی کرد. مامبا با اجازه دادن به پارامترهای داخلی مدل برای تغییر پویا بر اساس توکن ورودی خاصی که در حال خواندن آن بود، توانایی فیلتر کردن اطلاعات در زمان واقعی را به دست آورد.

این مدل بالاخره توانست بر آنچه مهم است تمرکز کند و بقیه را کنار بگذارد، و قابلیت‌های استدلال متراکم و بازیابی ترنسفورمرها را حفظ کند، در حالی که کارایی یک مدل متوالی را نیز دارا بود.[1]

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

آنها از یک عملیات «اسکن» استفاده کردند که حالت‌های متوالی را به صورت موازی مستقیماً در حافظه فوق‌سریع SRAM پردازنده گرافیکی محاسبه می‌کند. این ترفند مهندسی هوشمندانه از چرخه‌های کند خواندن و نوشتن حافظه استاندارد با پهنای باند بالای GPU جلوگیری می‌کند و به مامبا اجازه می‌دهد تا به همان سرعت یک ترنسفورمر آموزش ببیند.[1]

این معماری در اواسط سال ۲۰۲۴ با انتشار مامبا-۲ حتی بیشتر تکامل یافت. در این تکرار، محققان یک ارتباط ریاضی عمیق را کشف کردند: آنها ثابت کردند که یک زیرکلاس خاص از مدل‌های فضای حالت ساختاریافته در واقع معادل نوعی از توجه خطی است.

این مفهوم، که «دوگانگی فضای حالت» (SSD) نامیده می‌شود، نشان داد که دو معماری رقیب – SSMهای بازگشتی و مکانیزم‌های توجه موازی – در زیر سطح به طور بنیادی به هم مرتبط هستند و دو روی یک سکه ریاضی را نشان می‌دهند. این پیشرفت نظری، پایه‌ای برای طراحی مدلی بسیار کارآمدتر و مقیاس‌پذیرتر فراهم کرد.[2]

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

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

دستاوردهای عملکردی که توسط این معماری محقق شده، قابل توجه است. آزمایش‌های تجربی نشان می‌دهد که مامبا در طول استنتاج خودرگرسیو، تا ۵ برابر توان عملیاتی بالاتری نسبت به ترنسفورمرهای مشابه به دست می‌آورد. در حوزه‌هایی که ذاتاً به پنجره‌های متنی عظیم نیاز دارند – مانند مدل‌سازی شکل موج‌های صوتی خام یا تجزیه و تحلیل توالی‌های DNA به طول میلیون‌ها در ژنومیک – مامبا به طور مداوم از مدل‌های پیشرفته قبلی بهتر عمل کرده است.

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

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

در حالی که مامبا-۲ برخی از این ناکارآمدی‌های سخت‌افزاری را برطرف می‌کند، ترنسفورمرها همچنان اثبات‌شده‌ترین و آزموده‌شده‌ترین معماری برای آموزش مدل‌های مرزی چند تریلیون پارامتری هستند که چت‌بات‌های تجاری پیشرو امروزی را تقویت می‌کنند. برای بزرگترین شرکت‌های فناوری، خطر کنار گذاشتن یک معماری اثبات شده به نفع یک پارادایم جدید، همچنان مانع بزرگی برای پذیرش فوری است.[2][5]

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

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