ملخص تقني: التسطيح الطيفي هو كل ما يحتاجه Muon

بيان المشكلة

تتضمن تدريب الشبكات العصبية الحديثة عملية تحسين فوق تضاريس غير محدبة للغاية تتميز بمناطق مسطحة، واتجاهات حادة، ونقاط سرجية. ويعد اختيار المحسن (Optimizer) أمراً حاسماً للاستقرار والكفاءة. وبينما حسنت الطرق التكيفية المعتمدة على الإحداثيات مثل Adam وAdamW المتانة، إلا أنها لا تعالج بشكل كامل التحديات الهندسية التي تفرضها المعلمات ذات القيم المصفوفية. مؤخراً، برز Muon كبديل عالي الأداء يقوم بتعميد (Orthogonalization) مخزن الزخم قبل التحديثات، مما يسمح بمعدلات تعلم أكبر وتقارب أسرع. ومع ذلك، فإن الآليات النظرية التي تفسر لما reason يتفوق Muon على التدرج الاشتقاقي القياسي (SGD) من حيث أقصى أحجام الخطوات المستقرة ومعدلات التقارب لا تزال غير مستكشفة بشكل كافٍ. وتحديداً، لم يفسر النظرية الحالية كمياً كيف يغير التعميد (Orthogonalization) عنق الزجاجة الخاص بمعدل التعلم أو الرقم الشرطي الفعال لمشكلة التحسين.

المنهجية

تقدم الورقة تحليلاً نظرياً لـ Muon عبر تأطير قاعدة التحديث الخاصة به كشكل من أشكال التسطيح الطيفي (Spectral Flattening). طور المؤلفون نتيجتين نظريتين أساسيتين وتحققا منهما تجريبياً:

1. الإطار النظري:

   • تعريف Muon: يقوم Muon بتحديث الأوزان $W$ عن طريق حساب التدرج $G_t$، ثم تطبيق تكرارات Newton–Schulz لتقريب العامل القطبي $O_t \approx (G_t G_t^\top)^{-1/2} G_t = UV^\top$. هذه العملية تحافظ على المتجهات المفردة للتدرج بينما تستبدل القيم المفردة بواحدات.
   • تقريب هسيان (Hessian Approximation): يستخدم التحليل نموذج انحناءً مف de-factored (K-FAC)، حيث يقرب الهسيان كـ $H \approx X^\top X \otimes G_t G_t^\top$، حيث تمثل $X$ تنشيطات المدخلات و $G_t$ التدرج.
   • تحليل معدل التعلم: استنتج المؤلفون عتبات الهبوط في خطوة واحدة بدقة لكل من SGD وMuon. وقارنوا بين أقصى معدل تعلم مستقر ($\eta_{max}$) لكلا المحسنين، موضحين أن SGD مقيد بأكبر قيمة مفردة للتدرج ($\sigma_{max}$)، بينما يستند استقرار Muon إلى متوسط القيمة المفردة ($\frac{1}{m}\sum \sigma_i$).
   • تحليل التقارب: تمت إعادة صياغة Muon كطريقة تدرج مسبقة الشرط (Preconditioned Gradient Method) مع مسبق شرط $P = I_n \otimes (G_t G_t^\top)^{-1/2}$. وتحت شروط Polyak–Łojasiewicz (PL) والنعومة النسبية، حلل المؤلفون عامل التقارب. لقد أثبتوا أن نسبة الشرط الفعالة لـ Muon ($\tilde{\alpha}/\tilde{\beta}$) هي أفضل تماماً من نسبة SGD ($\alpha/\beta$) عندما يكون تباين التدرج غير متماثل (Anisotropic)، وتحديداً من خلال التحسن بعامل مرتبط بالانتشار الطيفي للتدرج.

2. التصميم التجريبي:

   • الإعداد: أُجريت التجارب على CIFAR-10 باستخدام CifarNet. ولعزل تأثير المحسن على المعلمات ذات القيم المصفوفية، تم استخدام استراتيجية المحسن المزدوج: حيث يتم تحديث النوى الالتفافية بواسطة المحسن المستهدف (Muon أو SGD)، بينما تُحدث الانحيازات (Biases) والمعلمات الأخرى باستخدام محسن مرجعي ثابت.
   • اختبارات الاستقرار: تم إجراء مسحات لمعدل التعلم بدون تطبيع الدفعات (Batch Normalization - BN) لمراقبة الاستقرار الجوهري.
   • اختبارات التقارب: تم تشغيل كلا المحسنين بمعدلات تعلم متطابقة (وبشكل منفصل عند أفضل معدلاتهما الخاصة) لقياس سرعة التقارب ونسب التقارب لكل خطوة ($r_t$).
   • مبدأ التطبيع: تم اقتراح واختبار مخطط تطبيع جديد، FrobNorm (التطبيع بواسطة معيار Frobenius)، للتحقق من الادعاء النظري بأن تقليل المقياس الطيفي للمدخلات ($\lambda_{max}(X^\top X)$) يتيح معدلات تعلم أعلى.

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

