ملخص تقني: Muown - التحكم في معيار الصف لتعظيم كفاءة Muon

بيان المشكلة

برزت Muon كمحسن تنافسي لتدريب النماذج اللغوية الكبيرة (LLM)، حيث تتفوق غالباً على AdamW من خلال الاستفادة من البنية المصفوفية للطبقات الخطية في Transformer. ومع ذلك، تشير الملاحظات الأخيرة إلى أن سلوك Muon عند النطاقات الكبيرة حساس لتقليص الأوزان (weight decay) وعرضة لعدم الاستقرار العددي، مثل تجاوز سعة bfloat16 وانفجار لوجيتات الانتباه (attention logits).

تتمثل إحدى المشكلات الحرجة التي تم تحديدها في الانجراف التصاعدي المستمر للمعيار الطيفي ($\|W_t\|_{S_\infty}$) لمصفوفات الأوزان أثناء التدريب عند استخدام Muon دون تقليص أوزان منفصل. هذا الانجراف ينتهك الشرط الطيفي المطلوب لتعلم الميزات المستقر ($\|W\|_{S_\infty} = \Theta(\sqrt{m/n})$). وبينما يمكن لتقليص الأوزان المنفصل أن يخفف من ذلك عن طريق تقليص مصفوفة الأوزان بشكل موحد، إلا أنه تدخل غير مستهدف يعمل على تقليص الطيف بأكمله بغض النظر عن الاتجاهات المحددة التي تدفع النمو، مما قد يؤدي إلى إبطاء التقارب في مراحل التدريب المبكرة. تسعى الورقة البحثية إلى تحديد المحرك التجريبي المحدد لانجراف المعيار الطيفي واقتراح آلية تحكم أكثر استهدافاً.

المنهجية: Muown

يقدم المؤلفون Muown، وهو بديل جاهز للاستخدام (drop-in replacement) لـ Muon يعالج انجراف المعيار الطيفي من خلال استراتيجية إعادة تمثيل (reparameterization) بدلاً من القيود الإجمالية.

1. تفكيك المعيار الطيفي

الرؤية الجوهرية هي تفكيك المعيار الطيفي إلى عاملين (Proposition 1):
$$\|W_t\|_{S_\infty}^2 = \|g_t\|_\infty^2 \cdot \lambda_{\max}(P_t C_t P_t)$$
حيث أن:

• $g_t = \|W_t\|_{row}$ هو متجه مقادير الصفوف (المعايير الإقليدية للصفوف).
• $D_t$ هي المصفوفة المعيرة صفياً ($W_t = \text{Diag}(g_t)D_t$).
• $C_t = D_t D_t^\top$ هي مصفوفة غرام (Gram matrix) للصفوف ذات المعيار الموحد.
• $P_t$ هي مصفوفة قطرية لمقادير الصفوف المعيرة.
• $\lambda_{\max}(P_t C_t P_t)$ يمثل عامل تماسك الصفوف (row-coherence).

يكشف التحليل التجريبي على نماذج بحجم 500 مليون معلمة أن عامل مقدار الصف ($\|g_t\|_\infty$) هو المحرك الرئيسي لانجراف المعيار الطيفي، حيث يتحرك بالتزامن مع $\|W_t\|_{S_\infty}$. في المقابل، يظل عامل تماسك الصفوف مستقراً ومحدوداً طوال فترة التدريب.

2. إعادة التمثيل والتحسين

بناءً على هذا التشخيص، تعيد Muown تمثيل مصفوفة الأوزان $W$ كالتالي:
$$W(g, R) = \text{Diag}\left(\frac{g}{\|R\|_{row}}\right) R$$
هنا، $g \in \mathbb{R}^m$ يمثل متجه مقدار الصف الصريح، و $R \in \mathbb{R}^{m \times n}$ يمثل مكون الاتجاه.

• مكون الاتجاه ($R$): يتم تطبيق Muon دون تغيير على $R$، حيث يتم التعامل معه تحت هندسة المعيار الطيفي (أقصى انحدار طيفي).
• مكون المقدار ($g$): يتم ترقية مقادير الصفوف $g$ لتصبح متغيرات محسنة صريحة. يتم تحديثها تحت هندسة $\ell_\infty$، وهي المعيار المزدوج الطبيعي لمعيار العمليات للمصفوفات القطرية.
• قاعدة التحديث: يستخدم التحديث لـ $g$ محسن Adam (يعمل كـ SignSGD مُنعم تحت هندسة $\ell_\infty$) للتعامل مع الضجيج العشوائي، بينما يتبع التحديث لـ $R$ إجراء Muon القياسي باستخدام تنظيم Newton-Schulz المتعامد.

