تکنیک های پیشرفته ای که هوش مصنوعی را بدون هزینه اضافی سرعت می بخشد

جدول پیوندها

چکیده و 1 مقدمه

2 مدل فضایی حالت

3 مدل فضایی حالت انتخابی و 3.1 انگیزه: انتخاب به عنوان وسیله فشرده سازی

3.2 بهبود SSM با انتخاب

3.3 اجرای کارآمد SSM های انتخابی

3.4 معماری ساده شده SSM

3.5 ویژگی های مکانیسم های انتخاب

3.6 جزئیات مدل اضافی

4 ارزیابی تجربی و 4.1 وظایف ترکیبی

4.2 مدل سازی زبان

4.3 مدل سازی DNA

4.4 مدل سازی و تولید صدا

4.5 معیار سرعت و حافظه

4.6 مدل ابلیشن

5 بحث

6 نتیجه گیری و مراجع

بحث: مکانیسم انتخاب

ب کارهای مرتبط

C مکانیک SSM های انتخابی

D الگوریتم آگاه از سخت افزار برای SSM های انتخابی

E جزئیات تجربی و نتایج اضافی

3.3 اجرای کارآمد SSM های انتخابی

معماری‌های سازگار با سخت‌افزار مانند کانولوشن‌ها (Krizhevsky, Sutskever, and Hinton 2012) و Transformers (Vaswani et al. 2017) کاربرد گسترده‌ای دارند. در اینجا هدف ما این است که SSM های انتخابی را روی سخت افزار مدرن (GPU) نیز کارآمد کنیم. مکانیسم انتخاب کاملاً طبیعی است، و کارهای قبلی سعی داشتند موارد خاصی از انتخاب را به کار ببرند، مانند اجازه دادن به تغییر Δ در طول زمان در SSMهای مکرر (Gu, Dao, et al. 2020). با این حال، همانطور که قبلاً ذکر شد، یک محدودیت اصلی در استفاده از SSMها، کارایی محاسباتی آنهاست، به همین دلیل است که S4 و همه مشتقات از مدل‌های LTI (غیرانتخابی) استفاده می‌کنند که معمولاً به شکل کانولوشن جهانی است.

3.3.1 انگیزه مدل های قبلی

ما ابتدا این انگیزه را مرور می کنیم و رویکرد خود را برای غلبه بر محدودیت های روش های قبلی مرور می کنیم.

• در سطح بالا، مدل های تکرارشونده مانند SSM ها همیشه تعادل بین بیان و سرعت را متعادل می کنند: همانطور که در بخش 3.1 بحث شد، مدل هایی با ابعاد حالت پنهان بزرگتر باید موثرتر اما کندتر باشند. بنابراین ما می خواهیم بعد حالت پنهان را بدون پرداخت هزینه های سرعت و حافظه به حداکثر برسانیم.

• توجه داشته باشید که حالت بازگشتی نسبت به حالت پیچشی انعطاف پذیرتر است، زیرا حالت دوم (3) از گسترش حالت اول (2) مشتق شده است (Gu, Goel, and Ré 2022; Gu, Johnson, Goel, et al. 2021). با این حال، این امر مستلزم محاسبه و تحقق حالت پنهان ℎ با شکل (B، L، D، N)، بسیار بزرگتر (با ضریب N، بعد حالت SSM) از ورودی x و خروجی y شکل (B، L، D). بنابراین حالت کانولوشن کارآمدتر معرفی شد که می‌توانست محاسبات حالت را دور بزند و یک هسته کانولوشن (3a) فقط (B، L، D) را تحقق بخشد.

• SSMهای LTI قبلی از اشکال دوگانه بازگشتی-پیچیده برای افزایش بعد حالت مؤثر با ضریب Nx (≈ 10-100)، بسیار بزرگتر از RNN های سنتی، بدون جریمه های کارایی استفاده می کنند.

3.3.2 مروری بر اسکن انتخابی: گسترش وضعیت آگاه از سخت افزار

مکانیسم انتخاب برای غلبه بر محدودیت های مدل های LTI طراحی شده است. در همان زمان، بنابراین ما نیاز به بررسی مجدد مشکل محاسباتی SSMها داریم. ما این موضوع را با سه تکنیک کلاسیک بررسی می کنیم: ترکیب هسته، اسکن موازی و محاسبه مجدد. ما دو مشاهدات اصلی را انجام می دهیم:

• محاسبات تکراری ساده از FLOPهای O(BLDN) استفاده می کند در حالی که محاسبات کانولوشنال از FLOPهای O(BLD log(L)) استفاده می کند و اولی ضریب ثابت کمتری دارد. بنابراین برای دنباله های طولانی و بعد حالت نه چندان بزرگ N، حالت بازگشتی در واقع می تواند از FLOP های کمتری استفاده کند.

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

ایده اصلی استفاده از ویژگی‌های شتاب‌دهنده‌های مدرن (GPU) برای تحقق وضعیت ℎ فقط در سطوح کارآمدتر سلسله مراتب حافظه است. به طور خاص، اکثر عملیات (به جز ضرب ماتریس) با پهنای باند حافظه محدود می شوند (دائو، فو، ارمون، و همکاران 2022؛ ایوانوف و همکاران 2021؛ ویلیامز، واترمن، و پترسون 2009). این شامل عملیات اسکن ما نیز می شود، و ما از ترکیب هسته برای کاهش مقدار IO های حافظه استفاده می کنیم که منجر به افزایش سرعت قابل توجهی در مقایسه با یک پیاده سازی استاندارد می شود.

برای جلوگیری از تکرار متوالی، مشاهده می‌کنیم که علیرغم خطی نبودن، می‌توان آن را با یک الگوریتم اسکن موازی کارآمد موازی کرد (بللوچ 1990؛ مارتین و کندی 2018؛ اسمیت، وارینگتون و لیندرمن 2023).

در نهایت، ما باید از ذخیره حالت های میانی که برای انتشار پس زمینه ضروری هستند نیز اجتناب کنیم. ما با دقت روش کلاسیک محاسبه مجدد را برای کاهش نیازهای حافظه به کار می‌بریم: حالت‌های میانی ذخیره نمی‌شوند، اما زمانی که ورودی‌ها از HBM به SRAM بارگیری می‌شوند، در گذر به عقب محاسبه می‌شوند. در نتیجه، لایه اسکن انتخابی ذوب شده همان نیازهای حافظه را دارد که اجرای ترانسفورماتور بهینه شده با FlashAttention.

جزئیات هسته ذوب شده و محاسبه مجدد در پیوست D است. لایه و الگوریتم کامل انتخابی SSM در شکل 1 نشان داده شده است.