Exclusive Self-Attention Beyond AdamW: Stability and Generalization under Muon Optimization

ملخص تقني: الانتباه الذاتي الحصري لما وراء AdamW

بيان المشكلة
تسمح آليات الانتباه الذاتي القياسية في نماذج Transformer لمخرجات الرمز (token) عند الموضع $i$ بأن تتضمن مزيجاً موزوناً من ناقلات القيم (value vectors)، بما في ذلك ناقل القيمة الخاص بالرمز نفسه ($v_i$). ورغم أن هذا التصميم يتمتع بقدرة تعبيرية عالية، إلا أنه قد يخلق "اختصاراً تمثيلياً" (representational shortcut)، حيث يظل مخرج الانتباه متوافقاً بقوة مع $v_i$، حتى وإن كانت المسارات المتبقية (residual pathways) تحفظ بالفعل المعلومات المحلية للرمز. هذه الظاهرة، التي تُسمى "تحيز تشابه الانتباه" (attention similarity-bias)، عالجها Zhai [2] من خلال "الانتباه الذاتي الحصري" (Exclusive Self-Attention - XSA)، والذي يقوم صراحةً بإزالة المكون الموازي لناقل قيمة الرمز نفسه من مخرج الانتباه.

ومع ذلك، فإن الدراسة الأصلية لـ XSA قيمت الطريقة حصرياً ضمن بيئة تدريب NanoGPT باستخدام مُحسن AdamW. وقد حدد مؤلفو العمل الأصلي صراحةً أن توافق XSA مع المُحسنات البديلة، مثل Muon، يظل سؤالاً مفتوحاً. ويمثل هذا فجوة حرجة لأن أبحاث النماذج اللغوية الصغيرة الحديثة (مثل stack الخاص بـ nanochat) تستخدم بشكل متزايد استراتيجيات تحسين هجينة (Muon لمعاملات المصفوفات، وAdamW للتضمينات/embeddings) وميزات معمارية مثل تضمينات القيم (value embeddings) وتشفير الموضع الدوراني (rotary position encoding). ولا يزال من غير الواضح ما إذا كانت القيود الهندسية لـ XSA تظل مستقرة أو مفيدة عند نقلها إلى أنظمة التدريب الحديثة القائمة على Muon.

المنهجية
تجري الدراسة تجربة إعادة إنتاج ونقل مضبوطة باستخدام بيئة تدريب nanochat. تتضمن المنهجية ما يلي:

1. التنفيذ: تم دمج نسخة مُعلمة (parameterized) من XSA في nanochat. تُعرف العملية كـ $z_i = y_i - \alpha \cdot \text{proj}_{v_i}(y_i)$، حيث $y_i$ هو مخرج الانتباه القياسي، و$v_i$ هو ناقل قيمة الرمز، و$\alpha \in [0, 1]$ يتحكم في قوة الإسقاط. عند $\alpha=1$، يتم إزالة المكون المتوافق ذاتياً بالكامل.
2. التحقق من الصحة: تم التحقق من التنفيذ عددياً لضمان أنه عند $\alpha=0$، يكون مسار XSA مكافئاً للانتباه القياسي. وتتحقق عمليات تدريب قصيرة من تطابق سلوك الخسارة المبكر بين النموذج الأساسي ومسار $\alpha=0$.
3. الإعداد التجريبي: تم تقييم عمقين للنموذج تحت ظروف تدريب متطابقة (نفس ترتيب البيانات، عائلة المُحسن، جدول معدل التعلم، وطول السياق):
   • d12: 12 طبقة، بُعد خفي 768، 12 رأساً.
   • d24: 24 طبقة، بُعد خفي 1024، 16 رأساً.
   • المُحسن: يتم تحسين معاملات المصفوفات باستخدام Muon؛ بينما تستخدم المعاملات المتبقية AdamW.
   • البيانات: ClimbMix-400B.
4. المقاييس: تتبع الدراسة خسارة التدريب، والبت لكل بايت (bpb) في التحقق، ودرجات التقييم الإجمالية اللاحقة (Base CORE وChatCORE)، وتشخيصات الانتباه (التوافق الجيبي بين المخرج والقيمة الذاتية، والسلوك القطري).

المساهمات الرئيسية

