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
LLMs के लिए द्वितीय-क्रम अनुकूलन की संभावना: पूर्ण गॉस-न्यूटन के साथ एक अध्ययन
यह पेपर बड़े भाषा मॉडल (LLM) पूर्व-प्रशिक्षण में मौजूदा द्वितीय-क्रम अनुकूलन विधियों के कम्प्यूटेशनल रूप से कुशल सन्निकटन द्वारा कितना प्रदर्शन खोया जाता है, इसका अध्ययन करता है। लेखकों ने 150M पैरामीटर के Transformer मॉडल पर पूर्ण गॉस-न्यूटन (GN) पूर्वशर्त लागू करके पुनरावृत्ति जटिलता के लिए व्यावहारिक ऊपरी सीमा स्थापित की है। प्रयोग दर्शाते हैं कि पूर्ण GN अपडेट SOAP और Muon जैसी मजबूत आधार रेखाओं की तुलना में प्रशिक्षण पुनरावृत्तियों में 5.4 गुना कमी प्राप्त करते हैं। इसके अलावा, परतों के बीच की जानकारी को अनदेखा करने वाली सटीक परत-दर-परत GN पूर्वशर्त लगभग पूर्ण GN विधि के प्रदर्शन तक पहुंचती है।
LLM की कम्प्यूटेशनल आवश्यकताओं में निरंतर वृद्धि के साथ, अनुकूलन विधियों में सुधार प्रशिक्षण दक्षता बढ़ाने की मूल रणनीति बन गया है। पारंपरिक प्रथम-क्रम विधियां (जैसे SGD और Adam) व्यापक रूप से उपयोग की जाती हैं, लेकिन द्वितीय-क्रम विधियां सैद्धांतिक रूप से तेजी से अभिसरण और बेहतर बड़े-बैच स्केलिंग क्षमता प्रदान करती हैं।
मौजूदा द्वितीय-क्रम विधियों की सीमाएं: वर्तमान द्वितीय-क्रम अनुकूलक (जैसे Shampoo, SOAP, Muon) कम्प्यूटेशनल व्यवहार्यता बनाए रखने के लिए Hessian के सन्निकटन का उपयोग करते हैं, लेकिन ये सन्निकटन कितना प्रदर्शन खोते हैं यह स्पष्ट नहीं है।
सिद्धांत और व्यवहार का अंतराल: हालांकि द्वितीय-क्रम विधियां सैद्धांतिक रूप से उत्तम हैं, पूर्ण Hessian के भंडारण और कम्प्यूटेशनल लागत के कारण, व्यावहारिक अनुप्रयोगों में सन्निकटन विधियों का उपयोग करना आवश्यक है।
मूल अनुसंधान प्रश्न: "LLM में द्वितीय-क्रम अनुकूलन की मौलिक प्रदर्शन सीमा क्या है? इन सीमाओं को प्राप्त करने के लिए Hessian के कौन से संरचनात्मक गुण आवश्यक हैं?"
प्रदर्शन ऊपरी सीमा स्थापित करना: पूर्ण गॉस-न्यूटन विधि के माध्यम से द्वितीय-क्रम अनुकूलन के लिए व्यावहारिक प्रदर्शन ऊपरी सीमा स्थापित की गई है, पुनरावृत्ति जटिलता में SOAP की तुलना में 5.4 गुना सुधार प्राप्त किया गया है।
मुख्य संरचना का खुलासा: पता चला कि परत-दर-परत Hessian संरचना में अधिकांश प्रदर्शन लाभ प्राप्त करने के लिए पर्याप्त जानकारी है, परतों के बीच की वक्रता जानकारी की सीमित महत्ता है।
सैद्धांतिक अंतर्दृष्टि: साबित किया कि GN सन्निकटन पूर्वशर्त के लिए अत्यधिक प्रभावी है, जो सुझाता है कि उच्च-क्रम हानि शर्तें अभिसरण गति के लिए महत्वपूर्ण नहीं हो सकती हैं।
बैच आकार स्केलिंग: महत्वपूर्ण रूप से महत्वपूर्ण बैच आकार का विस्तार किया, लगभग इष्टतम स्केलिंग प्रदर्शन प्रदर्शित किया।
मॉडल पैरामीटर θ, इनपुट x और लेबल y दिए गए हैं, हानि फ़ंक्शन L(f(θ,x), y) को परिभाषित करें। लक्ष्य अपेक्षित हानि को कम करना है, पुनरावृत्ति जटिलता (लक्ष्य हानि तक पहुंचने के लिए आवश्यक चरणों की संख्या) पर ध्यान केंद्रित करते हुए।
Hessian मैट्रिक्स को स्पष्ट रूप से संग्रहीत करने से बचने के लिए, Jacobian-वेक्टर गुणन (JVPs) का उपयोग करके कार्यात्मक रूप से समतुल्य विधि को लागू करें। मूल विचार हानि फ़ंक्शन L के द्वितीय-क्रम टेलर सन्निकटन और मॉडल f के प्रथम-क्रम टेलर सन्निकटन को अनुकूलित करना है।
यह पेपर अनुकूलन क्षेत्र के महत्वपूर्ण कार्यों का उद्धरण देता है, जिनमें शामिल हैं:
Martens (2010): Hessian-मुक्त अनुकूलन का अग्रणी कार्य
Gupta et al. (2018): Shampoo अनुकूलक
Jordan et al. (2024): Muon अनुकूलक
Vyas et al. (2025): SOAP अनुकूलक
समग्र मूल्यांकन: यह एक उच्च-गुणवत्ता वाला अनुसंधान पत्र है जो कठोर प्रयोगों के माध्यम से LLM प्रशिक्षण में द्वितीय-क्रम अनुकूलन की प्रदर्शन ऊपरी सीमा स्थापित करता है, इस क्षेत्र को महत्वपूर्ण सैद्धांतिक अंतर्दृष्टि और व्यावहारिक मार्गदर्शन प्रदान करता है। कम्प्यूटेशनल लागत और स्केल सीमाओं के बावजूद, इसका शैक्षणिक मूल्य और भविष्य के अनुसंधान के लिए मार्गदर्शन महत्वपूर्ण है।