توجه سریع با حساب تکه‌تکه

یعنی چی

یه انبار بزرگ دور دارید و یه میز کار کوچیک کنار دستتون. اگه برای هر کاری برید انبار و برگردید، روزتون تموم می‌شه.

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

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

سرعت از حافظه‌ست، نه از حساب

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

ضرب دو تا ماتریس بزرگ، به ازای هر عدد، حساب‌های زیادی داره؛ پس سر واحدهای حساب شلوغه و کارت با تمام توان کار می‌کنه. ولی سافت‌مکس و ماسک و دراپ‌اوت، به ازای هر عدد، فقط یکی دو تا حساب ساده دارن؛ پس بیشتر وقتشون منتظر حافظه‌ان. پژوهش ۲۰۲۲ نشون داد توی توجه معمولی، همین قدم‌های سبک بیشترین وقت رو می‌گیرن.

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

سافت‌مکس کل ردیف رو می‌خواد

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

یه مشکل دیگه هم هست. تابع نمایی خیلی زود عددهای خیلی بزرگ می‌سازه؛ e به توان ۹۰ از بزرگ‌ترین عددی که یه عدد اعشاری معمولی نگه می‌داره، بزرگ‌تره. برای همین همیشه پیش از سافت‌مکس، بزرگ‌ترین امتیاز ردیف رو از همه کم می‌کنن؛ جواب عوض نمی‌شه، ولی عددها کوچیک می‌مونن. ولی این هم یعنی باید بیشینه‌ی کل ردیف رو بدونیم، باز هم پیش از دیدن همه‌ی تکه‌ها.

پس به نظر میاد سافت‌مکس مجبورمون می‌کنه کل ردیف رو یه‌جا داشته باشیم. ترفند بعدی نشون می‌ده که نه.

خروجی هم همین‌جوری درست می‌شه

خروجی توجه، جمع وزن‌دار مقدارهاست. وقتی تکه‌ی اول رو دیدیم، یه خروجی موقت می‌سازیم؛ مقدارهای تکه‌ی اول، ضرب در عددهای نمایی همون تکه. هنوز تقسیم بر جمع نکردیم، چون جمع نهایی رو نمی‌دونیم.

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

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

کل دستور، قدم به قدم

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

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

ماسک، ارزون‌تر از همیشه

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

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

چند بار کمتر می‌ره انبار

پژوهش ۲۰۲۲ فقط نگفت «کمتر»؛ حساب کرد چقدر. فرض کنید طول متن n باشه، اندازه‌ی بردار هر سر d، و اندازه‌ی میز M عدد. توجه معمولی، چون جدول n در n رو می‌نویسه و می‌خونه، به اندازه‌ی n ضرب در n عدد با انبار رد و بدل می‌کنه. توجه سریع حدود n ضرب در n ضرب در d ضرب در d، تقسیم بر M.

این فرمول ترسناک به نظر میاد، ولی معنیش ساده‌ست. d معمولا ۶۴ یا ۱۲۸‌ه، پس d ضرب در d حدود چند هزاره، و M، یعنی تعداد عددهایی که روی میز جا می‌شن، ده‌ها هزار تا صدها هزار. پس کسر d ضرب در d تقسیم بر M خیلی کمتر از یکه، و توجه سریع چند برابر کمتر از توجه معمولی با انبار حرف می‌زنه. هر چی میز بزرگ‌تر، فاصله بیشتر. پژوهشگرها حتی نشون دادن که با یه میز ثابت، هیچ روش دقیقی نمی‌تونه به‌طور کلی از این خیلی کمتر بره انبار؛ یعنی این ترتیب، تقریبا بهترین ترتیب ممکنه.

تکه‌هایی که اصلا حساب نمی‌شن

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

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

دقیق، نه تقریبی

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

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

موقع یادگیری: دوباره حساب کن

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

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

یه برنامه به‌جای چند تا

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

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

چقدر سریع‌تر

پژوهش ۲۰۲۲ گزارش داد که برای آموزش یه مدل زبانی کوچیک با متن ۱۰۲۴ تکه‌ای، کل آموزش حدود سه برابر سریع‌تر شد، و برای یه مدل رمزگذار رایج، حدود ۱۵ درصد سریع‌تر از سریع‌ترین رکورد اون زمان. مهم‌تر از سرعت، این بود که با همون کارت، مدل‌ها تونستن متن‌های خیلی بلندتری ببینن، و کیفیتشون، چون متن بیشتری می‌دیدن، بهتر شد.

سال ۲۰۲۳، نسخه‌ی دوم اومد که کار رو بهتر بین بخش‌های موازی کارت پخش کرد و کارهای غیرضربی رو کمتر کرد، و حدود دو برابر نسخه‌ی اول سریع‌تر بود. سال ۲۰۲۴ هم نسخه‌ی سوم برای کارت‌های تازه‌تر اومد، که حساب و جابه‌جایی داده رو هم‌زمان انجام می‌ده و از عددهای کم‌دقت‌تر هم استفاده می‌کنه. فکر اصلی، یعنی تکه‌تکه و سافت‌مکس جاری، توی همه‌شون یکیه.

کجا و چطوری به کار می‌ره

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

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

از کجا اومد

سال ۲۰۱۸ دو تا پژوهشگر نشون دادن سافت‌مکس رو می‌شه یه‌دور، با یه بیشینه و یه جمع جاری، حساب کرد؛ اون موقع هدفشون فقط سریع‌تر کردن خود سافت‌مکس بود. سال ۲۰۲۱ یه پژوهش دیگه با عنوانی که می‌گفت خودتوجهی به حافظه‌ی مربعی نیاز نداره، نشون داد با همین ترفند، توجه رو می‌شه با حافظه‌ی خیلی کم حساب کرد.

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

اشتباه‌های رایج

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

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

تمرین

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

جواب‌ها. اولی، بیشینه ۳‌ه، پس یک به‌اضافه‌ی e به توان منفی ۲، یعنی حدود ۱٫۱۴. دومی، e به توان (۳ منهای ۵)، یعنی حدود ۰٫۱۴. سومی، چون جدول رو دوباره می‌سازه ولی گلوگاه رفت‌وآمد حافظه‌ست، و اون رو خیلی کم کرده. چهارمی، نه برابر، چون مربعیه.