Weight Decay Improves Language Model Plasticity

Resumen Técnico: El decaimiento de pesos mejora la plasticidad de los modelos de lenguaje

Planteamiento del problema
Las prácticas actuales para el entrenamiento de Grandes Modelos de Lenguaje (LLM) suelen desacoplar el preentrenamiento del post-entrenamiento (ajuste fino, alineación). La optimización de hiperparámetros y las leyes de escala durante el preentrenamiento se evalúan predominantemente basándose en la pérdida de validación del modelo base (entropía cruzada), bajo la suposición de que una menor pérdida de preentrenamiento se correlaciona con un modelo más capaz en tareas posteriores. Sin embargo, este enfoque pasa por alto la plasticidad del modelo: la capacidad de un modelo entrenado para adaptarse con éxito a nuevos datos y tareas mediante un entrenamiento posterior sin reinicialización. El artículo investiga si la optimización de los hiperparámetros de preentrenamiento únicamente para el rendimiento del preentrenamiento compromete inadvertidamente la capacidad de adaptación del modelo durante el ajuste fino posterior. Específicamente, los autores se centran en el decaimiento de pesos (weight decay), un hiperparámetro de regularización estándar, para determinar si su valor óptimo para minimizar la pérdida de preentrenamiento difiere de su valor óptimo para maximizar la adaptabilidad posterior.

Metodología
Los autores llevan a cabo un estudio empírico sistemático a través de dos familias de modelos (Llama-2 y OLMo-2) y múltiples tamaños de modelo (hasta 4B de parámetros). El diseño experimental varía el hiperparámetro de decaimiento de pesos ($\lambda$) durante el preentrenamiento mientras controla otros factores.

1. Regímenes de entrenamiento: Los modelos se entrenan en dos regímenes distintos:
   • Computacionalmente óptimo: 20 tokens por parámetro (TPP).
   • Sobreentrenado: 140 TPP.
2. Evaluación de ajuste fino: Los modelos preentrenados se someten a un ajuste fino en una suite diversa de tareas de seguimiento, que incluyen:
   • Seis conjuntos de datos de Cadena de Pensamiento (CoT) (p. ej., MetaMathQA, MMLUProCoT).
   • Cinco pruebas de comprensión del lenguaje y razonamiento de sentido común (p. ej., HellaSwag, WinoGrande).
   • Una tarea de alineación de seguridad.
3. Métricas: El rendimiento posterior se evalúa utilizando múltiples métricas, incluyendo Greedy (Pass@1), Pass@16, Voto de Mayoría (Maj@16), puntuación de Modelo de Recompensa (RM@16) y Ratio de Corrección.
4. Análisis Mecanicista: Para comprender por qué el decaimiento de pesos afecta la plasticidad, los autores analizan:
   • Sondeo Lineal (Linear Probing): Medición de la separabilidad lineal de las representaciones internas (clasificación de sentimiento y tema).
   • Rango de Atención: Análisis del pseudo-rango de las matrices de proyección de consulta-clave ($W_{QK}$) y de valor ($W_{VP}$).
   • Sobreajuste (Overfitting): Cálculo de la brecha entre la pérdida de entrenamiento y la pérdida de validación (Brecha Entreno-Validación).

Contribuciones Clave

1. El decaimiento de pesos como impulsor de la plasticidad: El artículo demuestra que el decaimiento de pesos es un factor crítico en la configuración de la plasticidad del modelo. Contrario a los valores predeterminados estándar, valores de decaimiento de pesos más grandes durante el preentrenamiento conducen consistentemente a un mejor rendimiento posterior tras el ajuste fino.
2. Desacoplamiento de objetivos: El estudio proporciona evidencia de que el valor del hiperparámetro que minimiza la pérdida de validación del preentrenamiento no es necesariamente el valor que maximiza el rendimiento posterior. En varios regímenes, los modelos con una mayor pérdida de preentrenamiento (debido a un mayor decaimiento de pesos) superan a los modelos con una menor pérdida de preentrenamiento tras el ajuste fino.
3. Perspectivas mecanicistas: Los autores identifican tres mecanismos específicos mediante los cuales el decaimiento de pesos mejora la plasticidad:
   • Fomenta representaciones linealmente separables, haciendo que las características internas sean más accesibles para las tareas posteriores.
   • Regulariza las matrices de atención, específicamente reduciendo el rango del producto consulta-clave ($W_{QK}$), lo que conduce a estructuras más comprimidas y robustas.
   • Reduce el sobreajuste en los datos de preentrenamiento, logrando efectivamente "olvidar" el ruido específico del entrenamiento y preservando la capacidad del modelo para aprender nueva información.

