Learning Rate Warmup/es
| Article | |
|---|---|
| Topic area | Optimization |
| Prerequisites | Stochastic Gradient Descent, Learning Rate Schedule, Adam |
Resumen
El warmup de la tasa de aprendizaje es una clase de programaciones de la tasa de aprendizaje en las que la tasa de aprendizaje se fija inicialmente en un valor pequeño y se incrementa gradualmente hacia una tasa pico objetivo durante las primeras cientos o miles de iteraciones del entrenamiento, antes de que comience cualquier fase de decaimiento estándar. La técnica apareció de modo informal en la práctica del aprendizaje profundo a mediados de la década de 2010, recibió una receta formal de Goyal y colaboradores para el entrenamiento de ImageNet con lotes grandes en 2017, y se convirtió en un componente casi universal de los pipelines de entrenamiento de los modelos Transformer tras su uso en los artículos originales del transformer y de BERT.[1][2]
La motivación es que las primerísimas actualizaciones del entrenamiento son inusualmente arriesgadas: los parámetros se encuentran en su inicialización aleatoria, los gradientes son grandes y están poco correlacionados con la dirección verdadera de descenso, y los optimizadores adaptativos como Adam aún no han acumulado estimaciones fiables del segundo momento. Una tasa de aprendizaje pequeña durante este período evita que estas actualizaciones tempranas empujen al iterado hacia una región desfavorable del paisaje de pérdida, después de lo cual una tasa más alta puede usarse con seguridad durante el grueso del entrenamiento.
Warmup lineal
La forma más común es el warmup lineal. Dada una tasa de aprendizaje pico objetivo $ \eta_{\max} $ y una longitud de warmup $ T_w $ medida en iteraciones (o a veces épocas), la tasa en la iteración $ t $ es
$ {\displaystyle \eta_t = \eta_{\max} \cdot \min\!\left(1, \frac{t}{T_w}\right).} $
Tras el paso $ T_w $ la programación continúa con cualquier decaimiento previsto: recocido del coseno, decaimiento de raíz cuadrada inversa, decaimiento por escalones o una tasa constante. Una longitud de warmup típica está entre el 1 % y el 10 % del entrenamiento total, siendo habituales entre 500 y 4000 iteraciones en el preentrenamiento de modelos de lenguaje.
Un primo cercano, el warmup constante, utiliza una tasa fija pequeña $ \eta_0 \ll \eta_{\max} $ durante los primeros $ T_w $ pasos y luego salta a $ \eta_{\max} $. Goyal y col. observaron que el warmup constante es frágil cuando la diferencia entre $ \eta_0 $ y $ \eta_{\max} $ es grande, y recomendaron el warmup gradual (la forma lineal anterior) como opción predeterminada más fiable para el entrenamiento con lotes grandes.
Warmup de raíz cuadrada inversa en transformers
El artículo original del transformer combinaba un warmup lineal con un decaimiento posterior de raíz cuadrada inversa, produciendo una programación con un único pico suavemente diferenciable:
$ {\displaystyle \eta_t = d_{\text{model}}^{-1/2} \cdot \min\!\left( t^{-1/2},\, t \cdot T_w^{-3/2} \right),} $
donde $ d_{\text{model}} $ es la dimensión del modelo transformer y $ T_w $ es la longitud del warmup (típicamente 4000 pasos). Las dos ramas del mínimo se cruzan en $ t = T_w $, donde la programación pasa de una fase de warmup que crece linealmente a un decaimiento de raíz cuadrada inversa. Esta receta fue heredada, con pequeñas modificaciones, por muchos sistemas posteriores secuencia a secuencia antes de que el warmup lineal seguido de decaimiento del coseno le ganara terreno en el preentrenamiento de modelos de lenguaje autorregresivos.
Por qué ayuda el warmup
Se han ofrecido varias explicaciones complementarias, ninguna de ellas plenamente asentada. Liu y colaboradores dieron el análisis más influyente: en Adam y sus variantes la estimación adaptativa del segundo momento $ v_t $ presenta una varianza muy alta durante las primeras actualizaciones, lo que a su vez hace que la tasa de aprendizaje efectiva por parámetro $ \eta / \sqrt{v_t + \epsilon} $ sea poco fiable.[3] Un período corto de warmup da tiempo a que $ v_t $ se estabilice antes de que el tamaño global del paso sea grande. El mismo artículo propuso RAdam, una variante de Adam que incorpora un término de rectificación de varianza y pretende hacer innecesario el warmup explícito; en la práctica, tanto RAdam como el warmup explícito siguen utilizándose.
Una segunda línea de razonamiento concierne al paisaje de pérdida cerca de la inicialización. Los pesos aleatorios producen activaciones y gradientes alejados de la trayectoria final de entrenamiento; pasos tempranos grandes pueden amplificar los desequilibrios entre capas (por ejemplo, entre las subcapas de atención y de avance directo en los transformers) y desencadenar inestabilidades numéricas como la explosión de los logits previos al softmax. El warmup limita el daño por paso durante este transitorio.
Una tercera explicación, específica del entrenamiento con lotes grandes, invoca la escala del ruido del gradiente. Con lotes muy grandes el gradiente estocástico es próximo al gradiente verdadero, por lo que el optimizador se comporta de modo casi determinista; una tasa de aprendizaje inicial alta puede entonces hacer que el iterado se pase de largo de manera persistente. Goyal y col. combinaron el warmup con la regla de escalado lineal (tasa pico proporcional al tamaño del lote) para extender el SGD con momento a lotes de 8192 imágenes en ImageNet sin pérdida de precisión.
Variantes y combinaciones
Más allá de las formas lineal y de raíz cuadrada inversa, en la práctica se utilizan varias otras formas.
- Warmup con coseno: la tasa sigue un medio coseno desde $ 0 $ hasta $ \eta_{\max} $, ofreciendo un arranque más suave que la rampa lineal. Esta forma se refleja después con un decaimiento del coseno, dando lugar a una única programación con forma de campana.
- Warmup exponencial: $ \eta_t = \eta_{\max} \cdot (1 - e^{-t / \tau}) $ para alguna constante de tiempo $ \tau $. Menos común en aprendizaje profundo, pero aparece en aprendizaje por refuerzo y metaaprendizaje.
- Warmup por capas: en LAMB y los optimizadores adaptativos relacionados de adaptación por capas, el mecanismo de razón de confianza produce un warmup implícito cuya longitud difiere por capa. A menudo se mantiene un warmup global explícito por encima de éste.
El warmup combina de forma natural con la mayoría de las programaciones de decaimiento. El patrón dominante en el preentrenamiento moderno de LLM es un warmup lineal seguido de decaimiento del coseno hasta una fracción (a menudo el 10 %) de la tasa pico. El warmup también se vuelve a aplicar habitualmente al inicio del ajuste fino, por la misma teoría: una cabeza de salida recién inicializada y un estado del optimizador reiniciado por el framework se benefician ambos de una breve fase de tasa baja.
Consideraciones prácticas
El warmup añade en principio dos hiperparámetros ($ T_w $ y la tasa inicial, a menudo cero), pero la mayoría de las recetas fijan la tasa inicial en cero o en una constante pequeña como $ 10^{-7} $ y ajustan únicamente $ T_w $. Guía empírica:
- Para el preentrenamiento de Transformer, fije $ T_w $ entre el 1 % y el 5 % del total de pasos de entrenamiento. Valores menores al 0,5 % producen con frecuencia picos de pérdida, especialmente en modelos muy anchos.
- Para el ajuste fino supervisado, son típicos entre 50 y 500 pasos de warmup; la longitud óptima escala con el tamaño del conjunto de datos.
- Aumente $ T_w $ al escalar el tamaño del lote o el tamaño del modelo; ambas intervenciones amplifican las inestabilidades de comienzo de entrenamiento que el warmup aborda.
- Cuando se utiliza precisión mixta o FSDP / paralelismo de pipeline, un warmup más largo ayuda a tolerar el desajuste inicial del factor de escalado de la pérdida y de las normas de capa.
El warmup interactúa con el decaimiento de peso en AdamW: como el decaimiento de peso desacoplado es independiente del gradiente, la fuerza efectiva del decaimiento durante la fase de warmup no se ve afectada, pero el tirón relativo hacia cero es mayor porque las actualizaciones del gradiente son más pequeñas. Algunas recetas, por tanto, escalan el decaimiento de peso por el mismo factor utilizado para la tasa de aprendizaje durante el warmup, aunque esto no es universal.
Limitaciones y alternativas
El warmup es empíricamente robusto pero está poco explicado teóricamente. Su eficacia depende de detalles que no están plenamente caracterizados: el tipo de optimizador, la arquitectura del modelo, la precisión aritmética y el esquema de inicialización. RAdam intenta absorber la fase de warmup en el propio optimizador; Adafactor y Lion, de modo similar, reducen la necesidad de warmup en algunos regímenes pero no la eliminan. Las correcciones arquitectónicas como los transformers Pre-Norm y la inicialización cuidadosa (Fixup, DeepNorm, ReZero) reducen la longitud de warmup necesaria para la estabilidad pero rara vez hacen que sea seguro omitirlo a gran escala.
En ausencia de una teoría única en forma cerrada, la recomendación que ha emergido a lo largo de varios grupos de investigación es mantener una fase de warmup explícita en cualquier pipeline de entrenamiento que utilice un optimizador adaptativo, lotes grandes o arquitecturas de tipo transformer, y tratar su longitud como un hiperparámetro digno de un pequeño barrido específico cuando se entrena por primera vez una nueva clase de modelo.