Shampoo: أمثَلة المُوتِّرات العشوائية بالتكييف المُسبَق

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

AdaGrad بمصفوفته الكاملة يفعل ذلك تماماً: يُراكم الجداءات الخارجية لمتّجهات الميل السابقة في مصفوفة واحدة $H_t = \sum_{s=1}^{t} g_s g_s^\top$ ثم يستخدم $H_t^{-1/2}$ لتكييف المتّجه. النتيجة خطوة تكيّفية لكل اتجاه في فضاء المعاملات تحقق ندماً مثالياً بشكل مُثبَت. لكن أين المشكلة؟ في شبكة بها $d$ معامل، حجم $H_t$ يصبح $d \times d$. طبقة عادية بها 1000 مدخل و1000 مخرج تعني $d = 10^6$ معامل، أي أن $H_t$ تصبح مصفوفة فيها تريليون عنصر. تخزينها وحده مستحيل، فما بالك بحساب معكوسها.

لذلك يلجأ الممارسون إلى التقريب القُطري: Adam مثلاً يتتبع فقط العزم الثاني لكل معامل على حدة. التكلفة زهيدة — ذاكرة بمقدار $O(d)$ فقط — لكن الثمن باهظ: نفقد كل المعلومات عن كيف تتفاعل المعاملات مع بعضها. الانحناء المتبادل بين المعاملات، وهو ما يجعل التكييف المُسبَق الكامل بتلك القوة، يضيع بالكامل.

الملاحظة المحورية وراء Shampoo بسيطة لدرجة مفاجئة: معاملات الشبكات العصبية مُنظَّمة أصلاً على هيئة مصفوفات ومُوتِّرات، لا متّجهات مسطّحة. أوزان الطبقة كاملة الاتصال مثلاً هي مصفوفة بأبعاد $m \times n$، وأوزان الطبقة الالتفافية مُوتِّر رباعي الأبعاد (قنوات-الدخل × العرض × الارتفاع × قنوات-الخرج). حين نُسطّح هذه البنى في متّجه واحد كما تفرض أساليب المصفوفة الكاملة، نُضيّع المعلومات البنيوية التي كان بالإمكان استغلالها لتسهيل التكييف المُسبَق.

هنا يأتي سؤال Shampoo الذكي: ماذا لو كيّفنا كل بُعد من أبعاد المُوتِّر على حدة؟ بالنسبة لـمصفوفة أوزان $W$ بأبعاد $m \times n$، بدلاً من مصفوفة تكييف واحدة ضخمة بحجم $mn \times mn$، نحتفظ بمصفوفتين أصغر بكثير — مصفوفة تكييف يسرى $L_t \in \mathbb{R}^{m \times m}$ تلتقط انحناء الصفوف (أي المخرجات)، ومصفوفة تكييف يمنى $R_t \in \mathbb{R}^{n \times n}$ تلتقط انحناء الأعمدة (أي المدخلات).

قاعدة التحديث تُراكم إحصائيات متّجه الميل بشكل مستقل لكل بُعد. في كل خطوة $t$ ومع وجود مصفوفة المتّجه $G_t$، يجري ما يلي:

• تحديث مصفوفة التكييف اليسرى: $L_t = L_{t-1} + G_t G_t^\top$ (ترصد الارتباطات بين الصفوف)
• تحديث مصفوفة التكييف اليمنى: $R_t = R_{t-1} + G_t^\top G_t$ (ترصد الارتباطات بين الأعمدة)
• اتخاذ الخطوة المكيَّفة: $W_{t+1} = W_t - \eta \, L_t^{-1/4} \, G_t \, R_t^{-1/4}$

لاحظ أن $G_t G_t^\top$ حجمها $m \times m$ فقط، و$G_t^\top G_t$ حجمها $n \times n$ — أي أن كلتيهما صغيرة وسهلة التعامل. النقطة الدقيقة هنا هي الأُس $-1/4$ بدلاً من $-1/2$ المعتاد في AdaGrad الكامل: بما أن لدينا مصفوفتَي تكييف تعملان معاً، كل واحدة تأخذ الجذر الرابع فقط، فيكون أثرهما المُركَّب مكافئاً للجذر التربيعي، مما يُنتج التناقص المألوف $O(1/\sqrt{t})$ في حجم الخطوة.

التبرير النظري لـShampoo ينبع من علاقة أنيقة بـجداء كرونيكر. حين نُسطّح مصفوفة الأوزان $W$ إلى متّجه $w = \overline{\text{vec}}(W)$، فإن تحديث Shampoo يتحوّل إلى خطوة انحدار تدريجي مكيَّفة بالمعنى المعتاد، حيث مصفوفة التكييف هي $H_t = L_t^{1/4} \otimes R_t^{1/4}$.

والنتيجة الأهم هي المبرهنة 8 من الورقة: جداء كرونيكر للإحصائيات المحسوبة لكل بُعد يحدّ من الأسفل مصفوفة الإحصائيات الكاملة. بصيغة رسمية، لمتّجهات ميل بـرتبة لا تتجاوز $r$:

