Convergence Rate Analysis of the AdamW-Style Shampoo: Unifying One-sided and Two-Sided Preconditioning

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

تتناول الورقة البحثية التحليل النظري لمعدل التقارب لمحسن AdamW-Style Shampoo، وهو خوارزمية متطورة فازت في مسابقة AlgoPerf لتدريب الشبكات العصبية. وبينما أظهر Shampoo (ومتغيراته مثل SOAP وMuon) أداءً تجريبيًا متفوقًا على طرق التمهيد القطرية (مثل Adam) من خلال استخدام ممهدات كاملة المصفوفة (full-matrix preconditioners)، فإن ضمانات التقارب النظرية لـ التحسين غير المحدب (nonconvex optimization) كانت إما مفقودة أو غير كافية ومحدودة.

تحديدًا، واجهت الأدبيات الحالية ثلاث فجوات رئيسية:

1. عدم التحدب (Nonconvexity): معظم التحليلات السابقة كانت مقتصرة على الإعدادات المحدبة (convex settings).
2. نطاق التمهيد (Scope of Preconditioning): كانت التحليلات السابقة لعدم التحدب (على سبيل المثال، Xie et al., 2026) محدودة بالتمهيد أحادي الجانب (one-sided) وتجاهلت التمهيد الأكثر تعقيدًا وفعالية المستخدم في التطبيقات العملية وهو ثنائي الجانب (two-sided).
3. المكونات العملية: أغفلت النظريات الحالية ميزات عملية بالغة الأهمية مثل الزخم (momentum) وتلاشي الوزن المنفصل (decoupled weight decay) بنمط AdamW، وهي عناصر لا غنى عنها لتدريب النماذج اللغوية الكبيرة الحديثة (LLMs).

تهدف الورقة إلى وضع معدل تقارب صارم لخوارزمية AdamW-Style Shampoo العامة، مع توحيد كل من التمهيد أحادي الجانب وثنائي الجانب تحت إطار عمل غير محدب.

2. المنهجية

يحلل المؤلفون الخوارزمية 1 (Algorithm 1)، التي تمثل محسن AdamW-Style Shampoo. تتضمن قاعدة التحديث ما يلي:

• المتوسطات المتحركة الأسية (EMA): تتبع العزوم الأولى ($M_k$) والعزوم الثانية ($L_k, R_k$) للتدرجات.
• التمهيد ثنائي الجانب: خطوة التحديث تُعطى بـ $X_{k+1} = (1-\lambda\eta)X_k - \eta L_{k,\epsilon}^{-\frac{1}{2p}} M_k R_{k,\epsilon}^{-\frac{1}{2q}}$، حيث $L_k$ و $R_k$ هما نواتج الضرب الخارجي المتراكمة للتدرجات (تقريبًا لـ $E[GG^T]$ و $E[G^TG]$).
• الأسس الموحدة: تسمح المعلمات $p$ و $q$ (حيث $1/p + 1/q = 1$) للخوارزمية بالتبديل بين التمهيد ثنائي الجانب ($p, q < \infty$) وأحادي الجانب ($p=1, q=\infty$ أو العكس).

الابتكارات التقنية الرئيسية:
للتعامل مع تعقيد التمهيد المصفوفي ثنائي الجانب في الإعدادات غير المحدبة، طور المؤلفون تقنيات إثبات جديدة:

1. متراجحة شاتن-p هولدر (Schatten-p Hölder Inequality): بدلاً من متراجحات هولدر القياسية القائمة على المتجهات المستخدمة في تحليل Adam، يستخدم المؤلفون متراجحة شاتن-p هولدر لتقييد النواة النووية ($\|\cdot\|_*$) للتدرج. هذا أمر بالغ الأهمية لأن التمهيد يعمل على القيم المفردة (singular values) لمعلمات المصفوفة.
2. متراجحة كوشي-شفارز المصفوفية (Matrix Cauchy-Schwarz Inequality): لتقييد المعيار الطيفي لخطوة التحديث ($L^{-1/2p} M R^{-1/2q}$)، طبقوا متراجحة كوشي-شفارز خاصة بالمصفوفات للتعامل مع التفاعل بين الممهدات اليسرى واليمنى.
3. تحليل الحد الأدنى للطيف (Spectral Floor Analysis): قدموا شرطًا (الشرط 4) يتعلق بالحد الأدنى للممهد ($\hat{\epsilon}$). وجادلوا بأن الحد الأدنى الطيفي الفعلي للممهد في الممارسة العملية أكبر بكثير من التنظيم العددي $\epsilon$ (مثل $10^{-12}$)، مما يربط الحدود النظرية بالملاحظات التجريبية.

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