يتم تنفيذ هذا النهج ضمنياً داخل حالة المحسن (optimizer state)، مما يترك عملية المرور الأمامي (forward pass) لبنية النموذج دون تغيير.

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

1. تشخيص الانجراف الطيفي: تقوم الورقة بتفكيك المعيار الطيفي لتحديد أن أقصى مقدار للصف ($\|g_t\|_\infty$) هو المحرك التجريبي لنمو المعيار الطيفي تحت Muon، مما يميزه عن عامل تماسك الصفوف المستقر.
2. خوارزمية Muown: محسن جديد يعامل مقادير الصفوف كمتغيرات قابلة للتدريب صراحةً يتم تحديثها تحت هندسة $\ell_\infty$، مع الحفاظ على تحديثات الاتجاه الخاصة بـ Muon. هذا يتجنب التقليص العشوائي لتقليص الأوزان ومن قبلويات التعامد (orthogonality priors) لإعادة التمثيل الطيفي.
3. ضمانات التقارب: يثبت المؤلفون أن Muown يحقق معدلات تقارب مثالية غير خطية في كل من النظم الحتمية والعشوائية. يتم إجراء التحليل في هندسة مزدوجة مختلطة ($\ell_1$ للمقادير، و Schatten-1 للاتجاهات) بما يتماشى مع هندسات التحديث. يشير الحد النظري إلى تقليل معامل الضجيج العشوائي الفعال مقارنة بـ Muon.
4. التحقق التجريبي: تُظهر التجارب المكثفة على FineWeb-Edu عبر أحجام نماذج تتراوح من 124 مليون إلى 2.7 مليار معلمة تحسينات متسقة.

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

• تحسينات الارتياب (Perplexity): تحقق Muown باستمرار ارتياباً أقل من Muon وSOAP وAdamW وLion عبر جميع أحجام النماذج المختبرة (124M إلى 2.7B) ومعدلات التعلم.
• المتانة تجاه المعلمات الفائقة (Hyperparameters):
  • معدل التعلم: توسع Muown هضبة معدلات التعلم القريبة من المثالية عبر مختلف عروض النماذج، مما يقلل من الحساسية لضبط معدل التعلم لكل مقياس.
  • تقليص الأوزان (Weight Decay): تخفف Muown من الحاجة إلى ضبط تقليص الأوزان؛ فهي تعمل جيداً بدون تقليص أوزان ($\lambda=0$)، بينما تتطلب Muon ضبطاً دقيقاً (الأمثل $\lambda=0.1$) لمنع التباعد أو الأداء دون المستوى.
• الاستقرار: تمنع Muown الانجراف المنهجي في المعيار الطيفي الملاحظ في Muon. في أنظمة معدل التعلم العالي (مثل $4 \times 10^{-3}$)، غالباً ما تتباعد Muon، بينما تظل Muown مستقرة وقريبة من المثالية.
• خصائص التدرج: يظهر التحليل أن Muown تحافظ على رتبة فعالة معيرة أعلى لتدرجات الدفعة المصغرة (minibatch gradients) وتُظهر معامل ضجيج عشوائي أقل طوال التدريب مقارنة بـ Muon.
• التكلفة الحسابية الإضافية: لا يضيف الأسلوب عبئاً زمنياً يُذكر. عندما يتم تقسيم المحسن عبر الرتب (في بيئة موزعة)، تصبح أوقات الخطوة لـ Muown وMuown متطابقة تقريباً. التكلفة في الذاكرة ضئيلة (خطية في عدد الصفوف)، حيث تضيف بضعة ميجابايت فقط لكل طبقة.

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

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

• تحسين التقارب: ارتياب أفضل واستقرار دون الحاجة إلى تقليص أوزان عدواني.
• الأساس النظري: معدلات تقارب مثالية في إطار عمل مدرك للهندسة يأخذ في الاعتبار صراحةً الطبيعة المختلطة للتحديثات.
• الكفاءة العملية: بديل جاهز للاستخدام يطابق وقت تشغيل Muon في البيئات الموزعة مع تقديم متانة فائقة تجاه خيارات المعلمات الفائقة.

يضع المؤلفون Muown كحل يسد الفجوة بين متطلبات الاستقرار النظري لتعلم الميزات والاحتياجات العملية لتدريب النماذج اللغوية الكبيرة واسعة النطاق، مما يقدم بديلاً أكثر استهدافاً لتقليص الأوزان المنفصل.