• إعادة الإنتاج في بيئة Muon: أول دمج لـ XSA في بيئة تدريب قائمة على Muon، مما يعالج مباشرة السؤال المفتوح بشأن توافق المُحسن الذي تركه البحث الأصلي.
• تنفيذ مُعلم (Parameterized): تنفيذ مرن يسمح بالانتباه القياسي، وXSA الكامل، والإسقاط الجزئي عبر معامل القوة $\alpha$، جنباً إلى جنب مع ضوابط اختيار الطبقات.
• التحقق والتدقيق: التحقق من حد $\alpha=0$ وتصحيح عيوب التسجيل (اصطلاحات مؤشر الخطوة والتنعيم) التي شوهت سابقاً المقارنات المبكرة للتدريب.
• تقييم بعمقين: تحليل مقارن عند أعماق d12 وd24 باستخدام مجموعة شاملة من المقاييس.
• التحليل الآلي والنظمي: تشخيصات التوافق الذاتي للانتباه وتقييم العبء الزمني (runtime overhead).

النتائج
تشير النتائج التجريبية إلى أن XSA مستقر تحت بيئة Muon ولكنه يعطي فوائد متباينة في المهام اللاحقة اعتماداً على عمق النموذج والمقياس المحدد:

• استقرار التحسين: تظل مسارات خسارة التدريب لنماذج XSA والنماذج الأساسية متقاربة بشدة عند كلا العمقين. لا يوجد دليل على التباعد، أو ظهور قيم NaN، أو عدم الاستقرار، مما يؤكد أن XXSA لا يخل بهندسة تحديثات Muon.
• أداء التحقق (البت لكل بايت - bpb): يحسن XSA باستمرار من bpb في التحقق (held-out validation):
  • d12: انخفض من 0.8484 إلى 0.8457 (تحسن نسبي قدره 0.32%).
  • d24: انخفض من 0.7193 إلى 0.7171 (تحسن نسبي قدره 0.31%).
• المقاييس اللاحقة (CORE):
  • d24: يحسن XSA كلاً من Base CORE (من 0.2593 إلى 0.2606، أي +0.50%) وChatCORE (من 0.3638 إلى 0.3682، أي +1.21%).
  • d12: بينما يتحسن bpb، ينخفض Base CORE من 0.1602 إلى 0.1508 (أي -5.87%).
• التشخيصات الآلية: تظهر تشخيصات الانتباه انخفاضاً في التوافق الذاتي وانخفاضاً في السلوك القطري تحت تأثير XSA، وهو ما يتوافق مع الآلية المقصودة لتعزيز المزج السياقي.
• العبء الزمني: وُصف العبء الزمني بأنه متواضع. وبينما أشارت تقديرات أولية إلى ~5%، إلا أن قياسات الإنتاجية المتكررة وضعت الفرق قريباً من التباين المنهجي للعمليات القصيرة، مما يشرح عدم وجود عقبة كبيرة.

الأهمية والادعاءات
تخلص الورقة إلى أن "الانتباه الذاتي الحصري" هو انحياز بنيوي متوافق مع بيئات التدريب القائمة على Muon، ولكن فائدته ليست عالمية. تكمن الأهمية الرئيسية في إثبات أن القيد التمثيلي لـ XSA ينتقل بسلاسة إلى نظام تحسين حديث دون زعزعة استقرار التدريب.

ومع ذلك، يؤكد المؤلف أن المكاسب اللاحقة تعتمد على المقياس والمقياس المستخدم. إن التحسن في bpb في التحقق عند كلا العمقين يشير إلى فائدة تمثيلية، لكن تدهور Base CORE عند d12 يشير إلى أن فوائد إزالة التوافق الذاتي قد تعتمد على عمق النموذج، أو قدرة مسارات المتبقي/القيمة على حفظ المعلومات المحلية، أو مهمة التقييم المحددة. لا تدعي الدراسة أن XSA هو بديل متفوق عالمياً للانتباه القياسي، بل هو قيد مستقر ومنخفض الاضطراب يقدم فوائد في سياقات محددة (لاسيما عند d24 في هذه الدراسة).

يعمل هذا العمل كفحص للمتانة (robustness check) بدلاً من كونه إعادة إنتاج مطابقة تماماً، مما يسلط الضال على التفاعلات المعقدة بين التعديلات المعمارية، وهندسة المُحسن، ومقاييس التقييم في أبحاث النماذج اللغوية الحديثة.