Resumen Técnico: Teoría de Calendarios de Tasa de Aprendizaje Óptimos y Leyes de Escalado para un Modelo de Características Aleatorias

Enunciado del Problema

Entrenar modelos de aprendizaje profundo requiere seleccionar hiperparámetros como la tasa de aprendizaje (LR), el tamaño de lote y el horizonte de entrenamiento. Aunque el "ensayo y error" empírico es común, existe una falta de comprensión teórica sobre cómo los hiperparámetros óptimos escalan con el horizonte de entrenamiento $T$. Las estrategias existentes para transferir hiperparámetros entre tamaños de modelo a menudo fallan al transferirse entre diferentes horizontes de entrenamiento. Además, aunque en la práctica se utilizan diversos calendarios de LR (por ejemplo, decaimiento lineal, calentamiento-estabilidad-decaimiento, recocido coseno), las condiciones teóricas que dictan qué calendario es óptimo para una tarea dada permanecen poco claras. Este artículo aborda la necesidad de un marco principista para derivar calendarios de LR óptimos que dependan del horizonte de entrenamiento y de la estructura específica de la tarea de aprendizaje.

Metodología

Los autores emplean un modelo de características aleatorias de ley de potencias entrenado con Descenso de Gradiente Estocástico (SGD) para derivar calendarios de hiperparámetros óptimos. El modelo asume datos de entrada $x$ con una función objetivo $y(x) = w^* \cdot \psi(x) + \sigma_0 z$, donde las características $\psi(x)$ exhiben un decaimiento espectral de ley de potencias ($\lambda_k \sim k^{-b}$) y los pesos del profesor escalan como $(w^*_k)^2 \lambda_k \sim k^{-a}$.

La metodología central implica teoría de control óptimo:

1. Formulación de Dinámicas: La pérdida de prueba esperada $L_t$ se deriva como una función del calendario de tasa de aprendizaje $\{\eta_t\}$ y del calendario de tamaño de lote $\{m_t\}$. Las dinámicas de la pérdida están gobernadas por una relación de recurrencia lineal para los errores de modo $c_{t,k}$.
2. Objetivo de Optimización: El objetivo es minimizar la pérdida de prueba final $L_T$ en un horizonte fijo $T$. Los autores formulan tres problemas de optimización:
   • LR constante óptima.
   • Calendario de LR óptimo (variante en el tiempo $\eta_t$).
   • Calendario conjunto óptimo de LR y tamaño de lote.
3. Enfoques de Solución:
   • Numérico: Utilizando un método de disparo único (implementado en CasADi) para resolver el problema de control óptimo en tiempo discreto para $T$ finito.
   • Analítico: Derivando aproximaciones en tiempo continuo para $T$ grande utilizando el teorema de Noether y el análisis asintótico del funcional de pérdida. Esto implica minimizar un funcional de la tasa de aprendizaje integrada $\chi(t) = \int_t^T \eta(s) ds$.

Contribuciones y Resultados Clave

1. Identificación de Dos Fases Distintas

El análisis revela dos regímenes distintos para los calendarios de tasa de aprendizaje óptimos, determinados por la relación entre los exponentes espectrales $a$ (dificultad de la tarea) y $b$ (complejidad de las características):

• La Fase Fácil ($a > b$):

  • El calendario óptimo sigue una forma de decaimiento polinómico: $\eta^*_T(t) \sim T^{b/a - 1} f(t/T)$, donde $f(z) \propto (1-z)^{2b-1}$.
  • La tasa de aprendizaje decae gradualmente para equilibrar la reducción de sesgo y varianza a lo largo de la trayectoria de entrenamiento.
  • El exponente de escalado de la pérdida es $\zeta_T = (a-1)/a$, coincidiendo con la tasa minimax-óptima para tamaños de lote fijos.

• La Fase Difícil ($a < b$):

  • El calendario óptimo se asemeja a Calentamiento-Estabilidad-Decaimiento (WSD). Mantiene una tasa de aprendizaje máxima constante ($\eta_{max}$) durante la mayoría de los pasos de entrenamiento, seguida de una fase de recocido rápida sobre una fracción desvaneciente del tiempo total ($1 - t_s/T \sim T^{-(b-a)/(2b-1)}$).
  • Esta estrategia prioriza la supresión del sesgo mediante una LR máxima constante, comprimiendo la reducción de varianza en los pasos finales.
  • El exponente de escalado de la pérdida mejora a $\zeta_T = (a-1)/b$, superando a las estrategias de LR constante.

