ملخص تقني: PowerStep: تحسين تكيفي موفر للذاكرة عبر التدرج المنحدر لـ $\ell_p$-norm

بيان المشكلة

أصبحت المحسنات التكيفية (Adaptive optimizers)، ولا سيما Adam وAdamW، هي المعيار لتدريب الشبكات العصبية واسعة النطاق (مثل Transformers) من خلال تهيئة تحديثات التدرج باستخدام إحصائيات العزم الأول والثاني المتراكمة. ومع ذلك، فإن هذه الطرق تسبب عبئًا كبيرًا على الذاكرة لأنها تحتفظ بمراكم العزم الثاني ($v_t$) بنفس أبعاد معاملات النموذج، مما يتطلب عادةً دقة 32-بت عائمة (fp32) كاملة لضمان الاستقرار. هذا "الجدار الذاكراتي" يضاعف بصمة حالة المحسن بالنسبة لـ Stochastic Gradient Descent (SGD) مع الزخم، مما يحد من القدرة على التوسع للنماذج التي تحتوي على مئات المليارات من المعلمات.

المحاولات الحالية للتخفيف من ذلك تشمل تحليل مصفوفة العزم الثاني، أو المشاركة عبر الكتل، أو تطبيق التكميم (quantization) منخفض الدقة. ومع ذلك، غالبًا ما يؤدي التكميم إلى عدم استقرار التدريب أو التباعد في الطرق التكيفية بسبب الحساسية العالية لمقدر العزم الثاني (تحديدًا عملية الجذر التربيعي العكسي) تجاه أخطاء التكميم. بدلاً من ذلك، تفشل بعض الطرق التي تتخلص تمامًا من العزم الثاني في مطابقة سرعة تقارب Adam في المهام واسعة النط scales.

المنهجية

يقدم المؤلفون PowerStep، وهو محسن موفر للذاكرة يحقق التكيف على مستوى الإحداثيات دون الحاجة لتخزين إحصائيات العزم الثاني. تم اشتقاق هذه الطريقة من مبادئ التدرج المنحدر لـ $\ell_p$-norm.

الأساس النظري

في الفضاءات المعيارية العامة، يكون اتجاه التدرج المنحدر $v^*$ الذي يقلل دالة الهدف $f(\theta)$ بأقصى قدر تحت قيد $\ell_p$-norm هو:
$$v^*_i \propto -\text{sign}(g_i) \cdot |g_i|^{\frac{1}{p}-1}$$
حيث $g_i$ هو مكون التدرج. يؤدي هذا إلى "تحويل قوة مُوقع" (signed power transform) $\Phi_\beta(x) = \text{sign}(x) \odot |x|^\beta$، حيث $\beta = 1/(p-1) \in [0, 1]$.

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

على عكس طريقة Powerball، التي تطبق هذا التحويل على التدرجات الخام، تقوم PowerStep بتطبيق التحويل على مخزن زخم الكرة الثقيلة ($m_t$). قاعدة التحديث هي:

1. تحديث الزخم: $m_t = \gamma m_{t-1} + g_t$
2. التحديث المتحول: $\theta_t = \theta_{t-1} - \eta_t \Phi_\beta(m_t)$

يوفر هذا التصميم مهيئًا قطريًا $D_t = \text{diag}(|m_t|^{\beta-1})$. بالنسبة لـ $\beta \in (0, 1)$، يكون الأس $\beta-1$ سالبًا، مما يعني أن حجم الخطوة يتناسب عكسيًا مع مقدار الزخم المتراكم. يتم تضخيم مقادير الزخم الصغيرة (الاتجاهات المسطحة)، بينما يتم تخميد المقاديد الكبيرة (الاتجاهات شديدة الانحدار). يوفر هذا تكيفًا على مستوى الإحداثيات مشابهًا لـ Adam ولكن يتم حسابه لحظيًا من مخزن الزخم، مما يلغي الحاجة إلى مقدر العزم الثاني.

التنفيذ منخفض الدقة

تصميم المخزن الواحد لـ PowerStep مهيأ طبيعيًا للتكميم الشديد. يقترح المؤلفون تخزين مخزن الزخم $m_t$ فقط بدقة int8 مع إجراء تحويل القوة المُوقع وتحديثات المعاملات بدقة كاملة. وهذا يتناقض مع Adam، حيث يؤدي تكميم مخزن العزم الثاني $v_t$ إلى تراكم كارثي للأخطاء بسبب عدم الخطية في عمليات الجذر التربيعي والقسمة.

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

