ملخص تقني: AMUSE (أي وقت مع Muon مع تقييم مستقر للتدرج)

بيان المشكلة

تعتمد عملية تحسين التعلم العميق الحديثة عادةً على مُحسّن AdamW مقترناً بجداول زمنية محددة لمعدل التعلم (مثل تلاشي جيب التمام - cosine decay). وقد تحدت التطورات الأخيرة هذه المكونات بشكل فردي: حيث يقوم التحسين "الخالي من الجداول" (Schedule-Free) بإزالة الجداول الزمنية الصريحة عبر تقييم التدرجات عند نقاط متداخلة بين المسارات الحالية والمتوسطة، بينما يعمل Muon على تحسين هندسة التحديث للمعلمات ذات القيم المصفوفية عن طريق تعامد الزخم (orthogonalizing momentum).

على الرغم من النجاح التجريبي لـ Muon في مرحلة التدريب المسبق للنماذج اللغوية الكبيرة (LLM)، إلا أن آليته الأساسية ليست مفهومة تماماً. يحدد البحث استقراراً محدداً في Muon: فبينما يعمل التعامد على تسريع التقدم في الاتجاهات "المسطحة" ("النهر" في مشهد الخسارة)، فإنه يتسبب دون قصد في تضخيم الضجيج في الاتجاهات "الحادة" عالية الانحناء ("جدران الوادي")، مما يؤدي إلى مسارات تذبذبية. وعلى العكس من ذلك، في حين أن طرق Schedule-Free تعمل على استقرار المسارات من خلال إبقائها قريبة من النهر، إلا أنها قد تفتقر إلى التقدم الكتلي السريع الذي يوفره تعامد Muon. المشكلة الجوهرية تكمن في تصميم مُحسّن يجمع بين سرعة تقارب Muon واستقرار طريقة Schedule-Free دون الحاجة إلى جداول زمنية يدوية لمعدل التعلم.

المنهجية: AMUSE

يقترح المؤلفون AMUSE (أي وقت مع Muon مع تقييم مستقر للتدرج)، والذي يدمج تحديثات مصفوفة Muon المتعامدة مع إطار عمل Schedule-Free باستخدام معامل تداخل متغير زمنياً.

الأساس النظري: مشهد النهر والوادي (River-Valley Landscape)

يستند التحليل إلى منظور مشهد خسارة النهر والوادي:

• النهر (The River): فضاء فرعي ذو أبعاد عالية وانحناء منخفض حيث يحدث التقدم المفيد في التدريب.
• جدران الوادي (The Valley Walls): فضاء مهيمن ذو أبعاد منخفضة وانحناء عالٍ يتسبب في حدوث تذبذبات.
• آلية Muon: يعمل تعامد مصفوفة الزخم على زيادة مكون الكتلة النسبي للتحديثات، مما يسرع التقدم على طول النهر. ومع ذلك، فإنه يمنع المكونات المهيمنة من التلاشي، مما يؤدي إلى تضخيم الضجيج الصغير في الاتجاهات المهيمنة ليتحول إلى تذبذبات عبر جدران الوادي.
• آلية Schedule-Free: يعمل تقييم التدرجات عند مسار متوسط (أقرب إلى النهر) على تصفية التقلبات عالية الانحناء قبل خطوة التحديث.

تصميم الخوارزمية

يستخدم AMUSE معامل تداخل متغير زمنياً $\beta_t$ لتحديد نقطة تقييم التدرج $Y_t$ بين التسلسل الأساسي السريع $Z_t$ (تحديثات Muon) والتسلسل المتوسط المستقر $X_t$:
$$Y_t = (1 - \beta_t)Z_t + \beta_t X_t$$

تم تصميم جدول $\beta_t$ لموازنة التكيف المبكر والاستقرار في المراحل المتأخرة:

1. مرحلة الإحماء ($t \le T_0$): يتم تثبيت $\beta_t$ عند قيمة صغيرة $\beta_1$. يسمح هذا للمُحسّن بتقييم التدرجات بالقرب من تسلسل Muon السريع ($Z_t$) من أجل تكيف أولي سريع.
2. ما بعد الإحماء ($t > T_0$): تزداد $\beta_t$ تدريجياً نحو 1 وفقاً لـ:
   $$\beta_t = 1 - \left( \frac{T_0 - 1}{t - 1} \right)^\rho (1 - \beta_1)$$
   هنا، يتحكم $\rho \in [0, 1]$ في معدل التحول. مع زيادة $\beta_t$، تنتقل نقطة تقييم التدرج نحو التسلسل المتوسط $X_t$. وهذا يقلل من المكونات عالية الانحناء (المهيمنة) للتدرج قبل أن تخضع لتعامد Mu، مما يؤدي إلى استقرار المسار.