$\epsilon I + \frac{1}{r} \sum_t g_t g_t^\top \;\leq\; L_T^{1/2} \otimes R_T^{1/2}$

ماذا يعني هذا عملياً؟ مصفوفة تكييف Shampoo الضمنية ليست مجرّد تقريب للمصفوفة الكاملة — بل هي أقوى منها أو تساويها (بعد تصحيح عامل الرتبة). القيم الذاتية الصغيرة، وهي الأكثر أهمية لأنها تمثّل الاتجاهات سيئة التكييف، تبقى محفوظة. بعبارة أخرى، Shampoo لا يقرّب المصفوفة الكاملة فحسب، بل يهيمن عليها بشكل مُثبَت رياضياً بمعنى المؤثرات.

الحالة المصفوفية واضحة ومباشرة، لكن قوة Shampoo الحقيقية تتجلّى عند التعامل مع مُوتِّرات أعلى درجةً. أوزان الطبقة الالتفافية مثلاً رباعية الأبعاد: قنوات-الدخل × العرض × الارتفاع × قنوات-الخرج. وShampoo يبني مصفوفة تكييف مستقلة لكل بُعد من هذه الأبعاد الأربعة.

التعميم يتمّ بسلاسة. لمُوتِّر من الدرجة $k$ بأبعاد $n_1 \times \cdots \times n_k$، يحتفظ Shampoo بعدد $k$ من مصفوفات التكييف $H_t^i \in \mathbb{R}^{n_i \times n_i}$. كل واحدة تُحدَّث عبر انقباض متّجه الميل مع نفسه على جميع الأبعاد ما عدا البُعد $i$: $H_t^i = H_{t-1}^i + G_t^{(i)}$ حيث $G_t^{(i)} = \text{mat}_i(G_t)\text{mat}_i(G_t)^\top$. والمتّجه المكيَّف يُحسب بضرب كل بُعد في $(H_t^i)^{-1/2k}$ باستخدام جداء المُوتِّر-المصفوفة $\times_i$.

هذا التعميم ليس بسيطاً رياضياً: إثبات التقارب يتطلب البرهنة على أن جداء كرونيكر لعدد $k$ من مصفوفات التكييف لا يزال يحدّ من الأسفل المصفوفة الكاملة. البرهان يعتمد على رتابة المؤثرات للمتوسطات الهندسية للمصفوفات شبه المُعرَّفة الموجبة المتبادلة الإبدال — وهي نتيجة عميقة من التحليل المصفوفي.

تحليل Shampoo النظري مبنيّ في إطار الأمثَلة المحدَّبة عبر الإنترنت. الفكرة أن المتعلم يواجه سلسلة من دوال الخسارة المحدَّبة $f_1, \ldots, f_T$ ويسعى لتقليل الندم — وهو الفرق بين خسارته الكلية وخسارة أفضل نقطة ثابتة كان يمكن اختيارها بأثر رجعي.

المبرهنة الرئيسية (المبرهنة 10) تضع حداً أعلى لندم النسخة المُوتِّرية العامة:

$\text{Regret}_T \leq \sqrt{2r}\, D \prod_{i=1}^{k} \text{Tr}\big((H_T^i)^{1/2k}\big)$

حيث $r$ هو المتوسط الهندسي لرتب متّجه الميل عبر الأبعاد، و$D$ قطر مسار المعاملات، و$H_T^i$ الإحصائيات المتراكمة للبُعد $i$. في ظل الافتراضات المعتادة، كل حدّ أثر يتدرّج بمقدار $O(T^{1/2k})$، وبالتالي جداء $k$ من هذه الحدود يعطي $O(T^{k \cdot 1/2k}) = O(\sqrt{T})$ — وهو المعدّل الأمثل نظرياً للأمثَلة المحدَّبة العشوائية.

الخوارزمية بشكلها الأساسي تحتاج إلى حساب القوى المصفوفية $(H_t^i)^{-1/2k}$ في كل خطوة، وهذا يستلزم تحليلاً ذاتياً أو تفكيكاً قيمياً مفرداً. ثلاث حيل عملية تجعل Shampoo قابلاً للمنافسة من حيث الزمن الفعلي:

• تأخير تحديث مصفوفات التكييف. بدلاً من إعادة حساب الجذور المصفوفية كل خطوة، نحسبها مرة كل 20 إلى 100 خطوة ونُعيد استخدام النتائج القديمة بينها. المؤلفون وجدوا أن هذا التأخير لا يكاد يؤثر في الدقة.
• الزخم. نُدمج زخماً معيارياً: $\bar{G}_t = \alpha \bar{G}_{t-1} + (1-\alpha) G_t$ مع $\alpha = 0.9$. مصفوفات التكييف تُحدَّث بـمتّجه الميل الخام كالمعتاد، لكن خطوة التحديث الفعلية تُتَّخذ بالمتّجه المُوسَّط بالزخم.
• التراجع إلى التقريب القُطري للأبعاد الكبيرة. حين يتجاوز بُعد ما حوالي 1200، ننتقل إلى تكييف قُطري لذلك البُعد فقط مع إبقاء المصفوفات الكاملة للأبعاد الأخرى. هذا يجعل Shampoo قابلاً للتطبيق حتى مع مُوتِّرات فيها بُعد ضخم.

