چطور توجه کار می‌کند: جذر d_k و هزینه‌ی n×n

توجه با سه ضرب کار می‌کند: ضرب داخلی کلید و پرس، نرم‌افزار روی امتیازها، و یک میانگین وزنی از مقدارها. هر سه در بیست خط پایتون بدون هیچ کتابخانه‌ای پیاده شدند. اندازه‌گیریِ کلیدی: بدون تقسیم بر جذر d_k بیشترین وزن به ۰٫۹۹۹۹۵۵ می‌رسد و آنتروپی به ۰٫۲۷ نات می‌افتد؛ با تقسیم، وزن ۰٫۴۷ و آنتروپی ۱٫۷۸ نات. زمان خواندن داده: ۲۷ سپتامبر ۲۰۲۶.

سه ضرب، سه ماتریس

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

def softmax(row):
    m = max(row)
    e = [math.exp(v - m) for v in row]
    s = sum(e)
    return [v / s for v in e]

def attention(q, k, v, scale=True):
    d_k = len(q[0])
    weights = []
    for qi in q:
        row = [sum(a * b for a, b in zip(qi, kj)) for kj in k]
        if scale:
            row = [x / math.sqrt(d_k) for x in row]   # the only difference
        weights.append(softmax(row))
    out = [[sum(w[i] * k2[i] for i in range(len(k2))) for k2 in zip(*v)]
           for w in weights]
    return weights, out

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

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

چرا جذر و نه عدد دیگر

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

mean of (q.k)^2 = 53.5  with d_k = 64
rms logit  = 7.31 unscaled, 0.91 after dividing by sqrt(d_k)

عدد ۷٫۳۱ یعنی امتیازها به‌طور معمول تا هفت نمره از صفر فاصله دارند. حالا نرم‌افزار را ببینید. اگر دو امتیاز ۷٫۳ و ۰٫۱ باشند، اختلافشان ۷٫۲ است و نمای اختلاف بیش از هزار می‌شود؛ یعنی نرم‌افزار عملاً همه‌چیز را به‌جز یک برنده صفر می‌کند. همان چیزی که در اجرای واقعی دیده شد:

d_k = 64, n = 8
scale=False  max weight 0.999955  mean entropy 0.2718 nats (uniform = 2.0794)
scale=True   max weight 0.473270  mean entropy 1.7842 nats (uniform = 2.0794)

آنتروپی ۰٫۲۷ نات در برابر بیشینه‌ی ۲٫۰۸ یعنی وزن‌ها عملاً تک‌قله‌ای‌اند. یک پرس به‌جای اینکه ترکیبی از هفت کلید را ببیند، یک کلید را انتخاب می‌کند. و بدتر از آن، گرادیان تقریباً صفر می‌شود: اگر نرم‌افزار جایی صفر باشد، مشتق جایی نیست. مدلی که آموزش می‌بیند اما چیزی یاد نمی‌گیرد، دقیقاً همین نشانه را دارد.

تقسیم بر جذر d_k انحراف معیار را به حدود یک می‌رساند، یعنی همان عددی که ۰٫۹۱ اندازه گرفت. آنتروپی به ۱٫۷۸ نات می‌رود، یعنی ۸۶ درصد بیشینه. وزن‌ها دیگر قله نیستند و گرادیان زنده می‌ماند.

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

N = 8, five seeds per row
d_k    max weight   entropy    share of the 2.0794 maximum
  16     0.99785     0.73 nats        35%
  64     1.00000     0.30 nats        14%
 256     1.00000     0.11 nats         5%

الگو همان چیزی است که استدلال واریانس پیش‌بینی می‌کرد: با بزرگ‌تر شدن ابعاد، فاصله از یکنواختی بدتر می‌شود. آنتروپی از ۳۵ درصد بیشینه به ۵ درصد می‌افتد، یعنی وزن‌ها تقریباً به یک اندیس فرو می‌ریزند. و بیشترین وزن در ابعاد ۶۴ و ۲۵۶ گرد شده و برابر یک است، یعنی نرم‌افزار عملاً یکنواختی را دور انداخته است.

اینجا به اصلاحی که خودم لازم دارم اشاره می‌کنم: در بخش قبل، با یک بذر، عدد ۰٫۲۷ نات برای ابعاد ۶۴ گزارش شد. با پنج بذر، همان اندازه ۰٫۳۰ نات می‌شود. عدد تک‌بذری برای نسبت ۸۶ درصدی در حالت مقیاس‌شده دقیق بود، ولی برای این جدولِ روند، میانگین پنج بذر لازم است. هر جا که یک عدد ممکن است به بذر حساس باشد، باید همان‌جا گفته شود.

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