Resultos Clave

• Desplazamiento del Decaimiento de Pesos Óptimo: Para el entrenamiento computacionalmente óptimo (20 TPP), el decaimiento de pesos óptimo para el rendimiento posterior es significativamente mayor que el estándar de 0.1. Por ejemplo, se encontraron valores óptimos de 0.5–1.0 para los modelos Llama-2 y OLMo-2 en este régimen. En el régimen sobreentrenado (140 TPP), el valor óptimo para el rendimiento posterior (0.3) sigue siendo superior al valor que minimiza la pérdida de preentrenamiento (0.1).
• Compensaciones Contraintuitivas: El artículo destaca una compensación donde modelos base "peores" (mayor pérdida de preentrenamiento) pueden convertirse en modelos finales "mejores" tras el ajuste fino. Por ejemplo, un modelo OLMo-2-1B-140x preentrenado con $\lambda=0.3$ o $1.0$ tuvo una mayor pérdida de preentrenamiento que la línea base de $\lambda=0.1$, pero logró una precisión en tareas posteriores notablemente mejor.
• Robustez: Estos hallazgos se mantienen a través de diferentes familias de modelos, tamaños e incluso cuando otros hiperparámetros (tasa de aprendizaje, tamaño de lote) varían durante tanto el preentrenamiento como el ajuste fino.
• Correlaciones Mecanicistas:
  • Un mayor decaimiento de pesos se correlaciona con una mayor precisión de sondeo lineal, indicando representaciones más estructuradas.
  • Un mayor decaimiento de pesos se correlaciona con un menor pseudo-rango en las matrices de atención (particularmente en $W_{QK}$), sugiriendo un movimiento hacia hipótesis de bajo rango y robustas.
  • Un mayor decaimiento de pesos se correlaciona con una menor Brecha Entreno-Validación, indicando un menor sobreajuste.

Significancia y Afirmaciones
El artículo argumenta que el campo debe reconsiderar la dependencia exclusiva de la pérdida de validación del preentrenamiento como la única métrica para la optimización de hiperparámetros. Los autores afirman que la plasticidad del modelo es una propiedad distinta y crucial que a menudo se sacrifica cuando se optimiza puramente para la pérdida de preentrenamiento.

La importancia de este trabajo radica en:

1. Reevaluación de los Valores Predeterminados: Sugiere que las elecciones estándar de hiperparámetros (p. ej., $\lambda=0.1$) pueden ser subóptimas para el ciclo de vida completo del desarrollo de LLM, particularmente cuando la adaptabilidad posterior es el objetivo final.
2. Limitaciones de la Entropía Cruzada: Destaca que la pérdida de entropía cruzada es un proxy imperfecto para la capacidad en tareas posteriores, ya que modelos con pérdidas similares pueden exhibir una plasticidad muy distinta.
3. Rol Multifacético de los Hiperparámetros: Subraya que un único hiperparámetro como el decaimiento de pesos desempeña un papel complejo y multifacético en la configuración del comportamiento del modelo, influyendo no solo en la generalización sino también en las propiedades estructurales (linealidad, rango) que facilitan el aprendizaje futuro.

Los autores mantienen la modestia, señalando que aunque sus resultados son robustos en las escalas probadas, los valores óptimos específicos pueden desplazarse en escenarios de escala extremadamente grande o fuertemente sobreentrenados más allá de su alcance experimental actual. Hacen un llamado a trabajos futuros para investigar más a fondo la compensación entre estabilidad y plasticidad, así como la aplicabilidad de estos hallazgos a otras modalidades de modelos fundacionales.