A Non-Monotone Preconditioned Trust-Region Method for Neural Network Training

Resumen Técnico: Un Método de Región de Confianza Precondicionado No Monótono para el Entrenamiento de Redes Neuronales

Planteamiento del Problema
El entrenamiento de redes neuronales (NN) profundas a escala presenta desafíos relacionados con el ajuste de hiperparámetros, la explotación de pipelines multi-dispositivo y el manejo de funciones objetivo no convexas y ruidosas. Si bien los métodos de primer orden como el Descenso de Gradiente Estocástico (SGD) y Adam son estándar, dependen en gran medida de la programación manual del tamaño de paso y del momento, y no aprovechan naturalmente el paralelismo de modelo donde una red se divide en bloques secuenciales a través de dispositivos. Por el contrario, los métodos de Descomposición de Dominio (DD) ofrecen una forma de particionar los parámetros en subdominios para actualizaciones paralelas, pero requieren estrategias de globalización robustas para garantizar la convergencia en entornos no convexos. Los enfoques anteriores, como la Estrategia de Región de Confianza Precondicionada Aditivamente (APTS), utilizan mecanismos de región de confianza (TR), pero a menudo sufren altas tasas de rechazo de pasos en el espacio grueso, lo que conduce a costosas recomputaciones de las direcciones de búsqueda.

Metodología: NAPTS
El artículo propone la Estrategia de Región de Confianza Precondicionada Aditivamente No Monótona (NAPTS), una extensión del marco APTS que integra principios de región de confianza no monótona (NTR). El método opera a través de tres fases distintas por iteración:

1. Actualizaciones Paralelas de Subdominios: El espacio de parámetros de la NN se particiona en subdominios no superpuestos $\{C_d\}_{d=1}^N$. En cada iteración $k$, un pase global hacia adelante y hacia atrás almacena en caché los valores de activación y los gradientes aguas abajo. Luego, cada subdominio realiza $\ell$ pasos locales de Adam con restricciones (CAdam) en paralelo. Estos pasos locales se calculan utilizando información de frontera almacenada en caché para aproximar los gradientes por bloques, asegurando que los subdominios operen de forma independiente mientras respetan la información global. Los pasos locales se elevan y combinan aditivamente para formar un paso de prueba global $s_k$.
2. Globalización No Monótona: A diferencia de los métodos TR estándar que requieren una disminución estrictamente monótona de la función objetivo, NAPTS emplea una ventana deslizante de las $\nu$ iteraciones exitosas más recientes. Define un índice de referencia $r(k)$ correspondiente al valor objetivo más alto dentro de esta ventana. La aceptación del paso se determina mediante una razón de acuerdo combinada $\rho_k = \max\{\rho^c_k, \rho^h_k\}$, donde:
   • $\rho^c_k$ mide el acuerdo entre la disminución real y la predicha para el paso actual.
   • $\rho^h_k$ mide el acuerdo relativo al valor de referencia histórico $f(\theta_{r(k)})$ y un término de historia acumulada $\sigma^h_k$.
     Este mecanismo permite la aceptación de pasos que aumentan temporalmente la función objetivo ($f(\theta_k + s_k) > f(\theta_k)$) si se alinean con tendencias de descenso a largo plazo, reduciendo así el rechazo de direcciones efectivas del espacio grueso.
3. Corrección y Suavizado: Si el paso de prueba es rechazado, se calcula un paso de corrección $c_k$ como una combinación convexa de un paso de descenso más pronunciado y la propuesta original, evitando el desperdicio de información de gradiente previamente calculada. Una actualización final NTR refina la iteración y ajusta el radio de la región de confianza $\Delta_k$.

Contribuciones Clave

• APTS No Monótono (NAPTS): La contribución principal es la integración de principios NTR en el marco APTS. Esto aborda específicamente el problema del rechazo agresivo de pasos en el espacio grueso en paisajes de entrenamiento de NN ruidosos y no convexos.
• No Monotonía Controlada: El artículo introduce un criterio de aceptación con ventana que permite aumentos controlados de la función objetivo. Esto evita la aceptación "indiscriminada" de pasos, permitiendo al mismo tiempo que el algoritmo persista con direcciones del espacio grueso que podrían ser temporalmente no mejoradoras pero globalmente beneficiosas.
• Paralelismo Eficiente: El método mantiene la estructura de Schwarz aditiva, permitiendo el entrenamiento paralelo de subdominios sin comunicación entre subdominios durante las iteraciones locales, requiriendo solo sincronización global para el almacenamiento en caché de gradientes y la agregación de pasos.

Resultados
Los experimentos se realizaron en el conjunto de datos CIFAR-10 utilizando una red neuronal convolucional (1.2M de parámetros) entrenada en cuatro GPUs Nvidia A100.

• Rendimiento: NAPTS logró un tiempo de entrenamiento de 115.88 segundos por época, lo que representa una reducción del 30% en el tiempo de CPU en comparación con el APTS estándar (169.14s/época) y una mejora adicional sobre la variante de aceptación agresiva APTS$_A$ (152.39s/época).
• Tasas de Rechazo: El método redujo significativamente el número de pasos rechazados. NAPTS registró aproximadamente 1,921 iteraciones rechazadas en total, en comparación con 13,972 para APTS y 6,648 para APTS$_A$. Esto indica que el mecanismo no monótono retiene con éxito propuestas efectivas del espacio grueso que de otro modo serían descartadas.
• Precisión: Mientras que los métodos TR y NTR de referencia lucharon por alcanzar una alta precisión de validación, las variantes precondicionadas (APTS y NAPTS) lograron precisiones de validación sustancialmente más altas y pérdidas de entrenamiento más bajas, lo que sugiere una optimización mejorada sin comprometer la generalización.

Significado
El artículo afirma que NAPTS ofrece una estrategia de globalización robusta específicamente adecuada para la naturaleza no convexa y ruidosa del entrenamiento de redes neuronales a gran escala. Al relajar el requisito de monotonía estricta, el método evita el costo computacional de rechazar y recomputar repetidamente las direcciones de búsqueda. Los autores postulan que este enfoque aprovecha eficazmente la descomposición de dominio para el paralelismo de modelo, mientras proporciona una adaptación automática del tamaño de paso a través de regiones de confianza, reduciendo la necesidad de ajuste manual de hiperparámetros asociado con los métodos de primer orden. Los resultados sugieren que NAPTS preserva la precisión de los métodos precondicionados mientras mejora sustancialmente la eficiencia computacional.