المفاجأة أنه رغم هذه الحسابات الإضافية، فإن زمن تنفيذ Shampoo لكل خطوة قريب من زمن الانحدار التدريجي العشوائي وAdam وAdaGrad في الممارسة العملية. السبب أن العمليات المصفوفية قابلة للتوازي بكفاءة على وحدات المعالجة الرسومية، ومصفوفات التكييف صغيرة بما يكفي لتحليل طيفي سريع.

الورقة تختبر Shampoo على تصنيف الصور (CIFAR-10/100 بشبكات ResNet وInception) ونمذجة اللغة (LM1B بنماذج انتباه). في جميع التجارب، تقارب Shampoo أسرع بفارق واضح من الانحدار التدريجي العشوائي وAdam وAdaGrad من حيث عدد خطوات التدريب.

على CIFAR-10 بشبكة ResNet من 32 طبقة مثلاً، وصل Shampoo إلى خسارة تدريب أقل بنحو نصف عدد الخطوات التي احتاجها Adam. وفي نمذجة لغة LM1B، حقّق حيرة اختبار مماثلة لـAdam في جزء بسيط من الخطوات. والأهم أن زمن التنفيذ لكل خطوة كان متقارباً — الحساب الإضافي لبناء مصفوفات التكييف وتطبيقها مهمل مقارنة بـالتمرير الأمامي والخلفي للشبكة.

النتيجة الأكثر لفتاً هي مقارنة السرعة الفعلية: رغم العمليات المصفوفية الإضافية التي تتضمن تفكيكاً قيمياً مفرداً، يعالج Shampoo العدد نفسه تقريباً من الدُّفعات في الثانية على وحدة Tesla K40. بل في بعض الإعدادات (ResNet-55 على CIFAR-100)، كان Shampoo فعلاً أسرع لكل خطوة من الأساليب المرجعية.

Shampoo وK-FAC كلاهما يستخدم مصفوفات تكييف مُحلَّلة بـجداء كرونيكر، لكن المنطلق مختلف تماماً:

• K-FAC ينطلق من تقريب مصفوفة معلومات فيشر للشبكة بافتراض استقلال إحصائي بين التنشيطات ومتّجهات الميل المنتشرة عكسياً. وهو مرتبط جوهرياً ببنية الانتشار العكسي في الشبكات الأمامية، ويحتاج إلى أخذ عيّنات من التوزيع التنبؤي للنموذج.
• Shampoo لا يفترض شيئاً عن بنية النموذج أصلاً. كل ما يفعله هو ملاحظة أن المعاملات تعيش في فضاءات مُوتِّرية واستخدام إحصائيات المتّجه نفسه لبناء مصفوفة تكييف لكل بُعد. وضمانات التقارب تسري لأي دالة خسارة محدَّبة، لا لخسائر الشبكات العصبية وحدها.

هذا الاستقلال عن البنية هو الميزة العملية الكبرى لـShampoo: يمكنك استخدامه كبديل مباشر لأي مُحسِّن دون تعديل النموذج أو شفرة التدريب. في المقابل، K-FAC يتطلب معرفة تفصيلية بنوع كل طبقة وبنيتها، مما يُصعّب استخدامه كأداة عامة.

فتح Shampoo الباب أمام عائلة كاملة من مُحسِّنات الدرجة الثانية العملية في التعلّم العميق. وتأثيره يمتدّ في عدة اتجاهات:

• SOAP (2024) أوضح أن Shampoo بالأُس 1/2 مكافئ لتشغيل Adafactor في القاعدة الذاتية لـمصفوفة التكييف، وحسّن استقراره باستخدام متوسطات متحركة أسِّية بأسلوب Adam بدلاً من المجاميع التراكمية.
• Muon (2024) يبسّط Shampoo بالتخلّي عن التراكم الزمني كلياً: الفكرة أننا نأخذ التفكيك القيمي المفرد لـمتّجه الميل الحالي فقط ونحوّله إلى تحليله القطبي $UV^\top$، فنحصل على تحديث يُسوّي جميع القيم المفردة. Muon يقرّب هذا بتكرارات نيوتن-شولتز الأسرع من التفكيك القيمي المفرد، مما يجعله عملياً حتى مع النماذج الضخمة.
• أبحاث المعايَرة القُصوى للتحديث (muP) كشفت أن Shampoo وK-FAC كليهما يتصلان بسؤال نظري أعمق: كيف نُعاير الشبكة العصبية حتى تنتقل المعاملات الفائقة المثلى بين أحجام النماذج المختلفة؟ المعايرة الصحيحة لمعاملات الإخماد في Shampoo تُمكّن من تدريب مستقل عن عرض الشبكة.

نسخة مطوّرة من Shampoo فازت بمسابقة AlgoPerf في نسختها الافتتاحية عام 2023، فأثبتت أن التكييف المُسبَق المُهيكَل يمكن أن يتفوّق على مُحسِّنات عائلة Adam في إطار معياري موحّد.