Three Phases of Expert Routing: How Load Balance Evolves During Mixture-of-Experts Training

إليك ملخص تقني مفصل لورقة البحث بعنوان "ثلاث مراحل لتوجيه الخبراء: كيف يتطور توازن الأحمال أثناء تدريب نماذج Mixture-of-Experts" للباحث شرف الدين مزوني.

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

تعمل بنيات Mixture-of-Experts (MoE) على توسيع سعة النموذج عبر توجيه الرموز (tokens) إلى مجموعة فرعية من شبكات الخبراء المتخصصة. ويتمثل أحد التحديات الهندسية الحرجة في توازن الأحمال (load balancing): فبدون تدخل، تميل الرموز إلى التركيز على عدد قน้อย من الخبراء عالي الجودة، مما يترك الآخرين خاملين (انهيار الخبراء - expert collapse).

• الوضع الحالي: تستخدم الحلول القياسية خسارة توازن إضافية (مثل Switch Transformers) مع معامل قابل للضبط $\alpha$ لمعاقبة تركيز الأحمال.
• الفجوة: بينما تعمل هذه الطرق بفعالية في الممارسة العملية، إلا أنه لا يوجد فهم نظري لـ كيفية تطور المقايضة بين توازن الأحمال وجودة الخبراء خلال عملية التدريب. التحليلات الحالية هي عادةً تحليلات ثابتة (بعد التدريب - post-hoc) وتفشل في التقاط المسار الديناميكي لآلية التوجيه.

2. المنهجية

قام المؤلفون بنمذجة توجيه رموز MoE كـ لعبة ازدحام (congestion game) ضمن إطار ألعاب المجال المتوسط (Mean-Field Games - MFG).

• صياغة اللعبة:
  • اللاعبون: الرموز (Tokens).
  • الموارد: شبكات الخبراء.
  • دالة التكلفة: يتحمل الرمز تكلفة تعتمد على جودة الخبير ($q_i$) والازدحام ($\gamma \mu_i$)، حيث $\mu_i$ هو الحمل على الخبير $i$ و $\gamma$ هو معامل الازدحام.
  • التوازن: يتم نمذجة توزيع التوجيه كالتوازن الفريد للعبة جهد (potential game) مع تنظيم الإنتروبيا (softmax).
• المعلم الأساسي: معامل الازدحام الفعال ($\gamma_{eff}$). يحدد هذا المعلم الوحيد قوة المقايضة بين التوازن والجودة في أي نقطة زمنية أثناء التدريب. وهو يلتقط كلاً من الخسارة الإضافية الصريحة والتوازن الضمني الذي تعلمه عبر التدرج الاشتقاقي.
• التحديد: يتم تقدير $\gamma_{eff}$ عن طريق مطابقة توازن MFG النظري مع توزيعات الأحمال المرصودة من نقاط التحقق (checkpoints) أثناء التدريب. يقوم المؤلفون بتقليل مسافة $L_1$ بين الحمل المرصود وخريطة الاستجابة المثلى النظرية.
• النماذج التي تم تحليلها:
  • OLMoE-1B-7B: 20 نقطة تحقق (أخذ عينات كثيفة في منطقة "الاندفاع" المبكرة).
  • OpenMoE-8B: 6 نقاط تحقق (بنية مختلفة: 32 خبيراً، 4 طبقات MoE).
• التوسعات:
  • MFG متعدد الأنواع: تجميع الرموز في $K$ من الأنواع لنمذجة عدم التجانس، مما يحسن التنبؤ بالحمل.
  • التفكيك: فصل $\gamma_{eff}$ إلى مكونين: صريح ($\gamma_{explicit} = \alpha M$) وضمني ($\gamma_{implicit}$).

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

1. مسار المراحل الثلاث: اكتشاف مسار غير رتيب على شكل حرف (U) مقلوب لـ $\gamma_{eff}$ أثناء التدريب المسبق، وهو أمر غير مرئي للتحليل الثابت.
2. الأمانة النظرية: إثبات أنه بالنسبة للنماذج المتوازنة جيداً، فإن نموذج MFG أحادي النوع يكافئ رياضياً الـ softmax المقيّد بالحرارة (temperature-scaled softmax). إن منظور نظرية الألعاب لا يحسن دقة التنبؤ الثابت، ولكنه يوفر تفسيراً هيكلياً لـ لماذا يعمل softmax وما الذي يمثله "معامل الحرارة" (معامل الازدحام).
3. تفكيك الازدحام الفعال: إثبات أن المُحسّن (optimizer) يستوعب التوازن داخلياً بما يتجاوز إشارة الخسارة الصريحة بكثير. عند التقارب، يكون $\gamma_{eff}$ أكبر بـ $\approx 13$ مرة من $\gamma_{explicit}$.
4. تشخيص نطاق العمل: تطوير مقاييس (مثل نسبة $K/M$ وانتشار الاستمرارية $\epsilon_l$) لتحديد مكان تطبيق نموذج MFG وأين ينهار (مثل الطبقات المتأخرة ذات التمايز العالي للرموز).

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

أ. ديناميكيات تدريب المراحل الثلاث

كشف تتبع $\gamma_{eff}$ عبر 20 نقطة تحقق لنموذج OLMoE-1B-7B عن ثلاث مراحل متميزة:

