El Potencial de la Optimización de Segundo Orden para LLMs: Un Estudio con Gauss-Newton Completo

Información Básica

• ID del Artículo: 2510.09378
• Título: The Potential of Second-Order Optimization for LLMs: A Study with Full Gauss-Newton
• Autores: Natalie Abreu (Harvard), Nikhil Vyas (Harvard/OpenAI), Sham Kakade (Harvard), Depen Morwani (Harvard)
• Clasificación: cs.LG cs.AI
• Fecha de Publicación: 10 de octubre de 2025 (preimpresión en arXiv)
• Enlace del Artículo: https://arxiv.org/abs/2510.09378

Resumen

Este artículo investiga cuánto rendimiento se pierde con las aproximaciones computacionalmente eficientes de los métodos de optimización de segundo orden existentes en el preentrenamiento de modelos de lenguaje grandes (LLMs). Los autores establecen límites superiores prácticos de complejidad iterativa aplicando el precondicionamiento completo de Gauss-Newton (GN) en un modelo Transformer de 150M parámetros. Los experimentos muestran que las actualizaciones completas de GN logran una reducción de 5.4 veces en las iteraciones de entrenamiento en comparación con líneas base sólidas como SOAP y Muon. Además, el precondicionador exacto de GN por capas, que ignora la información entre capas, alcanza casi el rendimiento del método GN completo.

Antecedentes y Motivación de la Investigación

Definición del Problema

Con el crecimiento continuo de los requisitos computacionales de los LLMs, la mejora de los métodos de optimización se ha convertido en una estrategia central para mejorar la eficiencia del entrenamiento. Aunque los métodos de primer orden tradicionales (como SGD y Adam) se utilizan ampliamente, los métodos de segundo orden poseen teóricamente velocidades de convergencia más rápidas y mejor escalabilidad en lotes grandes.

Motivación de la Investigación

1. Limitaciones de los métodos de segundo orden existentes: Los optimizadores de segundo orden actuales (como Shampoo, SOAP, Muon) utilizan aproximaciones del Hessiano para mantener la viabilidad computacional, pero no está claro cuánto rendimiento se pierde con estas aproximaciones.
2. Brecha entre teoría y práctica: Aunque los métodos de segundo orden son superiores teóricamente, debido al alto costo de almacenamiento y cálculo del Hessiano completo, la práctica requiere el uso de métodos aproximados.
3. Pregunta central de investigación: "¿Cuál es el límite fundamental de rendimiento de la optimización de segundo orden en LLMs? ¿Qué propiedades estructurales del Hessiano son necesarias para alcanzar estos límites?"

Contribuciones Principales

1. Establecimiento de límites de rendimiento: Se establece un límite superior de rendimiento práctico para la optimización de segundo orden mediante el método completo de Gauss-Newton, logrando una mejora de 5.4 veces en complejidad iterativa en comparación con SOAP.
2. Revelación de estructuras clave: Se descubre que la estructura del Hessiano por capas contiene información suficiente para lograr la mayoría de las ganancias de rendimiento, con importancia limitada de la información de curvatura entre capas.
3. Perspectivas teóricas: Se demuestra que la aproximación de GN es altamente efectiva para el precondicionamiento, sugiriendo que los términos de pérdida de orden superior podrían no ser críticos para la velocidad de convergencia.
4. Escalabilidad del tamaño de lote: Se extiende significativamente el tamaño crítico de lote, demostrando rendimiento de escalabilidad casi óptimo.

Explicación Detallada del Método

Definición de la Tarea

Dado los parámetros del modelo θ, entrada x y etiqueta y, se define la función de pérdida L(f(θ,x), y). El objetivo es minimizar la pérdida esperada, enfocándose en la complejidad iterativa (número de pasos necesarios para alcanzar la pérdida objetivo).

Principios del Método de Gauss-Newton

Fundamentos Matemáticos

La matriz Hessiana completa puede descomponerse como:

∇²θL(θ) = ∇θf(θ)ᵀ∇²zL(θ)∇θf(θ) + Σₐ(δL/δzₐ)∇²θ[f(θ)]ₐ