2. Superioridad sobre Puntos de Referencia

Los calendarios óptimos derivados superan estrictamente a los puntos de referencia estándar:

• LR Constante: El calendario óptimo logra mejores exponentes de escalado en ambas fases.
• Calendarios de Ley de Potencias: Incluso al optimizar los exponentes de los calendarios de ley de potencias ($\eta \sim T^{-\xi} t^{-\delta}$), recuperan asintóticamente el mismo comportamiento de escalado que la línea base de LR constante óptima, fallando en igualar el rendimiento de los calendarios óptimos dependientes del horizonte derivados mediante teoría de control.

3. Optimización Conjunta de Tasa de Aprendizaje y Tamaño de Lote

Los autores extienden el análisis para optimizar conjuntamente $\eta(t)$ y $m(t)$ bajo un presupuesto de datos fijo $B_{tot}$:

• Fase Fácil: Un "aumento de lote" (aumento del tamaño de lote con el tiempo) combinado con una tasa de aprendizaje máxima constante reduce el tiempo de reloj. La pérdida escala como $T^{-(a-1)/b}$, mejorando el exponente de reloj en comparación con tamaños de lote fijos, mientras que el exponente del presupuesto de muestras permanece en $(a-1)/a$.
• Fase Difícil: La optimización conjunta produce los mismos exponentes de escalado que el óptimo de lote fijo, con el tamaño de lote permaneciendo acotado.

4. Calendarios de Momento Óptimos

Para tareas con momento, los autores optimizan el parámetro de momento $\beta(t)$ junto con $\eta(t)$.

• En la Fase Fácil, el calendario de momento óptimo produce leyes de escalado similares a las del SGD óptimo.
• En la Fase Difícil, optimizar el calendario de momento proporciona una mejora en el exponente de escalado más allá de lo alcanzable con SGD solo. Las dinámicas de momento óptimas muestran una disminución transitoria inicial seguida de un aumento a medida que $t/T \to 1$.

5. Validación Empírica en Redes Profundas

Las fases teóricas se validan en puntos de referencia de aprendizaje profundo:

• ResNet en CIFAR-5M: Exhibe comportamiento de Fase Difícil. La tasa de aprendizaje base óptima para calendarios de recocido es aproximadamente constante a través de los horizontes, consistente con el perfil WSD.
• Transformadores estilo GPT-2 en C4: Exhiben comportamiento de Fase Fácil. La tasa de aprendizaje base óptima se desplaza (disminuye) con el horizonte $T$, incluso bajo recocido, consistente con el perfil de decaimiento polinómico.

Significado y Afirmaciones

El artículo afirma proporcionar una base teórica para comprender por qué los calendarios de tasa de aprendizaje varían entre diferentes tareas y horizontes de entrenamiento. Al utilizar la teoría de control óptimo en un modelo de características aleatorias resoluble, los autores demuestran que:

1. Dependencia del Horizonte: Las tasas de aprendizaje óptimas no son automáticamente transferibles entre horizontes de entrenamiento; el calendario debe diseñarse explícitamente en función de $T$.
2. Estructura de la Tarea: La "dificultad" de una tarea (caracterizada por los exponentes espectrales $a$ y $b$) dicta si una tarea cae en la fase fácil o difícil, determinando la forma del calendario óptimo (decaimiento polinómico vs. WSD).
3. Eficiencia: Los calendarios dependientes del horizonte y la optimización conjunta del tamaño de lote pueden mejorar significativamente las leyes de escalado y reducir el tiempo de entrenamiento de reloj en comparación con las heurísticas estándar.

Los autores reconocen limitaciones, señalando que su teoría actualmente se aplica a modelos lineales en el régimen perezoso/núcleo y aún no deriva analíticamente el exponente mejorado específico para el momento en la fase difícil ni da cuenta de optimizadores sofisticados como Adam o métodos precondicionados. Sin embargo, el trabajo ofrece una explicación principista de la efectividad de calendarios como WSD y sugiere que el ajuste empírico de las tasas de aprendizaje base a través de horizontes debe tener en cuenta las propiedades espectrales subyacentes del modelo y la tarea.