• معدل التقارب الموحد: تضع الورقة أول معدل تقارب لـ AdamW-Style Shampoo في التحسين غير المحدب، مغطيةً كلاً من التمهيد أحادي الجانب وثنائي الجانب.
• اشتقاق المعدل الأمثل: أثبتوا أن متوسط النواة النووية المتوقعة للتدرج يتقارب بمعدل:
  $$\frac{1}{K} \sum_{k=1}^K \mathbb{E}[\|\nabla f(X_k)\|_*] \leq O\left( \frac{\sqrt{m+n} \cdot C}{K^{1/4}} \right)$$
  حيث $m, n$ هما أبعاد المصفوفة، $K$ هو عدد التكرارات، و $C$ يقابل الثابت في معدل SGD الأمثل.
• التبرير النظري للإعدادات العملية: يوضح التحليل لماذا لا يعيق ضبط المنظم العددي $\epsilon$ إلى قيم صغيرة للغاية (مثل $10^{-12}$) عملية التقارب. تظهر النظرية أن حد التمهيد الفعلي يتحدد بواسطة مقياس ضوضاء التدرج وليس بواسطة $\epsilon$ العددي.
• العلاقة بين المعايير (Norm Relationship): أظهر المؤلفون أن معدل النواة النووية المستمد مشابه لمعدل معيار فروبينيوس الأمثل لـ SGD (وهو $O(K^{-1/4})$) تحت الشرط المثالي أن $\|\nabla f(X)\|_* \approx \Theta(\sqrt{\min(m,n)}) \|\nabla f(X)\|_F$.

4. النتائج

• الحدود النظرية:
  • يعتمد أسوأ سيناريو للحد على $\sqrt{m+n}$، وهو ما يعد أكثر إحكامًا بشكل ملحوظ من الاعتماد على $\sqrt{mn}$ (الإجمالي للأبعاد $d$) الموجود في بعض الطرق القطرية أو الحدود الدنيا للمصفوفات الكاملة.
  • يتوافق معدل التقارب مع الحد الأدنى المعروف للتحسين العشوائي غير المحدب ($O(K^{-1/4})$) حتى عامل $\sqrt{m+n}$.
• التحقق التجريبي (تدريب GPT-2):
  • نسبة معيار التدرج: أظهرت التجارب على تدريب GPT-2 أن نسبة $\|\nabla f(X)\|_* / \|\nabla f(X)\|_F$ تظل باستمرار قريبة من الحد الأعلى النظري $\sqrt{\min(m,n)}$. وهذا يؤكد صحة الافتراض بأن معدل النواة النووية يتوافق فعليًا مع معدل معيار فروبينيوس في الممارسة العملية.
  • الحد الأدنى للطيف (Spectral Floor): لوحظ أن أصغر القيم الذاتية للممهدات ($L_k, R_k$) أكبر بعدة مراتب من $\epsilon$ العددي ($10^{-12}$) وتماثل المقياس المرتبط بالضوضاء $\hat{\sigma}^2/(m+n)$. وهذا يؤكد أن افتراض "الحد الأدنى الطيفي الكبير" غير العملي ينطبق في تدريب نماذج LLM الحقيقية.

5. الأهمية

• الربط بين النظرية والتطبيق: يوفر هذا العمل أول أساس نظري صارم لنسخة Shampoo المحددة التي فازت في مسابقة AlgoPerf، ويشرح لماذا تؤدي هذا الأداء الجيد في الواقع.
• التقدم على الطرق القطرية: يثبت هذا العمل نظريًا تفوق التمهيد غير القطري (كامل المصفوفة) على الطرق القطرية (مثل Adam) بالنسبة للمعلمات ذات البنية المصفوفية، ويوضح أن الارتباطات بين المعلمات التي يلتقطها Shampoo تؤدي إلى اتجاهات تقارب أفضل.
• توجيه المعلمات الفائقة (Hyperparameters): يقدم التحليل تبريرًا نظريًا لاختيار المعلمات الفائقة (مثل العلاقة بين الزخم $\theta$ وعامل تلاشي العزم الثاني $\beta$، وكذلك دور $\epsilon$) ويوجه الممارسين في ضبط هذه المحسنات للنماذج واسعة النطاق.
• أساس للمحسنات المستقبلية: من خلال توحيد التمهيد أحادي وجانبي وثنائي الجانب ودمج الزخم وتلاشي الوزن، يضع هذا الإطار معيارًا جديدًا لتحليل المحسنات المتقدمة الشبيهة بالرتبة الثانية في التعلم العميق.

باختصار، نجحت الورقة في توحيد التحليل النظري لـ Shampoo من خلال إثبات أنه يحقق معدل تقارب أمثل للمسائل غير المحدبة مع التقاط الهياكل المصفوفية المعقدة المتأصلة في نماذج التعلم العميق الحديثة في آن واحد.