The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
Abreu, Vyas, Kakade et al.
Recent efforts to accelerate LLM pretraining have focused on computationally-efficient approximations that exploit second-order structure. This raises a key question for large-scale training: how much performance is forfeited by these approximations? To probe this question, we establish a practical upper bound on iteration complexity by applying full Gauss-Newton (GN) preconditioning to transformer models of up to 150M parameters. Our experiments show that full GN updates yield substantial gains over existing optimizers, achieving a 5.4x reduction in training iterations compared to strong baselines like SOAP and Muon. Furthermore, we find that a precise layerwise GN preconditioner, which ignores cross-layer information, nearly matches the performance of the full GN method. Collectively, our results suggest: (1) the GN approximation is highly effective for preconditioning, implying higher-order loss terms may not be critical for convergence speed; (2) the layerwise Hessian structure contains sufficient information to achieve most of these potential gains; and (3) a significant performance gap exists between current approximate methods and an idealized layerwise oracle.
academic
إمكانيات تحسين الرتبة الثانية لنماذج اللغة الكبيرة: دراسة مع Gauss-Newton الكامل
تبحث هذه الورقة عن مقدار الأداء الذي تفقده التقريبات الحسابية الفعالة للطرق الحالية من الرتبة الثانية في التدريب المسبق لنماذج اللغة الكبيرة (LLM). يقوم المؤلفون بتطبيق تحسين Gauss-Newton (GN) الكامل على نموذج Transformer بـ 150 مليون معامل، مما يؤسس حدوداً عملية لتعقيد التكرار. تُظهر التجارب أن تحديثات Gauss-Newton الكاملة حققت تقليلاً بمعامل 5.4 في تكرارات التدريب مقارنة بخطوط الأساس القوية مثل SOAP و Muon. علاوة على ذلك، فإن مُحسِّن Gauss-Newton الدقيق الموزع على كل طبقة، الذي يتجاهل المعلومات عبر الطبقات، يحقق أداءً قريباً جداً من طريقة Gauss-Newton الكاملة.
مع النمو المستمر في متطلبات الحوسبة لنماذج اللغة الكبيرة، أصبح تحسين طرق التحسين استراتيجية أساسية لزيادة كفاءة التدريب. بينما تُستخدم الطرق من الرتبة الأولى التقليدية (مثل SGD و Adam) على نطاق واسع، فإن طرق الرتبة الثانية تتمتع نظرياً بسرعة تقارب أسرع وقدرة أفضل على التوسع مع الدفعات الكبيرة.
قيود طرق الرتبة الثانية الحالية: تستخدم محسِّنات الرتبة الثانية الحالية (مثل Shampoo و SOAP و Muon) تقريبات Hessian للحفاظ على الجدوى الحسابية، لكن مقدار الأداء الذي تفقده هذه التقريبات لا يزال غير واضح.
الفجوة بين النظرية والممارسة: بينما تتفوق طرق الرتبة الثانية نظرياً، فإن التكاليف العالية للتخزين والحوسبة لـ Hessian الكامل تجعل استخدام التقريبات ضرورياً في التطبيقات العملية.
السؤال البحثي الأساسي: "ما هي حدود الأداء الأساسية لتحسين الرتبة الثانية في نماذج اللغة الكبيرة؟ أي الخصائص الهيكلية لـ Hessian ضرورية لتحقيق هذه الحدود؟"
تأسيس حدود الأداء: تأسيس حدود أداء عملية لتحسين الرتبة الثانية من خلال طريقة Gauss-Newton الكاملة، مع تحقيق تحسن بمعامل 5.4 في تعقيد التكرار مقارنة بـ SOAP.
الكشف عن الهياكل الرئيسية: اكتشاف أن هيكل Hessian الموزع على كل طبقة يحتوي على معلومات كافية لتحقيق معظم مكاسب الأداء، مع أهمية محدودة لمعلومات الانحناء عبر الطبقات.
الرؤى النظرية: إثبات أن تقريب GN فعال للغاية في التحسين المسبق، مما يشير إلى أن حدود الخسارة من الرتب الأعلى قد لا تكون حاسمة لسرعة التقارب.
توسع حجم الدفعة: توسيع كبير للحجم الحرج للدفعة، مما يُظهر أداء توسع قريب من الأمثل.
بالنظر إلى معاملات النموذج θ والمدخلات x والعلامات y، يتم تعريف دالة الخسارة L(f(θ,x), y). الهدف هو تقليل الخسارة المتوقعة، مع التركيز على تعقيد التكرار (عدد الخطوات المطلوبة للوصول إلى خسارة الهدف).
لتجنب تخزين مصفوفة Hessian بشكل صريح، يتم استخدام منتجات Jacobian-Vector (JVPs) لتنفيذ طريقة معادلة وظيفياً. الفكرة الأساسية هي تحسين تقريب Taylor من الرتبة الثانية للخسارة L وتقريب Taylor من الرتبة الأولى للنموذج f.
تستشهد هذه الورقة بأعمال مهمة في مجال التحسين، بما في ذلك:
Martens (2010): العمل الرائد في تحسين خالي من Hessian
Gupta et al. (2018): محسِّن Shampoo
Jordan et al. (2024): محسِّن Muon
Vyas et al. (2025): محسِّن SOAP
التقييم الشامل: هذه ورقة بحثية عالية الجودة تؤسس حدود أداء تحسين الرتبة الثانية في تدريب نماذج اللغة الكبيرة من خلال تجارب صارمة، مما توفر رؤى نظرية مهمة وتوجيهات عملية للمجال. على الرغم من وجود قيود في التكاليف الحسابية وحجم النماذج، فإن قيمتها الأكاديمية وأهميتها التوجيهية للأبحاث المستقبلية كبيرة جداً.