ملخص تقني: LionMuon

بيان المشكلة

يُعد تدريب النماذج اللغوية الكبيرة (LLMs) مشكلة تحسين ضخمة، حيث يمثل التكلفة الحسابية لكل خطوة (FLOPs) وسرعة التقارب المحددين الأساسيين لفاتورة الحوسبة النهائية. تواجه المحسنات الحالية مقايضة بين جودة التحديث والكفاءة الحسابية:

• الطرق القائمة على الإشارة (مثل Lion وSignum): تُنتج هذه الطرق تحديثات بسيطة تعتمد على إشارة الإحداثيات ($W_{t+1} = W_t - \eta \cdot \text{sign}(\hat{G}_t)$). وهي غير مكلفة حسابياً وفعالة من حيث الذاكرة، لكنها غالباً ما توفر اتجاهات تحديث أضعف مقارنة بالطرق الطيفية.
• الطرق الطيفية (مثل Muon): تقوم هذه الطرق بحساب تحديث إشارة المصفوفة عبر تكرارات Newton-Schulz ($W_{t+1} = W_t - \eta \cdot \text{msign}(\hat{G}_t)$). وبينما توفر اتجاهات طيفية أقوى تؤدي غالباً إلى تقارب أسرع، إلا أنها تتكبد تكلفة أعلى بكثير لكل خطوة بسبب عمليات ضرب المصفوفات المتعددة وزيادة عبء الاتصال في البيئات الموزعة (التي تتطلب عملية all-gather للمصفوفة كاملة).

السؤال المركزي الذي يتم معالجته هو: هل من الممكن الجمع بين التكلفة المنخفضة للتحديثات القائمة على الإشارة والجودة العالية للاتجاهات الطيفية لتحقيق مقايضة متفوقة بين التكلفة والجودة؟

المنهجية: LionMuon

يقترح المؤلفون LionMuon، وهو مُحسن يتناوب بين تحديثات الإشارة بأسلوب Lion والتحديثات الطيفية بأسلوب Muon خلال فترة ثابتة $P$.

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

1. الجدول الزمني المتناوب: تقوم الخوارزمية بخطوة Muon (تحديث طيفي Newton-Schulz) في كل خطوة رقم $P$، وتحديثات إشارة بأسلوب Lion في الخطوات الـ $P-1$ المتبقية.
2. الزخم المشترك: يشترك كلا نوعي التحديث في مخزن زخم واحد من نوع dual-Exponential Moving Average (EMA) وهو ($M_t$). يتم حساب الاتجاه $\hat{G}_t$ عبر استكمال Lion: $\hat{G}_t = \beta_1 M_{t-1} + (1-\beta_1)G_t$.
3. كفاءة الحالة: يحافظ LionMuon على مخزن زخم واحد فقط، مما يجعله يطابق البصمة الذاكرية لـ Lion/Signum، وهو بالضبط نصف مساحة AdamW.
4. الحالات الخاصة:
   • $P=1$: Muon خالص.
   • $P=\infty$: Lion خالص (أو Signum إذا كان $\beta_1=\beta_2$).
   • $P \in (1, \infty)$: الطريقة المتناوبة المقترحة.
   • SignMuon: نسخة متغيرة تستخدم زخمًا أحادي الـ EMA ($\beta_1 = \beta_2$) مع نفس الجدول الزمني المتناوب.

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

يقدم المؤلفون حدود تقارب لـ LionMuon تحت ظروف الضجيج ذو الذيول الثقيلة (heavy-tailed noise)، وهي حالة شائعة في تدريب النماذج اللغوية الكبيرة.

• الثوابت المستكملة: يعتمد حد التقارب على متوسط السلاسة للفترة ($\bar{L}$) ومستويات الضجيج ($\bar{\rho}$) التي تستكمل بين ثوابت Muon الصرف (النمط الطيفي) و Lion الصرف (نمط اللانهاية).
• الفترة المثلى: تشير النظرية إلى وجود فترة مثلى $P^*$ تقلل من عامل المقايضة الذي يتضمن نسب ثوابت السلاسة ($L_\infty / L_2$) ومستويات الضجيج ($\rho_1 / \rho_{nuc}$).
• التعقيد: تُظهر حدود التعقيد المستمدة أن LionMuon يمكنه تحقيق نفس دقة Muon أو Lion بعدد أقل من العمليات في تكوينات محددة، لا سيالما عندما يكون عامل المقايضة $\phi(P) < 1$.

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