1. المرحلة 1: الاندفاع (الخطوات 5K–50K):
   • يرتفع $\gamma_{eff}$ بحدة من $\approx 14$ إلى ذروة تبلغ 36–39 (عند الخطوات 30K–40K).
   • التفسير: يتعلم الموجه (router) موازنة الحمل بقوة. تزداد إنتروبيا التوجيه، ويعطي النظام الأولوية للاستخدام (utilization) على حساب جودة الخبير الخام.
2. المرحلة 2: الاستقرار (الخطوات 100K–400K):
   • يستقر $\gamma_{eff}$ عند $\approx 24–28$.
   • التفسير: يجد الموجه نقطة التشغيل الخاصة به. يبدأ الخبراء في التخصص (انخفاض تشتت الجودة $B_0$)، لكن قيد التوازن يظل ثابتاً.
3. المرحلة 3: الاسترخاء (الخطوات 400K–1.2M):
   • ينخفض $\gamma_{eff}$ بشكل كبير من $\approx 27$ إلى 8.5.
   • التفسير: بمجرد تمايز الخبراء بشكل كافٍ، "يرخي" الموجه قيد التوازن للتركيز على الجودة. يضحي النموذج بالانتظام لصالح الانتقائية.

• التكرار: تم تكرار نمط (U) المقلوب هذا على OpenMoE-8B، مما يؤكد أنه ليس ناتجاً عن بنية محددة.
• الضبط الدقيق (Fine-Tuning): نمط المراحل الثلاث خاص بـ التدريب المسبق (pretraining). أثناء الضبط الدقيق (annealing)، يظل $\gamma_{eff}$ مستقراً، مما يشير إلى أن توازن التوجيه قد استقر بالفعل.

ب. الأداء الثابت مقابل الديناميكي

• التكافؤ الثابت: في النموذج المتقارب، لا يمكن تمييز نموذج MFG أحادي النوع عن الـ temperature-scaled softmax إحصائياً (Held-out $L_1$: 0.199 مقابل 0.200). لا تضيف نظرية الألعاب قوة تنبؤية لقطة زمنية واحدة.
• القيمة الديناميكية: تكمن القيمة بالكامل في الديناميكيات. يكشف المسار غير الرتيب أن التدريب المبكر يعطي الأولوية للتوازن، بينما التدريب المتأخر يعطي الأولوقة للجودة.

ج. هيمنة التوازن الضمني

• عند التقارب، تساهم الخسارة الإضافية الصريحة بـ $\gamma_{explicit} = 0.64$.
• أما الازدحام الفعال الذي تم قياسه فهو $\gamma_{eff} = 8.5$.
• النتيجة: يستوعب المُحسّن (optimizer) ازدحاماً فعالاً أكبر بـ 13 مرة مما توفره الخسارة الصريحة. تعمل الخسارة الإضافية كـ "بذرة"، لكن ديناميكيات التدرج تزيد من حجم التوازن داخلياً.

د. التوسع متعدد الأنواع

• من خلال تجميع الرموز في أنواع (مثل $K=4$)، يحسن نموذج Multi-type MFG دقة التنبؤ بالحمل بنسبة 30% في المتوسط عبر جميع الطبقات الـ 16 مقارنة بالنموذج أحادي النوع.
• يؤكد هذا أن عدم تجانس الرموز كبير، رغم أن تأثير الاقتران (coupling effect) أقل بروزاً في النماذج المتوازنة جيداً.

5. الأهمية والآثار المترتبة

• أداة تشخيص جديدة: يعمل $\gamma_{eff}$ كمراقب آني لتدريب MoE. يمكن أن يشير الانخفاض المبكر في $\gamma_{eff}$ أو الفشل في الوصول إلى ذروة "الاندفاع" إلى انهيار الخبراء أو عدم استقرار التدريب.
• إعادة صياغة الخسارة الإضافية: تتحدى الورقة الرأي القائل بأن الخسارة الإضافية هي المحرك الرئيسي للتوازن. بدلاً من ذلك، تقترح أن الخسارة هي مجرد محفز لآلية توازن داخلية أكبر بكثير يقودها المُحسّن.
• استراتيجية التحسين: تشير النتائج إلى إمكانية وجود جداول زمنية للتوازن التكيفي (adaptive balance schedules). يمكن للممارسين نظرياً تقليل معامل التوازن الصريح ($\alpha$) خلال مرحلة "الاسترخاء" (المرحلة 3) للسماح للموجه بالتركيز على الجودة دون زعزعة استقرار النظام.
• رؤية نظرية: توفر الورقة تبريراً صارماً من نظرية الألعاب لشيوع استخدام softmax في موجهات MoE، بربط معلم "الحرارة" مباشرة بمعامل الازدحاد في لعبة القدر (potential game).

الخاتمة

تثبت الورقة أن تدريب MoE ليس مجرد تحسين ثابت لهدف ثابت، بل هو تطور ديناميكي عبر ثلاثة أنظمة متميزة: الاندفاع، الاستقرار، والاسترخاء. ومن خلال نمذجة التوجيه كلعبة ازدحام، يكشف المؤلفون أن "حرارة" الموجه هي معلم ديناميكي يعكس التوتر بين توازن الأحمال وتخصص الخبراء، وهي ظاهرة غير مرئية للتحليل القياسي لما بعد التدريب.