Layerwise LQR for Geometry-Aware Optimization of Deep Networks

Resumen Técnico: LQR por Capas para Optimización Consciente de la Geometría en Redes Profundas

1. Planteamiento del Problema

Los optimizadores conscientes de la geometría, como el método de Newton y el Descenso del Gradiente Natural (NGD), ofrecen un acondicionamiento y propiedades de convergencia superiores al utilizar información de segundo orden (por ejemplo, matrices Hessiana o de Información de Fisher). Sin embargo, estos métodos son computacionalmente prohibitivos para el aprendizaje profundo a gran escala debido a que las matrices de curvatura son densas y acoplan los parámetros a través de todas las capas mediante la regla de la cadena. Resolver directamente la ecuación de actualización $H\Delta\theta = -g$ es inviable.

Las aproximaciones escalables existentes, como K-FAC, Shampoo y otros precondicionadores relacionados, abordan esto imponiendo restricciones estructurales (por ejemplo, bloque-diagonal, factorizadas por Kronecker) en la matriz de curvatura tempranamente en la derivación. Si bien esto hace que la inversión sea tratable, descarta las interacciones entre capas antes de que incluso se resuelva el problema de optimización que define la actualización. El artículo argumenta que esta imposición estructural prematura limita la capacidad de estos optimizadores para capturar la verdadera geometría del paisaje de pérdida, particularmente los acoplamientos intercapas inducidos por el gráfico de computación de la red.

2. Metodología: LQR por Capas (LLQR)

Los autores proponen LQR por Capas (LLQR), un marco que reformula el paso de actualización consciente de la geometría como un problema de Regulador Lineal Cuadrático (LQR) de horizonte finito. Este enfoque separa la dinámica de la red de la elección de la geometría de descenso, permitiendo una relajación escalable que aprende precondicionadores estructurados mientras retiene el objetivo acoplado por capas.

Insight Teórico Central:
El artículo establece una equivalencia exacta entre el paso de descenso más pronunciado bajo una amplia clase de modelos cuadráticos inducidos por divergencias (incluyendo Newton, Gauss-Newton, Fisher/gradiente natural y métricas de capas intermedias) y un problema de LQR de horizonte finito.

• Dinámica: El paso hacia adelante de la red neuronal define dinámicas de perturbación lineales: $\delta x_{i+1} = A_i \delta x_i + B_i \delta \theta_i$, donde $A_i$ y $B_i$ son las jacobianas de los mapas de la capa.
• Costo: La divergencia elegida (por ejemplo, divergencia KL para NGD, brecha de Bregman para Newton) define las matrices de costo cuadrático ($Q_i, R_i, M_i$) asociadas con las perturbaciones de estado y control.
• Solución Exacta: La actualización exacta consciente de la geometría puede recuperarse resolviendo este problema LQR mediante recursiones de Riccati hacia atrás, que calculan matrices de ganancia locales y adjuntos sin formar el Hessiano global denso.

Relajación Escalable:
Aunque la solución exacta de Riccati sigue siendo computacionalmente costosa para redes grandes debido a cantidades dependientes de la jacobiana, los autores introducen una relajación escalable. En lugar de resolver para la actualización exacta $\delta \theta$, parametrizan la actualización como un gradiente precondicionado:
$$\Delta \theta_i = -U_i \nabla_{\theta_i} L(\theta)$$
donde $U = \text{diag}(U_0, \dots, U_{N-1})$ es un precondicionador inverso estructurado aprendido (por ejemplo, diagonal, factorizado por Kronecker o E-KFAC).

Crucialmente, la estructura de bloques se impone en el precondicionador aprendido $U$, no en la matriz de curvatura en sí misma. El precondicionador se aprende minimizando el objetivo LQR (Ecuación 15) sobre un minibatch. Esto permite que el optimizador aproxime la geometría densa y acoplada por capas utilizando bloques estructurados, intercambiando efectivamente expresividad por escalabilidad mientras mantiene una conexión principista con la geometría de segundo orden original.

Implementación Algorítmica:
El método envuelve optimizadores estándar (por ejemplo, SGDM, AdamW). Periódicamente (cada $n$ iteraciones), el algoritmo:

1. Linealiza las dinámicas de la red ($A_i, B_i$) y forma bloques de costo locales ($Q_i, R_i, M_i$) basados en la divergencia elegida.
2. Resuelve un problema de optimización interno para actualizar el precondicionador $U$ utilizando un optimizador estándar (por ejemplo, SGDM) para minimizar el objetivo LQR relajado.
3. Aplica un Promedio Móvil Exponencial (EMA) para estabilizar $U$.
4. Utiliza la $U$ actualizada para precondicionar los gradientes para los pasos subsiguientes del bucle externo.