من المهم ملاحظة أن هذا الجدول مستقل عن الأفق الإجمالي للتدريب، مما يتيح التدريب في أي وقت (يمكن إيقاف التدريب في أي نقطة دون تدهور في الأداء). لا يتطلب AMUSE أي جدول زمني لتلاشي معدل التعلم.

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

1. التحليل الهندسي لـ Muon: يقدم البحث تفسيراً هندسياً لكفاءة Muon وعدم استقراره. ويوضح أن التعامد يعزز التحديثات الموجهة نحو الكتلة (مما يسرع التقدم في النهر) ولكنه يضخم أيضاً ضجيج الاتجاه المهيمن (مما يسبب تذبذبات جدران الوادي).
2. الاستقرار عبر التداخل: يظهر المؤلفون أن تقييم التدرجات بالقرب من المسار المتوسط (عبر Schedule-Free) يقلل من المكونات المهيمنة قبل التعامد، مما يخفف من التذبذبات.
3. خوارزمية AMUSE: مُحسّن مبتكر يغير ديناميكياً نقطة تقييم التدرج من مسار Muon السريع إلى المسار المتوسط المستقر. يجمع هذا بين التقارب السريع لـ Muon واستقرار طرق Schedule-Free.
4. دعم التدريب في أي وقت: من خلال استخدام جدول زمني متغير زمنياً مستقل عن إجمالي الخطوات وتجنب تلاشي معدل التعلم، يدعم AMUSE آفاق تدريب مرنة.

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

تم تقييم AMUSE عبر اختبارات متنوعة، بما في ذلك تصنيف الصور (CIFAR-10/100, SVHN, ImageNet-1k)، وتقسيم الصور (ISIC 2018)، والتدريب المسبق لنماذج LLM (نماذج Llama-style بـ 124 مليون، 720 مليون، و1 مليار معلمة على FineWeb).

• حد باريتو للأداء مقابل التكرار: يتفوق AMUSE باستمرار على (Schedule-Free) AdamW وMuon في جميع المهام.
• مكاسب الكفاءة:
  • في التدريب المسبق لنموذج 720M LLM، يصل AMUSE إلى أداء Muon النهائي بـ 1.51$\times$ خطوات أقل.
  • في ImageNet مع ResNet-50، يحقق نفس الأداء بـ 1.12$\times$ خطوات أقل.
  • في ImageNet مع ضبط دقيق لـ ViT، يتطلب 3.08$\times$ خطوات أقل.
• الاستقرار: على عكس Muon التقليدي، الذي يستفيد بشكل كبير من متوسط الأوزان الأسي اللاحق (EWA) أو تلاشي معدل التعلم، يحافظ AMUSE على مسار مستقر طوال التدريب، حيث يظهر تحسناً طفيفاً عند تطبيق هذه التقنيات اللاحقة.
• القابلية للتوسع: يتوسع AMUSE بفعالية مع أحجام الدفعات الكبيرة (تصل إلى 2048) ونطاقات النماذج الكبيرة (1 مليار معلمة)، حيث تواجه بعض متغيرات Schedule-Free صعوبات أحياناً.

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

يزعم البحث أن AMUSE يمثل نهجاً مبدئياً لتصميم المُحسّنات مشتقاً من فهم منهجي لديناميكيات التدريب وهندسة مشهد الخسارة. ومن خلال معالجة المقايضة صراحةً بين التقدم الكتلي السريع واستقرار جدران الوادي، يحقق AMUSE أداءً فائقاً دون الحاجة إلى جدولة يدوية لمعدل التعلم.

يشير المؤلفون بتواضع إلى أنه بينما يتفوق AMUSE على المُحسّنات الموجودة، فإنه يتسبب في عبء ذاكرة طفيف (نسخة حالة إضافية لمتتالية القاعدة $Z_t$) مقارنة بـ Muon التقليدي، رغم أنه يظل مشابهاً لـ AdamW وSchedule-Free AdamW. ويقترحون أن العمل المستقبلي يمكن أن يركز على تحقيق استقرار وأداء مماثلين دون هذه التكلفة الإضافية في الذاكرة. ويؤكد البحث أن منظور "النهر والوادي" يوفر إطاراً قوياً لتصميم مُحسّنات تتسم بالكفاءة والاستقرار في آن واحد.