هزینه‌ی درجه‌دو: ماتریس مربعی

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

دو برابر کردن طول زمینه، هزینه را چهار برابر می‌کند. سطر آخر را ببینید: زمینه‌ی ۱۲۸ هزار توکنی برای نگه‌داشتن ماتریس وزن‌ها به‌تنهایی ۶۵ گیگابایت حافظه می‌خواهد، و این فقط برای یک سرِ توجه و یک سر. با هشت سر، این عدد بزرگ‌تر می‌شود، و اگر چند لایه داشته باشیم، ضرب می‌شود. خانواده‌ای مثل روفرمر هر سه ضرب را نگه می‌دارد و فقط جای ورودی را عوض می‌کند، پس این هزینه پابرجا می‌ماند.

• نرم‌افزار بدون مقیاس، در آنتروپی ۰٫۲۷ نات قفل می‌شود: وزن‌ها تک‌قله‌ای و گرادیان نزدیک صفر.
• نرم‌افزار با مقیاس، در آنتروپی ۱٫۷۸ نات از بیشینه‌ی ۲٫۰۸ فاصله می‌گیرد و ترکیب واقعی می‌سازد.
• هر دو از همان داده و همان بذر آمدند؛ تنها فرق، یک تقسیم بود.

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

آنچه در این پست نیامد

یک نکته درباره‌ی خود پیاده‌سازی باقی مانده که به آموزش درست توجه می‌کند. در کد بالا حلقه‌ی داخلی ضرب داخلی را با zip نوشته‌ام، یعنی در هر سطر n ضرب اسکالر انجام می‌شود. برای d_k برابر ۶۴ و n برابر ۸ می‌شود ۴٬۰۹۶ ضرب. این برای فهم درست است و برای اجرا بی‌معنا.

# پایتون خالص: درست، کند
row = [sum(a * b for a, b in zip(qi, kj)) for kj in k]

# همان نتیجه با ضرب ماتریسی، که GPU انجام می‌دهد
logits = q @ k.T                      # (n, d_k) @ (d_k, n) -> (n, n)
logits = logits / math.sqrt(d_k)
weights = softmax(logits, axis=-1)
out = weights @ v                     # (n, n) @ (n, d_v) -> (n, d_v)

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

همین‌جا دلیل آن است که اعداد ۶۵ گیگابایت برای زمینه‌ی ۱۲۸ هزاری جدی گرفته می‌شوند. راه‌حل‌های عملی یا ماتریس را تکه‌تکه حساب می‌کنند یا از تقریب استفاده می‌کنند، ولی هر دو یعنی دیگر همه‌ی n² خانه‌ها را یک‌جا نگه نمی‌دارند.

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

دو چیز را هم اندازه نگرفتم و نباید ادعا کنم: رابطه‌ی دقیق جذر d_k با پیش‌ماس در عمل چقدر به هم وابسته‌اند، و اینکه در کدام d_k دقیقاً فرو می‌ریزد. آنچه اندازه گرفتم یک نقطه است، d_k برابر ۶۴. استدلال واریانس پیش‌بینی می‌دهد مشکل با رشد d_k بدتر می‌شود، ولی عدد دقیق آستانه را اندازه نگرفته‌ام.

عددهای تک‌بذری را جدا کنید از میانگین‌ها. نخستین اجرا با یک بذر ثابت بود و در چند اجرای تکراری همان نتیجه را داد؛ نسبت ۸۶ درصدی هم با همان داده پایدار ماند. ولی جدولِ روند ابعاد را با پنج بذر ساختم، چون در همان‌جا تفاوت دیده شد: ۰٫۲۷ در برابر ۰٫۳۰ نات. هر عددی که ممکن است به بذر حساس باشد باید میانگین چند بذر باشد، وگرنه دارید یک تصادف را گزارش می‌کنید.

منابع

1. مقاله‌ی ترنسفورمر — تعریف توجه و دلیل تقسیم بر جذر d_k
2. جای‌گذاری چرخشی — حذف موقعیت مطلق از کلید و پرس
3. درون‌یابی موقعیت — کوچک کردن اندیس موقعیت به اندازه‌ی زمینه‌ی آموزش
4. روش یارن — افزایش پنجره‌ی زمینه با کمترین آموزش
5. مرجع چرخشی ترنسفورمرز — انواع مقیاس‌پذیری و کلید پیکربندی
6. مدل روفرمر — همان سه ضرب با ورودی چرخشی
7. پرچم‌های خط فرمان لام‌سی‌پی‌پی — تعداد دنباله‌ی هم‌زمان و اندازه‌ی پنجره در اجرای محلی
8. پست موقعیت چرخشی — اندازه‌گیری همین روش در عمل