3. Contribuciones Clave

1. Formulación de Control Óptimo por Capas: El artículo demuestra que el descenso más pronunciado bajo una amplia clase de modelos cuadráticos inducidos por divergencias puede escribirse exactamente como un problema de LQR de horizonte finito. Esto proporciona una nueva referencia teórica para actualizaciones conscientes de la geometría que separan explícitamente la dinámica de la red de la elección de la métrica.
2. Relajación Escalable mediante Precondicionadores Aprendidos: Los autores proponen aprender precondicionadores inversos estructurados directamente minimizando el objetivo LQR. Esto genera una familia de optimizadores que pueden utilizar estructuras diagonales, factorizadas por Kronecker o E-KFAC, mientras preservan el objetivo acoplado por capas inducido por el modelo denso original.
3. Envoltorio de Optimizador Práctico: La actualización LLQR relajada se implementa como un envoltorio para arquitecturas modernas (ResNets, Transformers) que reutiliza precondicionadores aprendidos a través de iteraciones, evitando la inversión explícita de curvatura y añadiendo una sobrecarga computacional moderada.
4. Validación Empírica: Experimentos extensivos muestran que LLQR mejora la dinámica de optimización y el rendimiento final de prueba en benchmarks de clasificación de imágenes (CIFAR, ImageNet) y traducción automática (IWSLT14). También acelera el "grokking" en Transformers.

4. Resultados Experimentales

• Validación de Juguetes: En la función de Rosenbrock, la solución LQR exacta (mediante recursión de Riccati) coincide perfectamente con el método de Newton. El LLQR relajado con precondicionadores bloque-diagonales converge más rápido que el descenso de gradiente estándar y rastrea la trayectoria de Newton más de cerca que las aproximaciones de Hessiano diagonal, validando la capacidad del método para capturar acoplamientos intercapas.
• CIFAR-10/100: En ResNet-18, LLQR con estructura E-KFAC mejora consistentemente la precisión Top-1 sobre las líneas base (SGDM, AdamW) con solo un aumento moderado en el tiempo de reloj (por ejemplo, $\times 1.03$ a $\times 1.15$). Los precondicionadores diagonales mostraron menos mejora, sugiriendo que las estructuras de Kronecker son necesarias para capturar la curvatura.
• ImageNet: Entrenando ResNet-50 durante 100 épocas, LLQR+E-KFAC con NGD alcanzó una precisión Top-1 del 78.05% en comparación con el 77.42% de la línea base SGDM, con una sobrecarga computacional de solo $\approx 1.03\times$.
• Transformers (IWSLT14): LLQR+E-KFAC mejoró las puntuaciones BLEU de 34.24 a 34.51 en la traducción del alemán al inglés con una ralentización de $1.16\times$.
• Grokking: En conjuntos de datos algorítmicos, LLQR aceleró consistentemente el inicio del grokking (generalización repentina) en términos de conteo de iteraciones y tiempo de reloj en comparación con las líneas base.
• Comparación de Eficiencia: Al compararse con AdaFisher y otros métodos de segundo orden bajo presupuestos de tiempo de reloj coincidentes, LLQR logró una mayor precisión, demostrando que estructuras de precondicionador más ricas (E-KFAC) pueden hacerse prácticas a gran escala.

5. Significado y Afirmaciones

El artículo posiciona a LLQR como un marco práctico para métodos de segundo orden conscientes de la geometría que cierra la brecha entre la optimalidad teórica y la escalabilidad.

• Aproximación Principista: A diferencia de los métodos que aproximan primero la matriz de curvatura, LLQR deriva el objetivo de actualización desde la geometría densa y luego restringe la clase de precondicionador. Esto asegura que el precondicionador aprendido se optimice en presencia de acoplamientos entre capas codificados por la dinámica LQR.
• Flexibilidad: El marco es agnóstico a la divergencia (soportando Newton, NGD, etc.) y agnóstico a la estructura (soportando diagonal, Kronecker, E-KFAC).
• Eficiencia: Al amortizar el costo de aprender el precondicionador y aplicarlo sin inversión, LLQR desplaza el precondicionamiento expresivo de una opción teóricamente atractiva pero a menudo impráctica a un régimen computacionalmente viable para el aprendizaje profundo a gran escala.

Los autores reconocen limitaciones, señalando que LLQR introduce sobrecarga de memoria y computación para almacenar y reajustar el precondicionador $U$. Sin embargo, argumentan que este costo es controlable mediante perillas de implementación (frecuencia de actualización, tamaño de fragmento) y está justificado por las ganancias de rendimiento y la capacidad de utilizar estructuras más ricas que las aproximaciones diagonales estándar.