1. SignMuon: مُحسن متناوب يؤدي خطوة Muon كل $P$ دورة، وخطوة إشارة بأسلوب Signum في الدورات الأخرى. يوضح المؤلفون أن هذا المتغير أحادي الـ EMA يحسن أيضاً من مقايضة الخسارة مقابل FLOPs مقارنة بـ Muorn الصرف عند $P=2$.
2. LionMuon: المساهمة الرئيسية، حيث يستبدل Signum بقاعدة dual-EMA الخاصة بـ Lion. وهو يحافظ على توفير التكلفة لـ SignMuon مع تحسين التقارب بشكل أكبر. يتطلب فقط معلمة فائقة عددية إضافية واحدة ($P$) وهو بديل سهل الاستخدام لـ Lion.
3. الحدود النظرية: حدود تعقيد حادة تحت ضجيج الذيول الثقيلة والتي تنمذج صراحةً الاستكمال بين أنظمة Muon و Lion، مما يتنبأ بالفترة $P$ المثلى للحوسبة.
4. التحقق التجريبي: تجارب واسعة النطاق عبر مجموعات بيانات متعددة (FineWeb, SlimPajama, WikiText-103) وبنيات مختلفة (GPT-2, LLaMA) بمقاييس 124 مليون، 355 مليون، و720 مليون معلمة.

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

• مقياس 124 مليون: يتفوق LionMuon (تحديداً عند $P=2$) على Muon و Lion و Signum و AdamW في جميع مجموعات البيانات والبنيات المختبرة. فهو يحقق خسارة تحقق (validation loss) أقل مع حوالي 5.5% أقل من إجمالي عمليات تدريب FLOPs مقارنة بـ Muon الصرف.
• التوسع (355 مليون و720 مليون): تستمر ميزة جودة التقارب عند المقاييس الأكبر. عند 355 مليون، يتفوق LionMuon و SignMuon ($P=2$) على Muon الصرف. وعند 720 مليون، لا يزال SignMuon ($P=2$) يتفوق على Muon الصرف، رغم أن ميزة dual-EMA لـ LionMuon على SignMuon لم تتكرر في هذه الميزانية المحددة من التدريب غير المكتمل.
• التدريب الموزع: توفر الطريقة وفورات هيكلية في الاتصال. بما أن خطوات Lion هي عنصرية (element-wise) ومحلية للشرائح، فإن ضبط $P=2$ يقلل بمقدار النصف من تكرار عمليات all-gather المكلفة المطلوبة لخطوة Newton-Schulz في Muon، بغض النظر عن حجم الدفعة (batch size).

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

يزعم البحث أن LionMuon يقدم نهجاً متفوقاً هيكلياً لتدريب النماذج اللغوية الكبيرة من خلال فصل تكرار التحديثات الطيفية المكلفة عن عدد الدورات.

• الكفاءة: يحافظ على فعالية خطوات Muon مع تقليل متوسط تكلفة الدورة، مما يطابق كفاءة الذاكرة للطرق القائمة على الإشارة.
• المتانة: الطريقة متينة عبر أحجام نماذج ومجموعات بيانات مختلفة، حيث تحدد باستمرار جبهة باريتو (Pareto frontier) للخسارة مقابل FLopS.
• المواءمة النظرية: تتوافق النتائج التجريبية مع التوقعات النظرية فيما يتعلق بالفترة المثلى $P$، مما يشير إلى أن القيم المتوسطة (مثل $P=2$ أو $P=5$) غالباً ما توفر أفضل توازن بين الاتجاه الطيفي المكلف واتجاه الإشارة الرخيص.

يخلص المؤلفون إلى أن LionMuon هو امتداد عملي لتدريب النماذج اللغوية الكبيرة الحديثة، حيث يحسن بشكل صارم من الحالة الراهنة (Muon) من حيث جودة التقارب والكفاءة الحسابية على حد سواء. ويُقترح في العمل المستقبلي اختيار $P$ بشكل تكيفي والتوسع لما بعد 720 مليون معلمة.