ByteDance Seed propone AIR, un calendario de learning rate: la pérdida final es 0,01 menor que con cosine

Investigadores de ByteDance Seed (Qifan Yu, Xinyu Ma, Xingyan Bin, Di He) proponen AIR, un calendario de tasa de aprendizaje. En el preentrenamiento de un modelo Mixture-of-Experts (MoE) con 128 mil millones de tokens, AIR redujo la pérdida final sobre los datos de entrenamiento (training loss) en aproximadamente 0,01 respecto al decaimiento coseno (cosine). La comparación se hizo con la misma suma de pasos a lo largo de todo el entrenamiento. Aquí, la tasa de aprendizaje (learning rate, LR) es el tamaño del paso con el que el optimizador modifica los pesos del modelo. Suele reducirse durante el entrenamiento, casi siempre siguiendo una curva coseno.

Por qué se empezó a hablar de la normalización de pesos

Para controlar el tamaño de los pesos se suele usar weight decay (WD), que en cada paso los empuja ligeramente hacia cero. Recientemente se ha empezado a considerar como alternativa la weight normalization (WN), un entrenamiento con pesos normalizados. Varios trabajos han reportado una mejora notable del loss en el preentrenamiento.

Los autores comprobaron de dónde procede esa mejora y llegaron a otra explicación. Con WN, el LR nominal se traduce en un tamaño real de actualización distinto del que resulta con WD. Así se obtiene una trayectoria propia del learning rate efectivo (effective learning rate, ELR). Los investigadores ajustaron el ELR de ambos métodos para que las trayectorias coincidieran. Con ello, las curvas de loss de WN y WD casi se fundieron. Por tanto, la mejora la daba la forma del ELR, no la normalización por sí misma.

De aquí se deduce una conclusión práctica: el calendario debe construirse para el ELR, y no solo para el LR nominal.

Cómo funciona AIR

El nombre significa Affine Inverse-square-Root decay, es decir, decaimiento afín según la raíz cuadrada inversa. El calendario tiene dos partes.

• Base. El paso disminuye en proporción a 1/√t, donde t es el número del paso de entrenamiento. Los autores derivaron esta forma de manera teórica. • Corrección. Una transformación lineal (afín). El entrenamiento dura un tiempo finito y un decaimiento 1/√t puro no da los valores necesarios al principio ni al final. La corrección permite fijar de antemano el LR máximo y el LR final.

Igual que con cosine, el usuario indica dos puntos extremos: el máximo y el final.

Hay también una observación sobre la fase tardía. Si varias ejecuciones comparten la misma cola del calendario de ELR, esa cola atenúa la influencia de las elecciones de LR anteriores. Los autores consideran que la planificación tardía del ELR es una de las principales fuentes de la mejora.

Qué mostró el experimento

La prueba se realizó en el preentrenamiento de un modelo MoE con 128 mil millones de tokens. En un MoE, para cada token trabaja solo una parte de los bloques de expertos. AIR se comparó con cosine y con otros calendarios candidatos. Los resultados:

• El training loss es menor con AIR. La pérdida final es aproximadamente 0,01 menor que con cosine, con la misma suma de LR a lo largo de todo el entrenamiento. • AIR también es mejor en calidad en tareas aplicadas (downstream). • El calendario es más robusto frente a la elección de hiperparámetros.

En un experimento aparte, la misma trayectoria de ELR se aplicó al weight decay habitual. La dinámica de entrenamiento se reprodujo. Los autores concluyen que el beneficio no se limita al entrenamiento con norma de pesos fija y señalan el ELR como el objeto que conviene planificar.

Qué implica esto en la práctica

Una diferencia de 0,01 en el loss es pequeña, pero se logró cambiando un solo calendario de paso. El trabajo ofrece tres pautas prácticas:

• comprobar cómo su esquema de pesos y su optimizador convierten el LR nominal en efectivo; • probar el calendario AIR, en el que solo se fijan el LR máximo y el final; • prestar más atención al final del entrenamiento, porque una cola común del calendario atenúa la influencia de los ajustes anteriores.

En el resumen (abstract) del artículo, los experimentos se describen con preentrenamiento de un MoE con 128 mil millones de tokens. No informa sobre otros tamaños de modelo ni volúmenes de datos.

El artículo se publicó en arXiv con el número 2610.40362 y entró en la selección de Hugging Face Daily Papers.