donde el primer término es la matriz de Gauss-Newton G y el segundo término es la curvatura del modelo.

Implementación del Algoritmo

Algoritmo 1: Método de Gauss-Newton

1. Realizar expansión de Taylor de primer orden del modelo: f⁽¹⁾θₜ(θ,x) := f(θₜ,x) + ∇f(θₜ,x)ᵀ(θ-θₜ)
2. Convexificar la pérdida: L̃θₜ(θ) := (1/b)Σ₍ₓ,ᵧ₎∈B ℓ(f⁽¹⁾θₜ(θ,x), y)
3. Construir aproximación de Taylor de segundo orden: L̃⁽²⁾θₜ(θ)
4. Resolver el problema de mínimos cuadrados: θ̂ = argminθ L̃⁽²⁾θₜ(θ)
5. Búsqueda lineal: θₜ₊₁ ← θₜ + α*(θ̂ - θₜ)

Implementación Viable en Memoria

Para evitar el almacenamiento explícito de la matriz Hessiana, se utilizan productos Jacobiano-vector (JVPs) para implementar un método funcionalmente equivalente. La idea central es optimizar la aproximación de Taylor de segundo orden de la función de pérdida L y la aproximación de Taylor de primer orden del modelo f.

Métodos Variantes

Método GN-prox-linear

Minimizar directamente la pérdida en el modelo linealizado: θ* = argminθ L̃θₜ(θ), utilizado para investigar el impacto de los términos de pérdida de orden superior.

Gauss-Newton por Capas

Para cada capa l de forma independiente:

1. Calcular la expansión de Taylor de primer orden de esa capa f⁽¹⁾θₗ,ₜ(θₗ)
2. Resolver: θₗ,ₜ₊₁ = argminθₗ L̃⁽²⁾θₗ,ₜ(θₗ)
3. Combinar las actualizaciones de todas las capas y aplicar búsqueda lineal

Configuración Experimental

Conjunto de Datos y Modelo

• Modelo: Arquitectura LLaMA con 45M y 150M parámetros
• Conjunto de datos: Conjunto de datos C4
• Longitud de secuencia: 1024

Métodos de Línea Base

• AdamW: El optimizador más ampliamente utilizado para LLMs
• Muon: Método que utiliza ortogonalización Newton-Schulz
• SOAP: Variante más reciente de Shampoo

Configuración Experimental

• Optimizador interno: Utilizar Muon para resolver el problema de mínimos cuadrados
• Tamaño de lote: Controlado mediante acumulación de gradientes, bᵢₙₜₑᵣₙₒ = 32(45M) / 128(150M)
• Programación de tasa de aprendizaje: Tres estrategias: coseno global, coseno global+interno, constante+coseno interno
• Regularización: Múltiples estrategias incluyendo decaimiento de pesos y búsqueda lineal

Resultados Experimentales

Resultados Principales

Complejidad Iterativa

En el experimento de alcanzar una pérdida de 3.25:

• Gauss-Newton: 54 pasos
• SOAP: 292 pasos (diferencia de 5.4 veces)
• Muon: aproximadamente 16 veces más
• GN por capas: 78 pasos (diferencia de solo 1.4 veces)

Escalabilidad del Tamaño de Lote

En entrenamiento fijo de 3B tokens:

• Gauss-Newton mantiene buen rendimiento incluso con tamaño de lote de 120M (pérdida 3.45)
• AdamW muestra degradación severa de rendimiento con el mismo tamaño de lote (pérdida >4.4)
• El tamaño crítico de lote se extiende significativamente, aproximándose a la tendencia de escalabilidad óptima

Estudios de Ablación

GN vs GN-prox-linear

Ambos métodos muestran rendimiento casi idéntico, indicando que los términos de pérdida de orden superior contribuyen limitadamente a la mejora de rendimiento.

GN Completo vs GN por Capas

El método por capas se aproxima al rendimiento de GN completo en la mayoría de configuraciones, demostrando que la importancia de la información de curvatura entre capas es limitada.

