Efficient Construction of Model Family through Progressive Training Using Model Expansion

Título: Construcción Eficiente de una Familia de Modelos mediante Entrenamiento Progresivo usando Expansión de Modelos

1. El Problema

La industria de los Modelos de Lenguaje Grandes (LLM) ha estandarizado la entrega de "familias de modelos" (conjuntos de modelos con la misma arquitectura pero diferentes tamaños de parámetros, ej. 1B, 7B, 70B) para adaptarse a diversas restricciones computacionales.

• Ineficiencia Actual: El enfoque convencional consiste en entrenar cada modelo de la familia de forma independiente desde cero (inicialización aleatoria).
• Costo Computacional: Este método genera un costo computacional aditivo. Entrenar una familia de modelos requiere sumar los costos de cada modelo individual, lo que implica un gasto masivo de recursos (miles de días de GPU).
• Objetivo: Reducir significativamente el costo total de entrenamiento de una familia completa de modelos sin sacrificar el rendimiento de los modelos individuales, y preferiblemente mejorar la consistencia entre ellos.

2. Metodología: Entrenamiento Progresivo

Los autores proponen un método llamado Entrenamiento Progresivo, que utiliza la expansión de modelos para construir una familia de manera incremental.

• Concepto Central: En lugar de entrenar cada modelo desde cero, se entrena el modelo más pequeño y, a medida que se avanza, se expanden sus parámetros para inicializar el siguiente modelo más grande.
• Proceso:
  1. Se entrena el modelo inicial más pequeño ($X_1$) desde cero.
  2. Para el siguiente modelo ($X_{i+1}$), se utiliza un método de expansión de modelos (en este caso, bert2BERT) para inicializar los parámetros $\theta_{i+1}$ basándose en los parámetros entrenados del modelo anterior $\theta_i$.
  3. El modelo expandido se entrena con un nuevo conjunto de tokens.
• Asignación de Tokens (Restricción de Costo): El objetivo es que el costo total de entrenar toda la familia sea equivalente al costo de entrenar solo el modelo más grande ($X_n$) desde cero.
  • Se calcula la asignación de tokens ($T^{prog}_i$) para cada etapa restando el costo computacional ya gastado en las etapas anteriores del costo total permitido.
  • Esto permite obtener modelos de 1B, 2B, 4B y 8B con un costo total igual al de entrenar solo el de 8B.
• Ajuste de Tasa de Aprendizaje: Se introduce una estrategia donde la tasa de aprendizaje máxima se ajusta según el tamaño del modelo (mayor para modelos pequeños, menor para grandes) para equilibrar la eficiencia de optimización y la estabilidad del entrenamiento.

3. Contribuciones Clave

• Reducción de Costos: Demostración empírica de que el entrenamiento progresivo reduce el costo computacional total en aproximadamente un 25% en comparación con el entrenamiento independiente, manteniendo un rendimiento comparable o superior.
• Mejora de Rendimiento: Al ajustar la tasa de aprendizaje según el tamaño del modelo, el método supera a los modelos entrenados independientemente en diversas métricas de evaluación.
• Consistencia de Comportamiento: Se identifica que los modelos entrenados progresivamente exhiben una mayor consistencia en sus distribuciones de probabilidad entre diferentes tamaños, lo cual es crucial para técnicas como la decodificación especulativa.
• Validación de Datos Únicos: Se demuestra que las mejoras no se deben a un mayor volumen de datos únicos, sino a la eficiencia de la inicialización y el entrenamiento progresivo (validado mediante un experimento de control con datos fijos).

4. Resultados Experimentales

Los experimentos se realizaron en una familia de modelos de 1B a 8B parámetros utilizando la arquitectura Llama y datos de FineWeb-Edu.

• Eficiencia Computacional:
  • Entrenar independientemente la familia [1B, 2B, 4B, 8B] bajo la ley de Chinchilla (2x) requeriría ~20.4 ZFLOPs.
  • El enfoque progresivo redujo esto a ~15.4 ZFLOPs (costo equivalente solo al modelo de 8B), logrando una reducción del 25%.
  • Con el ajuste de la tasa de aprendizaje, la reducción de FLOPs para alcanzar el mismo loss fue del 31%.
• Rendimiento (Perplejidad y Tareas Downstream):
  • Los modelos progresivos igualaron o superaron a los modelos independientes en perplejidad y tareas como razonamiento común (WinoGrande, PIQA) y respuesta a preguntas (ARC, OBQA).
  • El modelo de 8B progresivo mostró mejoras notables en perplejidad y precisión en tareas de lenguaje.
• Consistencia (Divergencia KL):
  • Se midió la Divergencia Kullback-Leibler (KL) entre modelos adyacentes. Los modelos progresivos mostraron una menor divergencia KL (ej. 0.3378 para 4B→8B vs 0.4584 en independientes), indicando que comparten mecanismos internos más similares.
• Decodificación Especulativa:
  • La mayor consistencia se tradujo en beneficios prácticos. Al usar un modelo "borrador" (4B) y un "generador" (8B) ambos entrenados progresivamente, la tasa de aceptación de tokens aumentó al 93.2% (vs 87.2% con modelos independientes), reduciendo el tiempo de inferencia en un 14.8%.

5. Significado e Impacto

• Accesibilidad: Este método hace viable la creación de familias de modelos coherentes para organizaciones con recursos limitados, permitiendo obtener modelos de múltiples escalas con el presupuesto de uno solo.
• Despliegue Flexible: La consistencia en el comportamiento entre tamaños facilita la adaptación dinámica de modelos (escalado hacia arriba o abajo) sin necesidad de reentrenar usuarios o revalidar sistemas extensivamente.
• Eficiencia en Inferencia: La alta consistencia entre modelos de diferentes tamaños mejora significativamente técnicas de aceleración como la decodificación especulativa, haciendo que la inferencia sea más rápida y eficiente en producción.
• Paradigma de Entrenamiento: Cambia el enfoque de "entrenar cada modelo desde cero" a un enfoque de "crecimiento incremental", aprovechando el conocimiento adquirido en modelos más pequeños para acelerar y estabilizar el entrenamiento de modelos más grandes.

En resumen, el artículo presenta una solución robusta y eficiente para la construcción de familias de LLM, logrando ahorros masivos en costos computacionales mientras mejora la calidad y la coherencia interna de los modelos resultantes.