‏Wall Attention: واریانتی از توجه با فراموشیِ کانال‌به‌کانال

Wall Attention یه واریانتِ مکانیزمِ توجهه که ایدهٔ اصلیش یه زوالِ ضربیِ کانال‌به‌کانال و گام‌به‌گام ـه که مستقیم داخلِ ضربِ داخلیِ Q و K جاسازی شده. جایی که توجهِ استاندارد امتیازِ یه زوج موقعیت رو با مجموعِ سادهٔ ضربِ کانال‌ها می‌سازه، Wall Attention هر کانال رو با یه زوالِ یادگرفته‌شده که بین دو موقعیت انباشته می‌شه وزن می‌ده. نتیجه اینه که هر کانالِ کوئری یه نرخِ فراموشیِ مستقل و وابسته به محتوا پیدا می‌کنه.

به گفتهٔ توسعه‌دهنده، این کار گیتینگِ اسکالر (مثل FoX) و زوال‌های سبکِ RoPE رو به کلِ بُعدِ کانال تعمیم می‌ده. نکتهٔ جالب اینه که اگه زوال (g) رو صفر بذاری، دقیقاً به همون توجهِ softmaxِ معمولی برمی‌گردی؛ پس این مکانیزم یه حالتِ کلی‌تره که توجهِ استاندارد یه نمونهٔ خاصش حساب می‌شه.

این ریپو دوتا کرنل که در عمل استفاده می‌شن رو، هرکدوم جداگانه، بسته‌بندی کرده. یکی برای آموزش/prefill (به اسم wall_attn): یه کرنلِ Triton با forward و backwardِ ادغام‌شده به سبکِ FlashAttention با گرادیان‌های تحلیلی برای q، k، v و g. اون یکی برای decode (به اسم wall_attn_decode): یه کرنلِ تک‌گامه که از یه کشِ KVِ از پیش بازمقیاس‌شده می‌خونه، طوری که تولیدِ هر توکن به‌جای محاسبهٔ دوبارهٔ کلِ پیشوند، فقط یه پاسِ کوچیکِ شبیهِ GEMV هزینه داره. رابطِ کاربریِ آموزش خیلی ساده‌ست:

from wall_attn import wall_attn
# decay g is per-channel; g=0 recovers vanilla softmax attention
o = wall_attn(q, k, v, g, scale=K**-0.5)  # [B, T, HQ, V]

از قابلیت‌های دیگه‌ش می‌شه به پشتیبانی از GQA (که تعدادِ هدهای کوئری از هدهای KV بیشتر باشه)، گیتِ اسکالرِ اختیاری، attention sink، پنجرهٔ لغزان و بسته‌بندیِ varlen اشاره کرد. کشِ decodeش با تاکردنِ زوال داخلِ کلیدها و استفاده از لنگرِ هر-تکه (per-chunk anchor)، از نظر عددی تا زمینهٔ طولانی پایدار می‌مونه. هر مسیرِ کرنل در برابرِ یه مرجعِ PyTorchِ ساده تست شده و گرادیان‌ها با تفاضلِ محدودِ مرکزی راستی‌آزمایی شدن. کرنل‌ها روی ماشینِ توجهِ موازیِ کتابخانهٔ flash-linear-attention ساخته شدن و تحتِ لایسنسِ MIT منتشر شدن.

نکات کلیدی:

• زوالِ ضربیِ کانال‌به‌کانال و گام‌به‌گام، داخلِ ضربِ داخلیِ Q و K
• هر کانالِ کوئری نرخِ فراموشیِ مستقل و وابسته به محتوا داره؛ با g=0 به softmaxِ معمولی برمی‌گرده
• تعمیمِ گیتینگِ اسکالر (FoX) و زوال‌های RoPE به کلِ بُعدِ کانال
• دو کرنلِ Triton: آموزش/prefill و decode با کشِ KVِ از پیش بازمقیاس‌شده
• پشتیبانی از GQA، attention sink، پنجرهٔ لغزان و varlen؛ لایسنس MIT