Hallazgos Clave

1. Importancia de la programación de tasa de aprendizaje: La programación coseno global muestra el mejor rendimiento en lotes pequeños y medianos
2. Necesidad de búsqueda lineal: Es crítica para la convergencia estable del método GN
3. Selección del optimizador interno: Muon supera a AdamW como optimizador interno

Trabajo Relacionado

Métodos de Optimización de Segundo Orden

• Optimización libre de Hessiano: Método de gradiente conjugado propuesto por Martens (2010)
• Aproximaciones diagonales del Hessiano: Métodos ligeros como AdaHessian y Sophia
• Aproximaciones por capas: Idea central de la serie de métodos Shampoo

Desarrollo de Optimizadores para LLMs

• Métodos tradicionales: Series SGD y Adam
• Métodos modernos de segundo orden: Shampoo ganó en la competencia AlgoPerf con 28% de mejora
• Métodos especializados: Muon y SOAP diseñados específicamente para LLMs

Conclusiones y Discusión

Conclusiones Principales

1. Establecimiento de límites de rendimiento: El método GN completo proporciona un objetivo de rendimiento claro para la optimización de segundo orden
2. Importancia de la estructura: La estructura del Hessiano por capas contiene información suficiente para lograr la mayoría de las ganancias
3. Efectividad de aproximaciones: Existe una brecha de rendimiento significativa entre los métodos aproximados actuales y el oráculo por capas idealizado

Limitaciones

1. Costo computacional: La implementación actual es 4-5 veces más lenta que el entrenamiento estándar
2. Restricción de escala: Los experimentos se limitan a modelos de 150M parámetros
3. Practicidad: Funciona principalmente como herramienta de análisis en lugar de optimizador práctico directo

Direcciones Futuras

1. Implementación eficiente: Desarrollar métodos de segundo orden exactos computacionalmente eficientes
2. Mejores aproximaciones: Mejorar los métodos de aproximación del Hessiano por capas
3. Escalabilidad: Validar los hallazgos en modelos más grandes

Evaluación Profunda

Fortalezas

1. Profundidad teórica: Proporciona perspectivas teóricas importantes sobre los límites de rendimiento de la optimización de segundo orden
2. Rigor experimental: Búsqueda exhaustiva de hiperparámetros y múltiples estrategias de regularización
3. Valor práctico: Proporciona objetivos claros para mejorar los métodos de segundo orden existentes
4. Innovación metodológica: Uso ingenioso de JVPs para evitar almacenamiento explícito del Hessiano

Deficiencias

1. Costo computacional: El alto costo computacional limita la aplicación práctica
2. Limitaciones de escala: No se ha validado en LLMs verdaderamente a gran escala
3. Análisis teórico: Falta explicación teórica profunda sobre por qué las aproximaciones por capas son tan efectivas

Impacto

1. Contribución académica: Proporciona un punto de referencia importante para la investigación en optimización de segundo orden
2. Orientación práctica: Indica direcciones para mejorar los métodos existentes
3. Valor metodológico: Establece un nuevo marco para evaluar métodos de segundo orden

Escenarios Aplicables

• Análisis teórico de métodos de optimización de segundo orden
• Punto de referencia de rendimiento para nuevos algoritmos de optimización
• Selección de optimización para escenarios de entrenamiento en lotes grandes

Referencias

Este artículo cita trabajos importantes en el campo de la optimización, incluyendo:

• Martens (2010): Trabajo pionero en optimización libre de Hessiano
• Gupta et al. (2018): Optimizador Shampoo
• Jordan et al. (2024): Optimizador Muon
• Vyas et al. (2025): Optimizador SOAP

Evaluación General: Este es un artículo de investigación de alta calidad que establece rigurosamente los límites de rendimiento de la optimización de segundo orden en el entrenamiento de LLMs mediante experimentos rigurosos, proporcionando perspectivas teóricas importantes y orientación práctica para el campo. Aunque existen limitaciones en costo computacional y escala, su valor académico y significado orientador para investigaciones futuras son considerables.