Layerwise LQR for Geometry-Aware Optimization of Deep Networks

ملخص تقني: LQR الطبقي للتحسين المدرك للهندسة في الشبكات العميقة

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

توفر المحسنات المدركة للهندسة (Geometry-aware optimizers)، مثل طريقة نيوتن (Newton's method) وتدرج الطبيعة (Natural Gradient Descent - NGD)، خصائص أفضل من حيث التكييف (conditioning) والتقارب عبر استخدام معلومات الدرجة الثانية (مثل مصفوفات هسيان أو مصفوفات معلومات فيشر). ومع ذلك، فإن هذه الطرق مكلفة حوسبياً للتعلم العميق واسع النطاق لأن مصفوفات الانحناء تكون كثيفة وتربط المعلمات عبر جميع الطبقات عبر قاعدة السلسلة (chain rule). إن حل معادلة التحديث $H\Delta\theta = -g$ بشكل مباشر أمر غير ممكن.

تتعامل التقريبات القابلة للتوسع الموجودة، مثل K-FAC وShampoo والممهدات (preconditioners) ذات الصلة، مع هذه المشكلة من خلال فرض قيود هيكلية (مثل الهيكل القطري أو العامل كروينكر - Kronecker-factored) على مصفوفة الانحناء في مرحلة مبكرة من الاشتقاق. وبينما يجعل هذا عملية الارتداد (inversion) قابلة للتتبع، إلا أنه يتخلص من التفاعلات بين الطبقات قبل حل مشكلة التحسين التي تحدد التحديث. تجادل الورقة بأن هذا الفرض المبكر للهيكل يحد من قدرة هذه المحسنات على التقاط الهندسة الحقيقية لمشهد الخسارة (loss landscape)، لا سيما الارتباطات بين الطبقات الناتجة عن مخطط حساب الشبكة.

2. المنهجية: LQR الطبقي (LLQR)

يقترح المؤلفون LQR الطبقي (Layerwise LQR - LLQR)، وهو إطار عمل يعيد صياغة خطوة التحديث المدركة للهندسة كمسألة منظم تربيعي خطي (Linear Quadratic Regulator - LQR) ذات أفق زمني محدد. يفصل هذا النهج بين ديناميكيات الشبكة واختيار هندسة الهبوط، مما يسمح بتخفيف قابل للتوسع يتعلم ممهدات مهيكلة مع الاحتفاظ بالهدف المرتبط بالطبقات.

الرؤية النظرية الجوهرية:
تثبت الورقة وجود تكافؤ دقيق بين خطوة الهبوط الأشد (steepest-descent step) تحت فئة واسعة من النماذج التربيعية المستحثة بالتباعد (بما في ذلك نيوتن، Gauss-Newton، فيشر/التدرج الطبيعي، ومقاييس الطبقات المتوسطة) ومسألة LQR ذات الأفق الزمني المحدود.

• الديناميكيات: تمرير الشبكة للأمام (forward pass) يحدد ديناميكيات الاضطراب الخطي: $\delta x_{i+1} = A_i \delta x_i + B_i \delta \theta_i$، حيث $A_i$ و $B_i$ هما جاكوبيان (Jacobians) لخرائط الطبقات.
• التكلفة: التباعد المختار (مثل تباعد KL لـ NGD، أو فجوة Bregman لـ Newton) يحدد مصفوفات التكلفة التربيعية ($Q_i, R_i, M_i$) المرتبطة باضطرابات الحالة والتحكم.
• الحل الدقيق: يمكن استعادة التحديث الدقيق المدرك للهندسة عن طريق حل مسألة LQR هذه عبر تراجعات ريكاتي (Riccati recursions) الخلفية، والتي تحسب مصفوفات الكسب المحلية والملحقات (adjoints) دون تشكيل مصفوفة هسيان كثيفة عالمية.

التخفيف القابل للتوسع:
بينما لا يزال حل ريكاتي الدقيق مكلفاً حوسبياً للشبكات الكبيرة بسبب الكميات المعتمدة على الجاكوبيان، يقدم المؤلفون تخفيفاً قابلاً للتوسع. بدلاً من الحل لـ $\delta \theta$ الدقيق، يقومون بتمثيل التحديث كـ تدرج ممهد:
$$\Delta \theta_i = -U_i \nabla_{\theta_i} L(\theta)$$
حيث $U = \text{diag}(U_0, \dots, U_{N-1})$ هو ممهد عكسي مهيكل مُتعلم (على سبيل المثال: قطري، أو عامل كروينكر، أو E-KFAC).

من المهم ملاحظة أن الهيكل يتم فرضه على الممهد العكسي المُتعلم $U$، وليس على مصفوفة الانحونة نفسها. يتم تعلم الممهد عن طريق تقليل هدف LQR (المعادلة 15) عبر دفعة مصغرة (minibatch). وهذا يسمح للمحسن بتقريب الهندسة الكثيفة المرتبطة بالطبقات باستخدام كتل مهيكلة، مما يوازن فعلياً بين القدرة التعبيرية والقابلية للتوسع مع الحفاظ على اتصال مبدئي بالهندسة من الدرجة الثانية الأصلية.

التنفيذ الخوارزمي:
يقوم هذا الأسلوب بتغليف (wrap) المحسنات القياسية (مثل SGDM وAdamW). وفي كل $n$ من التكرارات، تقوم الخوارزمية بما يلي:

1. خطية (linearize) ديناميكيات الشبكة ($A_i, B_i$) وتشكيل كتل التكلفة المحلية ($Q_i, R_i, M_i$) بناءً على التباعد المختار.
2. حل مسألة تحسين داخلية لتحديث الممهد $U$ باستخدام محسن قياسي (مثل SGDM) لتقليل هدف LQR المخفف.
3. تطبيق متوسط متحرك أسي (EMA) لتثبيت $U$.
4. استخدام $U$ المحدث لتمهيد التدرجات لخطوات الحلقة الخارجية اللاحقة.

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

1. صياغة التحكم الأمثل الطبقية: تثبت الورقة أن الهبوط الأشد تحت فئة واسعة من النماذج التربيعية المستحثة بالتباعد يمكن كتابته بدقة كمسألة LQR ذات أفق زمني محدد. يوفر هذا مرجعاً نظرياً جديداً للتحديثات المدركة للهندسة التي تفصل صراحة بين ديناميكيات الشبكة واختيار المقياس (metric).
2. التخفيف القابل للتوسع عبر الممهدات المتعلمة: يقترح المؤلفون تعلم ممهدات عكسية مهيكلة مباشرة عن طريق تقليل هدف LQR. ينتج عن ذلك عائلة من المحسنات التي يمكنها استخدام هياكل قطرية، أو عامل كروينكر، أو E-KFAC مع الحفاظ على الهدف المرتبط بالطبقات الناتج عن النموذج الكثي الأصلي.
3. غلاف محسن عملي: يتم تنفيذ تحديث LLQR المريح كغلاف للمعماريات الحديثة (ResNets, Transformers) التي تعيد استخدام الممهدات المتعلمة عبر التكرارات، مما يتجنب الارتداد الصريح للانحناء ويضيف عبئاً حوسبياً طفيفاً.
4. التحقق التجريبي: تظهر التجارب المكثفة أن LLQR يحسن ديناميكيات التحسين وأداء الاختبار النهائي في مهام تصنيف الصور (CIFAR, ImageNet) والترجمة الآلية (IWSLT14). كما أنه يسرع عملية "الاستيعاب" (grokking) في نماذج Transformers.

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

• التحقق من النماذج التجريبية (Toy Validation): في دالة Rosenbrock، يطابق حل LQR الدقيق (عبر تراجع ريكاتي) طريقة نيوتن تماماً. كما أن LLQR المخفف مع ممهدات قطرية يتلاقى بشكل أسرع من تدرج الهبوط القياسي ويتتبع مسار نيوتن بشكل أوثق من تقريبات هسيان القطرية، مما يؤكد القدرة على التقاط الارتباطات بين الطبقات.
• CIFAR-10/100: في ResNet-18، أدى LLQR مع هيكل E-KFAC باستمرار إلى تحسين دقة Top-1 مقارنة بالنماذج المرجعية (SGDM, AdamW) مع زيادة طفيفة فقط في وقت التنفيذ (مثلاً: $\times 1.03$ إلى $\times 1.15$). أظهرت الممهدات القطرية تحسناً أقل، مما يشير إلى أن الهياكل الكروينكر ضرورية لالتقاط الانحناء.
• ImageNet: حقق تدريب ResNet-50 لمدة 100 حقبة باستخدام LLKR+E-KFAC مع NGD دقة Top-1 بلغت 78.05% مقارنة بـ 77.42% لنموذج SGDM المرجعي، مع عبء حوسبي قدره $\approx 1.03\times$ فقط.
• Transformers (IWSLT14): حسن LLKR+E-KFAC درجات BLEU من 34.24 إلى 34.51 في الترجمة من الألمانية إلى الإنجليزية مع تباطؤ قدره $1.16\times$.
• الاستيعاب (Grokking): في مجموعات البيانات الخوارزمية، سرع LLKR باستمرار بداية الاستيعاب (التعميم المفاجئ) من حيث عدد التكرارات ووقت التنفيذ مقارنة بالنماذج المرجعية.
• مقارنة الكفاءة: عند مقارنته بـ AdaFisher وطرق الدرجة الثانية الأخرى تحت ميزانيات زمنية متساحة، حقق LLCR دقة أعلى، مما يثبت أن الهياكل الأكثر ثراءً للممهدات (E-KFAC) يمكن جعلها عملية على نطاق واسع.

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

تضع الورقة LLKR كإطار عمل عملي للمحسنات مدركة للهندسة من الدرجة الثانية يجسد الفجوة بين المثالية النظرية والقابلية للتوسع.

• التقريب المبدئي: على عكس الطرق التي تقرب مصفوفة الانحناء أولاً، يشتق LLKR هدف التحديث من الهندسة الكثيفة ثم يضع قيوداً على فئة الممهد. يضمن هذا أن الممهد المتعلم يتم تحسينه في ظل وجود ارتباطات بين الطبقات مشفرة بواسطة ديناميكيات LQR.
• المرونة: الإطار مستقل عن نوع التباعد (يدعم Newton، NGD، إلخ) ومستقل عن الهيكل (يدعم القطرية، عامل كروينكر، E-KFAC).
• الكفاءة: من خلال توزيع تكلفة تعلم الممهد وتطبيقه دون الحاجة للارتداد الصريح، ينقل LLCR التمهيد ذو القدرة التعبيرية العالية من خيار جذاب نظرياً ولكنه غير عملي غالباً إلى نظام قابل للتطبيق حوسبياً للتعلم العميق واسع النطاق.

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