1. تصميم الخوارزمية: تم اشتقاق PowerStep من مبادئ التدرج المنحدر لـ $\ell_p$-norm، مما يثبت أن تطبيق تحويل القوة المُوقع على مخزن زخم الكرة الثقيلة يحقق تكيفًا على مستوى الإحداثيات دون إحصائيات العزم الثاني.
2. التحليل النظري: تثبت الورقة أن PowerStep يحقق معدل التقارب الأمثل $O(1/\sqrt{T})$ للتحسين العشوائي غير المتصل (non-convex stochastic optimization) تحت الافتراضات المنتظمة القياسية (L-smoothness، التدرجات المحدودة، إل_خ).
3. التحقق التجريبي: تُظهر التجارب المكثفة على نماذج Transformer تتراوح من 124 مليون إلى 235 مليار معلمة (بما في ذلك البنى الكثيفة وبنى Mixture-of-Experts) أن PowerStep يطابق سرعة تقارب AdamW مع استخدام نصف ذاكرة حالة المحسن فقط.
4. القوة تجاه التكميم: عند الجمع مع التكميم الشديد int8، تظل PowerStep مستقرة عددياً وتقلل من ذاكرة المحسن بنحو 8 أضعاف مقارنة بـ AdamW كامل الدقة. في المقابل، يتباعد Adam تحت نفس ظروف التكميم.

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

• سرعة التقارب: عبر نماذج من 124 مليون إلى 235 مليار معلمة، تطابق PowerStep مسارات خسارة التدريب لـ AdamW. أظهرت الطرق ذات الكفاءة في الذاكرة المرتبطة (pbSGDM، SignSGD، AdamS) تقاربًا أبطأ أو عدم استقرار في النماذج الأكبر.
• كفاءة الذاكرة: تقلل PowerStep من بصمة ذاكرة حالة المحسن بنسبة 50% مقارنة بـ AdamW (fp32) لأنها تخزن مخزنًا واحدًا بدلًا من اثنين.
• أداء التكميم: تحت تكميم int8 لحالة المحسن:
  • AdamW: ينهار التدريب فورًا بسبب تراكم خطأ التكميم في مقدر العزم الثاني.
  • PowerStep: تظل مستقرة وتطابق تقارب نظيرتها كاملة الدقة.
  • تقليل الذاكرة: تحقق PowerStep (int8) تقليلاً بنحو 8 أضعاف في ذاكرة المحسن مقارنة بـ AdamW (fp32). على سبيل المثال، في نموذج بـ 235 مليار معلمة، تنخفض الذاكرة لكل وحدة معالجة عصبية (NPU) من 6768 ميجابايت (AdamW) إلى 872.4 ميجابايت (PowerStep int8).
• الحساسية للمعلمات الفائقة: تتميز الطريقة بالمتانة تجاه معامل الزخم $\gamma$ في النطاق $[0.85, 0.95]$. يوفر أس القوة $\beta = 0.1$ أفضل توازن بين التقدم السريع في البداية والاستقرار طويل الأمد؛ بينما يؤدي $\beta=0$ (المكافئ لـ SignSGD) إلى الانهيار، و $\beta=0.2$ يؤدي إلى تقارب بطيء جدًا.

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

تدعي الورقة أن PowerStep توفر بديلًا مبدئيًا، وقابلًا للتوسع، وفعالًا في الموارد للتدريب واسع النطاق. من خلال إنشاء اتصال مباشر بين التحسين التكيفي وهندسة $\ell_p$-norm، تحقق PowerStep المثالية النظرية للطرق التكيفية دون تكلفة ذاكرة تقدير العزم الثاني.

يؤكد المؤلفون أن التخلص من مخزن العزم الثاني لا يقلل الذاكرة للنصف فحسب، بل يحسن بشكل أساسي المتانة تجاه التكميم منخفض الدقة. وهذا يسمح بتقليل إضافي في ذاكرة المحسن بمقدار رتبة مقدار (إجمالي ~8 أضعاف مقابل AdamW) دون التضحية بسرعة التقارب أو جودة النموذج النهائية. يشير العمل إلى أن دمج زخم الكرة الثقيلة مع التحويلات غير الخطية يمكن أن يحل محل آليات العزم الثاني المعقدة لـ Adam، مما يوفر مسارًا أبسط وأكثر قابلية للتوسع لتدريب النماذج المستقبلية واسعة النطاق.

تشمل اتجاهات العمل المستقبلي التي حددها المؤلفون دمج PowerStep مع التحديثات المصفوفة المهيكلة (مثل محسن Muon) واستكشاف تكميم أكثر عدوانية (مثل int4).