1. اشتقاق الحد الأقصى لمعدل التعلم: تثبت الورقة أن أقصى حجم خطوة مستقر لـ Muon يتناسب مع متوسط القيمة المفردة للتدرج بدلاً من أكبر قيمة لها.

   • بالنسبة لـ SGD: $\eta_{max} \propto \frac{2}{\lambda_{max}(H)} \approx \frac{2}{\lambda_{max}(X^\top X) \sigma_{max}(G)}$.
   • بالنسبة لـ Muon: $\eta_{max} \propto \frac{2}{\lambda_{max}(H)} \frac{\sum \sigma_i}{m} \approx \frac{2}{\lambda_{max}(X^\top X) \sigma_{max}(G)} \frac{\sum \sigma_i}{m \sigma_{max}(G)}$.
   • يوضح هذا أن التسطيح الطيفي يمنع اتجاهاً مفردًا مهيمناً واحداً من فرض قيود على حجم الخطوة، مما يسمح لـ Muon بتحمل معدلات تعلم أكبر بكثير.

2. تسريع معدل التقارب: يظهر المؤلفون أن Muon يعمل كمسبق شرط (Preconditioner) يحسن نسبة الشرط الفعالة.

   • تحت نموذج K-FAC، يكون عامل التقارب لـ Muon هو $1 - \frac{\tilde{\alpha}}{\tilde{\beta}}$، حيث $\frac{\tilde{\alpha}}{\tilde{\beta}} = \frac{\alpha}{\beta} \sqrt{\frac{\lambda_{min}(G_t G_t^\top)}{\lambda_{max}(G_t G_t^\top)}}$.
   • بما أن $\frac{\lambda_{min}}{\lambda_{max}} < 1$، فإن الحد $\sqrt{\dots}$ أقل من 1، لكن الاشتقاق يوضح أن $\frac{\tilde{\alpha}}{\tilde{\beta}} > \frac{\alpha}{\beta}$ في سياق النظام مسبق الشرط، مما يؤدي إلى معدل تقارب خطي أسرع بشكل صارم عندما يكون تباين التدرج سيء الشرط (Ill-conditioned).

3. مبدأ تطبيع جديد: يكشف التحليل أن أقصى معدل تعلم مستقر يتناسب عكسياً مع $\lambda_{max}(X^\top X)$. يقترح المؤلفون أن طبقات التطبيع يجب أن تهدف إلى تقليل هذا المقياس الطيفي. وقد تحققوا من ذلك باستخدام FrobNorm، حيث أظهر أنه يسمح بالتدريب عند معدلات تعلم عالية (مثل 0.25) حيث يفشل التدريب القياسي، خاصة بالنسبة لـ SGD.

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

• الاستقرار: في غياب تطبيع الدفعات (BN)، تسبب SGD في انفجار (Divergence) خلال الدورات الأولى عند معدلات التعلم التي ظل فيها Muon مستقراً واستمر في خفض الخسارة. أظهر Muon نمواً محكوماً في معايير (Norms) المعلمات والتدرج، بينما أظهر SGD انفجاراً كارثياً.
• سرعة التقارب: عند معدلات تعلم متطابقة (مثلاً $\eta=0.05$)، وصل Muon إلى معالم الدقة (70%–91%) قبل عدة دورات (Epochs) من SGD. كان نسبة التقارب التجريبية $r_t$ أقل باستمرار لـ Muon، مما يؤكد التنبؤ النظري بمعدل خطي أسرع.
• تأثير التطبيع: سمح تطبيق FrobNorm لكلا المحسنين بالتدريب باستقرار عند معدلات تعلم عالية (0.25 لـ CifarNet، و0.2 لنموذج Transformer على طراز GPT-2)، مما يؤكد الفرضية القائلة بأن التحكم في المقياس الطيفي لمدخلات الطبقة هو مبدأ عام لتمكين أحجام خطوات أكبر.

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

تدعي الورقة تقديم تفسير هندسي مبدئي لنجاح Muel التجريبي. وهي تفترض أن ميزة Muon ليست مجرد وسيلة استدلالية (Heuristic) بل هي نتيجة مباشرة لـ التسطيح الطيفي، الذي:

1. يفك الارتباط بين قيد معدل التعلم وأكبر قيمة مفردة، مما يسمح بأحجام خطوات أكبر.
2. يعمل كمسبق شرط أحادي الجانب يحسن الرقم الشرطي الفعال لتضاريس التحسين، مما يؤدي إلى تقارب أسرع.

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