مکانیسم خوداتنشن چندسر (Multi-Head Self-Attention) موتور محرک اصلی مدلهای زبانی بزرگ امروزی است. بسیاری از آموزش‌ها این بلوک را صرفاً در حد فرمول‌های ریاضی توضیح می‌دهند یا کد پیاده‌سازی آن را با سه لایهی خطی جداگانه می‌نویسند که بازدهی پایینی روی سخت‌افزار دارد. درک شیوه‌ی بهینه‌ی چیدمان ابعاد تنسورها و ادغام تبدیل‌های پرسش، کلید و مقدار (QKV) در یک لایه، مرز میان یک مدل کند و یک معماری استاندارد است.
در این درس یاد می‌گیرید که چگونه اتنشن چندسر را با یک لایهی تصویر خطی یکپارچه، ماسک علّی (Causal Mask) و مقیاسگذاری درست بسازید. این پیاده‌سازی سبک و ساخت‌یافته نیازی به GPUهای ابری سنگین ندارد و به‌راحتی روی CPU سیستم شخصی یا شتاب‌دهنده‌های محلی اجرا میشود و پایهی ساخت ترنسفورمر کامل را شکل می‌دهد.
یک لایهی تصویر کردن خطی، سه دیدگاه، H سرِ موازی، یک ماسک. بلوک اتنشن (attention) دقیقاً همان‌طور که مدل در عمل از آن استفاده میکند.
اهداف یادگیری
- پیاده‌سازی دسته‌ایِ (batched) تصویر کردنِ پرسش/کلید/مقدار (Query/Key/Value) در قالب یک لایهی خطیِ یکپارچه که به H سر تقسیم میشود.
- محاسبه‌ی اتنشن ضرب داخلی مقیاس‌شده (scaled dot-product attention) با نرمال‌سازی صحیح و مدیریت درست انواع داده (dtype).
- اعمال یک ماسک علّی (causal mask) که مانع از توجه هر موقعیت به موقعیت‌های بعدی میشود.
- بررسی وزن‌های اتنشن به‌ازای هر سر برای یک ورودی ثابت و تحلیل اینکه هر سر به چه چیزی نگاه میکند.
- آموزش یک بلوک اتنشن کوچک روی یک وظیفه‌ی ساختگی (toy task) و مشاهده‌ی افت تابع زیان (loss) هم‌زمان با تخصصی‌شدن عملکرد سرها.
قاب کلی
اتنشن تابعی است که به بازنمایی یک توکن (token) اجازه می‌دهد اطلاعات را از سایر توکنهای همان توالی بیرون بکشد. خوداتنشن (self-attention) یعنی پرسش‌ها، کلیدها و مقدارها همگی از یک ورودی یکسان به دست می‌آیند. چندسر بودن (multi-head) هم به این معناست که تصویر کردن به H مسئله‌ی اتنشن موازی تقسیم میشود؛ سپس خروجی آنها به هم الحاق (concatenate) شده و دوباره تصویر می‌شوند.
الگوی پیاده‌سازی بهینه، استفاده از یک لایهی خطی است که ابعاد را از D به 3 * D تصویر میکند، سپس به سه دیدگاه (view) برش می‌خورد و در نهایت به H سر، هر کدام با اندازه‌ی D // H، تغییر شکل می‌یابد. ضرب ماتریسی، سافتمکس (softmax) و جمع وزن‌دار همگی به‌صورت عملیات‌های تنسوری دسته‌ای انجام می‌شوند تا سرها روی شتابدهنده به‌طور موازی اجرا شوند.
در این درس همین بلوک را می‌سازیم. همچنین ماسک علّی را اضافه می‌کنیم تا همین کد بتواند به‌عنوان لایهی اتنشن در یک مدل زبانی فقط-دیکودر (decoder-only) کار کند. در درس بعد، این بلوک را برای ساخت یک ترنسفورمر کامل روی هم می‌چینیم و در درس پس از آن، مدل را آموزش خواهیم داد.
قرارداد ابعاد
ابعاد ورودی (B, T, D) است. ابعاد خروجی نیز (B, T, D) خواهد بود. ابعاد ماسک (T, T) است یا می‌توان آن را به این ابعاد پخش (broadcast) کرد. داخل بلوک، تنسورهای میانی ابعاد (B, H, T, d_head) دارند که در آن d_head = D // H است. قید اصلی این ساختار، برقراری رابطه‌ی D % H == 0 است.
دو لایهی خطی (تصویر کردن QKV و تصویر کردن خروجی) تنها پارامترهای این بلوک هستند. ماسک، سافتمکس، ضرب‌های ماتریسی و تغییر شکل‌ها همگی بدون پارامتر (parameter-free) هستند.
تفکیک QKV
پیاده‌سازی ابتدایی سه لایهی خطی مجزا دارد؛ یکی برای Q، یکی برای K و یکی برای V. در مقابل، نسخه‌ی بهینه فقط یک لایه دارد که 3 * D ویژگی در خروجی تولید میکند و سپس نتیجه را برش می‌زند. این دو روش از نظر ریاضی کاملاً معادل‌اند؛ زیرا سه ضرب ماتریسی مجزا در وزن‌های (D, D)، دقیقاً برابر است با یک ضرب ماتریسی در یک وزن (3D, D) که از روی‌هم چیده شدن همان ماتریس‌ها تشکیل شده است.
نسخه‌ی بهینه سریع‌تر است چون شتابدهنده بهجای سه بار، فقط یک بار ضرب ماتریسی را اجرا میکند. مقداردهی اولیهی آن هم ساده‌تر است، زیرا هر سه زیرماتریس داخل یک تنسور پارامتر مشترک قرار دارند و می‌توان آنها را با هم مقداردهی اولیه کرد.
تغییر شکل سرها
پس از تفکیک، ابعاد هر یک از Q ،K و V برابر با (B, T, D) است. برای تبدیل آنها به H مسئله‌ی اتنشن موازی، ابعاد را به (B, T, H, d_head) تغییر می‌دهیم و سپس به (B, H, T, d_head) ترانهاده (transpose) می‌کنیم. با این کار، بُعد سر کنار بُعد دسته (batch) قرار می‌گیرد و PyTorch با اتنشنِ هر سر مثل یک عملیات دسته‌ای روی B * H نمونه‌ی مستقل برخورد میکند.
بُعد d_head در جایگاه آخر باقی می‌ماند تا ضرب ماتریسی امتیازها یعنی Q @ K.transpose(-2, -1) آن را منقبض و ادغام (contract) کند. خروجی این عملیات، امتیازهای اتنشن به‌ازای هر سر با ابعاد (B, H, T, T) است.
مقیاسگذاری
امتیازها پیش از ورود به سافتمکس بر sqrt(d_head) تقسیم می‌شوند. بدون این مقیاسگذاری، با بزرگتر شدن d_head، حاصل ضرب‌های داخلی رشد میکند و سافتمکس را به سمتی push میکند که در آن تقریباً تمام وزن روی یک درایه متمرکز شده و بقیه به صفر میل می‌کنند. گرادیانها در چنین وضعیتی فوق‌العاده ناچیز می‌شوند و یادگیری متوقف میشود. تقسیم بر sqrt(d_head) واریانس امتیازها را در اندازه‌های مختلف سر تقریباً ثابت نگه می‌دارد.
ماسک علّی
یک مدل زبانی فقط-دیکودر هنگام پیشبینی توکن بعدی، صرفاً می‌تواند بر گذشته شرطی‌سازی شود. ماسک این محدودیت را تحمیل میکند. به‌طور مشخص، پیش از سافتمکس، تمام درایههای بالای قطر اصلی در ماتریس امتیازِ (T, T) با منفی بینهایت جایگزین می‌شوند. در نتیجه پس از سافتمکس، وزن این موقعیت‌ها صفر خواهد شد.
ماسک را هنگام ساخت لایه به‌عنوان یک بافر ثبت می‌کنیم تا روی همان دستگاه (device) مدل قرار بگیرد و بخشی از گراف گرادیان نباشد. این ماسک بیشترین طول بافتاری را که بلوک با آن مواجه خواهد شد پوشش می‌دهد. در زمان گذر رو به جلو (forward pass)، گوشه‌ی بالا-چپ آن را با ابعاد (T, T) برش می‌زنیم.
تصویر کردن خروجی
پس از به دست آمدن بردارهای بافتار (context vectors) به‌ازای هر سر با ابعاد (B, H, T, d_head)، آنها را دوباره به (B, T, H, d_head) ترانهاده می‌کنیم، به (B, T, D) تغییر شکل می‌دهیم و یک تصویر کردن خطی نهایی با ابعاد (D, D) روی آنها اعمال می‌کنیم. لایهی تصویر کردن خروجی به مدل اجازه می‌دهد سرها را با هم ترکیب کند. بدون آن، سرهای Hگانه فقط از طریق لایه‌های بعدی شانس ترکیب مجدد پیدا می‌کردند و بلوک دچار محدودیت مصنوعی می‌شد.
بررسی وزن‌های اتنشن
در این درس فلگ return_weights=True روی گذر رو به جلو تعبیه شده است. وقتی فعال باشد، بلوک علاوه بر خروجی، وزن‌های اتنشن هر سر را با ابعاد (B, H, T, T) برمی‌گرداند. بخش دمو نقشه‌ی حرارتیِ (heatmap) وزن‌های یکی از سرها را روی یک ورودی کوتاه نمایش می‌دهد تا ساختار مثلثی علّی و تمرکز مدل روی هر موقعیت را ببینید.
در یک مدل آموزش‌دیده، سرهای مختلف الگوهای متفاوتی را یاد می‌گیرند. بعضی از سرها به توکنِ دقیقاً قبلی نگاه می‌کنند، برخی به ابتدای توالی توجه دارند و دسته‌ای دیگر اتنشن را تقریباً یکنواخت پخش می‌کنند. این هوکِ بررسی، نقطه‌ی ورود به کارهای تفسیرپذیری (interpretability) است.
دموی آموزش
دموی انتهای فایل main.py، بلوک اتنشن را به یک سرِ مدل زبانی (LM head) کوچک متصل کرده و کل ساختار را روی یک وظیفه‌ی تکرار آموزش می‌دهد. هر سطر از ورودی، یک شناسه‌ی تصادفی است که در طول بافتار تکرار شده است. هدف، همان ورودی است که یک گام شیفت خورده؛ بنابراین مدل باید یاد بگیرد که توکن بعدی همان توکن قبلی است. تابع زیان، آنتروپی متقاطع (cross-entropy) است. با H=4 ،D=32 ،T=12 و اندازه‌ی واژگان 64، تابع زیان طی سه دور (epoch) روی CPU از مقدار تصادفی (حدود log(64) ~ 4.16) به مقداری بسیار کمتر از 1.0 می‌رسد.
هدف دمو ساختن یک مدل کاربردی نیست؛ بلکه می‌خواهیم مطمئن شویم گرادیانها در تمام بخش‌های بلوک جریان پیدا می‌کنند و سرها روی مسئله‌ای با جواب کاملاً واضح، یاد می‌گیرند.
مواردی که در این درس بررسی نمی‌شوند
این درس شامل بلوک پیش‌خور (feed-forward) نمی‌شود. لایهی ترنسفورمر در یک مدل واقعی شامل اتنشن و به دنبال آن یک MLP دولایه است که هر کدام اتصالات باقیمانده (residual connection) و نرمال‌سازی لایه را در اطراف خود دارند. در درس بعدی این بخش‌ها را اضافه می‌کنیم.
کدگذاری مکانی روتاری (rotary) یا AliBi پیاده‌سازی نمی‌شود. هر دو روش در مرحله‌ی تصویر کردن QKV در همین بلوک اعمال می‌شوند، اما مبحث آموزشی مستقلی هستند. بلوکی که در اینجا می‌سازیم با هر دو سازگار است؛ فقط کافی است پیش از ضرب ماتریسی، تبدیل لازم را روی Q و K پیاده کنید.
کش کلید-مقدار (KV cache) برای استنتاج پیاده‌سازی نمی‌شود. کش کردن کلیدها و مقدارها در میان گذرهای رو به جلو، همان بهینه‌سازی است که رمزگشایی (Decode) خودبازگشتی (autoregressive) را سرعت می‌بخشد. این کار قرارداد ابعاد تنسورهای K و V را تغییر می‌دهد اما اثری روی Q ندارد؛ جایگاه این مبحث در درس استنتاج است.
راهنمای خواندن کدها
فایل main.py کلاس MultiHeadSelfAttention را تعریف میکند. این کلاس دو لایهی خطی و یک بافر ماسکِ ثبت‌شده دارد. گذر رو به جلو عملیات‌های تصویر کردن، تغییر شکل، امتیازدهی، ماسک‌گذاری، سافتمکس، وزن‌دهی، تغییر شکل مجدد و تصویر کردن دوباره را انجام می‌دهد. دموی انتهای فایل، مدل کوچکی می‌سازد که اتنشن را در کنار امبدینگ‌های توکن و مکانی و یک LM head در بر می‌گیرد، آن را برای سه دور روی وظیفه‌ی کپی آموزش می‌دهد و در نهایت منحنی زیان و نقشه‌ی حرارتی اتنشن هر سر را چاپ میکند. تست‌های موجود در code/tests/test_attention.py قرارداد ابعاد، خاصیت علّی، ویژگی‌های سافتمکس، تفکیک درست سرها و جریان گرادیان را تثبیت و بررسی می‌کنند.
ابتدا دمو را اجرا کنید. سپس مقدار n_heads را از 4 به 8 برسانید (با ثابت نگه داشتن d_model=32 تا d_head=4 شود) و تغییرات نقشه‌ی حرارتی را ببینید.