ملخص تقني: MONA – مُحسِّن Muon مع تسارع نيستروف

1. بيان المشكلة

تتطلب عملية تدريب النماذج اللغوية الكبيرة (LLMs) بشكل متزايد مُحسّنات ذات كفاءة عينة فائقة. وبينما ظل AdamW هو المعيار السائد لأكثر من عقد من الزمن، ظهرت بدائل حديثة مثل Muon. يعمل Muon على تحسين الطرق من الدرجة الأولى عبر التعامل مع مصفوفات الأوزان كوحدات هندسية واحدة وتطبيق التعامد المصفوفي (عبر تكرار نيوتن-شولز) على مخزن الزخم (momentum buffer). يتماشى هذا النهج مع الانحدار الشديد تحت المعيار الطيفي، مما يوفر كفاءة حسابية تبلغ حوالي $2\times$ مقارنة بـ AdamW في تدريب نماذج خليط الخبراء (MoE) واسعة النطاق.

ومع ذلك، يشترك Muon في قيد جوهري مع جميع طرق الدرجة الأولى: فهو يفتقر إلى آليات صريحة لاستكشاف مشهد الخسارة (loss landscape) بما يتجاوز معلومات التدرج المحلية. ونتيجة لذلك، يمكن لـ Muon أن يعلق في النقاط الدنيا الحادة (sharp local minima)، لا سيما في سيناريوهات التدريب ذات الدفعات الكبيرة (large-batch training). ترتبط النقاط الدنيا الحادة بضعف التعميم، بينما تميل النقاط الدنيا المسطحة (flatter minima) إلى تقديم أداء أفضل في المهام اللاحقة. وبينما قدم المُحسِّن ALTO حداً للتسارع يعتمد على المتوسط المتحرك الأسي (EMA) لفروق التدرجات للهروب من النقاط الدنيا الحادة، فإن تصميمه كمحول عام لـ Adam و LAMB يفرض تحديات في التكامل مع عملية التعامد غير الخطية للغاية في Muon.

2. المنهجية: MONA

يقترح المؤلفون MONA (مُحسِّن Muon مع تسارع نيستروف)، والذي يدمج بسلاسة حداً للتسارع مدركاً للانحناء (curvature-aware) مباشرة في إطار عمل التعامد الخاص بـ Muon.

الخوارزمية الأساسية

يعدل MONA مسار Muon القياسي عن طريق إدخال حد التسارع قبل تراكم الزخم والتعامد. تسير عملية التحديث لمعلم المصفوفة $W$ عند التكرار $k$ كما يلي:

1. حساب فرق التدرج: حساب الفرق بين التدرج العشوائي الحالي $G_k$ والتدرج السابق $G_{k-1}$:
   $$D_k = G_k - G_{k-1}$$
2. تحديث مخزن التسارع: الحفاظ على متوسط متحرك أسي لهذه الفروقات:
   $$A_k = \beta_a A_{k-1} + (1 - \beta_a) D_k$$
   هنا، يتحكم $\beta_a$ في ذاكرة التسارع.
3. التدرج المدرك للانحناء: إنشاء تدرج مُعزز $\tilde{G}_k$ عن طريق إضافة حد التسارع المقيّس إلى التدرج الخام:
   $$\tilde{G}_k = G_k + \alpha A_k$$
   يتحكم المعامل $\alpha$ في مقياس التسارع. تشير الدوافع النظرية إلى أن $G_k - G_{k-1}$ يقرب $H_k(\theta_k - \theta_{k-1})$، حيث $H_k$ هي مصفوفة هسيان (Hessian). وبالتالي، فإن هذا الحد يشير بعيداً عن النقاط الدنيا الحادة (حيث تكون تغيرات التدرج كبيرة) ويظل صغيراً في المناطق المسطحة.
4. تعامد Muon: يتم تغذية التدرج المعزز في خطوات زخم وتعامد Muon القياسية:
   $$M_k = \mu M_{k-1} + \tilde{G}_k$$
   $$O_k = \text{Newton-Schulz}(M_k)$$
   $$W_{k+1} = W_k - \eta (\gamma O_k + \lambda W_k)$$

المنطق التصميمي

يعد تطبيق حد التسارع قبل عملية التعامد أمراً بالغ الأهمية. فإذا طُبق بعد ذلك، فإنه سيخل بالبنية المتعامدة التي يعتمد عليها Muon لتنظيم المعيار الطيفي. ومن خلال تحويل التدرج الخام إلى اتجاه مدرك للانحناء أولاً، يلتقط مخزن الزخم هذه الاتجاهات الغنية، والتي تتم معالجتها بعد ذلك عبر مسار Muon القياسي.

بالنسبة للمعلمات ذات القيم المتجهة (مثل التضمينات أو الانحيازات)، يعود MONA افتراضياً إلى AdamW، ملتزماً ببروتوكولات Muon الحالية.

3. التحليل النظري

يقدم البحث تحليلاً للتقارب لـ MONA تحت افتراضات قياسية (النعومة L، التدرجات غير المنحازة مع تباين محدود، والمحاذاة الاتجاهية المتوقعة).

• معدل التقارب: يثبت المؤلفون أن MONA يحقق معدل تقارب قدره $O(K^{-1/2})$ في الإعدادات غير المحدبة، وهو ما يطابق المعدلات التقاربية لـ SGD و AdamW. يؤثر حد التسارع على الثوابت في الحد العلوي ولكن ليس على المعدل التقاربي.
• الهروب من النقاط الدنيا الحادة: يوضح التحليل كيف يقاوم حد التسارع النقاط الدنيا الحادة بشكل انتقائي. بالقرب من نقطة ثابتة، يكون التسارع المتوقع $E[A_k]$ متناسباً مع مجموع الهسيان للاتجاهات الماضية. بالنسبة للاتجاهات ذات القيم الذاتية الكبيرة (النقاط الدنيا الحادة)، يكون مقدار التسارع كبيراً، مما يدفع المُحسِّن بعيداً. أما بالنسبة للقيم الذاتية الصغيرة (الاتجاهات المسطحة)، فيكون التسارع صغيراً، مما يسمح بتقارب مستقر. هذه الآلية تحافظ على تنظيم المعيار الطيفي لـ Muon مع إضافة الوعي بالانحناء.

4. النتائج التجريبية

تم تقييم MONA عبر ثلاثة مقاييس لنماذج تدريب MoE (بـ 1 مليار، 6 مليارات، و68 مليار معلمة) ومرحلة الضبط الدقيق الخاضع للإشراف (SFT).

أداء ما قبل التدريب (Pretraining)

• النماذج: MOE-1B-A0d2B (400 مليار توكن)، MOE-6B-A0d5B (1.2 تريليون توكن)، و MOE-68B-A3B (700 مليار توكن).
• المقاييس: تمت مراقبة خسارة التحقق (validation loss) عبر مجالات البرمجة، الاستدلال الرياضي، اللغة الإنجليزية العامة، والنصوص الأكاديمية الصينية.
• النتائج: حقق MONA باستمرار خسارة تحقق أقل من كل من Muon و AdamW عبر جميع مقاييس النماذج الثلاثة وجميع المجالات. كانت فجوة الأداء واضحة بشكل خاص في النموذج الأكبر 68B، حيث أظهر MONA استقراراً أكبر في التقارب.

اختبارات المهام اللاحقة (MOE-68B-A3B)

تم تقييم نموذج 68B المدرب باستخدام MONA عبر مجموعة شاملة من الاختبارات:

• القدرة العامة: حقق MONA أعلى متوسط درجات (0.4557) في اختبارات تشمل MMLU، MMLU-Pro، CMMLU، CEVAL، و BBH، متفوقاً على Muon (0.4478) و AdamW (0.4382).
• الاستداد الرياضي وتوليد الكود: أظهر MONA أداءً فائقاً في GSM8K، GPQA، ومختلف اختبارات توليد الكود (HumanEval+، MBPP+، BigCodeBench، LiveCodeBench). والجدير بالذكر أن MONA قدم مكاسب كبيرة في المهام التي تتطلب تمثيلات برمجية قابلة للنقل.
• الضبط الدقيق الخاضع للإشراف (SFT): عندما خضع النموذج المدرب بـ MONA لعملية SFT على بيانات برمجية، حافظ على خسارة تدريب أقل طوال العملية مقارنة بأساس Muon. وأكدت تقييمات ما بعد SFT على مجموعة BigCode أن النماذج التي بدأت بـ MONA حققت درجات أعلى في 6 من أصل 7 مهام مقارنة بالنماذج التي بدأت بـ Muon.

الكفاءة والنشر

لمعالجة العبء الإضافي للذاكرة الناتج عن تخزين التدرج السابق ($G_{k-1}$) ومخزن التسارع ($A_k$)، قدم المؤلفون MONA-Lite:

• الكمية (Quantization): تخزين الحالات المساعدة بتنسيق bfloat16 (BF16) بدلاً من float32.
• الحوسبة المتدفقة (Streaming Computation): حساب فروق التدرج في مكانها (in-place) لإلغاء الحاجة إلى مخزن مخصص لـ $G_{k-1}$.
• النتيجة: قللت هذه الاستراتيجيات من عبء الذاكرة الإضافي بنسبة 75% تقريباً دون التضحية بجودة التدريب. وأظهر تحليل العبء الحسابي أن MONA أبطأ بنسبة ~1% فقط عند مستوى خطوة المُحسّن، مع عدم وجود فرق عملي في وقت التكرار النهائي (end-to-end iteration time) مقارنة بـ Muon.

5. الأهمية والادعاءات

يزعم البحث أن MONA نجح في سد الفجوة بين الكفاءة الهندسية لـ Muon والحاجة إلى التسارع المدرك للانحناء. تكمن أهميته الأساسية في:

1. التكامل الخوارزمي: إثبات أنه يمكن دمج حدود التسارع في أطر عمل التعامد المصفوفي دون الإخلال بخصائصها الهندسية.
2. القابلية للتوسع: التحقق من المنهج عبر نطاق واسع من أحجام النماذج (من 1 مليار إلى 68 مليار معلمة) وعدد التوكنات (يصل إلى تريليون توكن)، مما يظهر تحسينات ثابتة على كل من AdamW و Muon.
3. التعميم: تقديم أدلة تجريبية على أن الهروب من النقاط الدنيا الحادة أثناء ما قبل التدريب يؤدي إلى تعميم أفضل في المهام اللاحقة، وخاصة في الاستدلال المعقد وتوليد الكود.
4. العملية: تقديم نسخة فعالة من حيث الذاكرة (MONA-Lite) تجعل المُحسّن قابلاً للاستخدام في البيئات محدودة الموارد.

يخلص المؤلفون إلى أن MONA يمثل حالة متقدمة (state-of-the-art) للمُحسّنات المستخدمة في تدريب النماذج اللغوية الكبيرة واسعة النطاق، حيث يحقق تقارباً فائقاً وأداءً لاحقاً متميزاً مع الحفاظ على الكفاءة الحسابية لإطار عمل Muon.