Mixture of Training: Recombining Small-Scale Scaffolded Pretraining Runs into a Larger Language Model

Resumen Técnico: Mezcla de Entrenamiento (MoT)

Planteamiento del Problema
El preentrenamiento de modelos de lenguaje de gran tamaño (LLM) se organiza actualmente como una única ejecución de optimización monolítica y de extremo a extremo. Esta arquitectura crea un cuello de botella de escalabilidad donde todas las capas deben entrenarse simultáneamente, lo que hace que el proceso sea frágil ante fallos (donde un solo error detiene toda la ejecución) e inflexible para la investigación iterativa (donde las mejoras requieren reiniciar o continuar un sistema masivo acoplado). El artículo plantea si el preentrenamiento puede descomponerse en trabajos más pequeños, entrenables de forma independiente, que puedan ser posteriormente recompuestos en un modelo coherente y más grande sin sacrificar la calidad.

Metodología: Mezcla de Entrenamiento (MoT)
Los autores introducen la Mezcla de Entrenamiento (MoT), un procedimiento de preentrenamiento modular y andamiado diseñado para particionar un Transformer objetivo en bloques de capas contiguas. El método opera en tres etapas:

1. Etapa 0 (Preparación): Se selecciona o prepara un alineador ($A$) preentrenado. Tanto el Transformer objetivo ($F$) como el alineador se particionan en $K$ bloques de capas contiguas. El alineador es compatible en forma con el objetivo, compartiendo la anchura global, la dimensionalidad de las cabezas de atención, la anchura de la red feed-forward, los embeddings de tokens y las cabezas de salida.
2. Etapa 1 (Entrenamiento Andamiado): El modelo objetivo se descompone en $K$ submodelos ($f_1, \dots, f_K$). Para cada submodelo $f_i$, se construye una red andamiada $S_i$ insertando $f_i$ en la estructura del alineador congelado:
   $$S_i = a_K \circ \dots \circ a_{i+1} \circ f_i \circ a_{i-1} \circ \dots \circ a_1$$
   En esta configuración, solo el bloque objetivo $f_i$ es entrenable; todos los fragmentos del alineador están congelados. Cada andamio se optimiza de forma independiente utilizando la pérdida estándar de predicción del siguiente token (next-token prediction loss). Crucialmente, no hay intercambio de gradientes entre los bloques objetivo durante esta fase, pero todos aprenden dentro del contexto representacional estable proporcionado por el alineador compartido.
3. Etapa 2 (Recomposición y Adaptación): El alineador se descarta. Los bloques objetivo entrenados se recomponen en el modelo completo $\hat{F} = f_K \circ \dots \circ f_1$. Este modelo de "composición en frío" (cold-composed) puede evaluarse directamente, o puede someterse a una breve pasada de adaptación de extremo a extremo para resolver cualquier desajuste representacional restante.

Contribuciones Clave

1. Entrenamiento Modular Andamiado: Un procedimiento que entrena cortes de profundidad objetivo de forma independiente mientras preserva una interfaz representacional compartida mediante un alineador congelado, permitiendo la recomposición sin necesidad de unir capas post-hoc.
2. Prueba de Mecanismo: Un estudio a pequeña escala que demuestra que los cortes de profundidad entrenados de forma independiente pueden recomponerse con éxito en un modelo de lenguaje de 1.3 mil millones de parámetros utilizable.
3. Contabilidad de Cómputo y Programación: Un análisis detallado de los compromisos entre cómputo (FLOPs), exposición de tokens y la trayectoria crítica ideal equivalente a capas. Esto incluye el tratamiento explícito del coste del alineador, distinguiendo entre costes de entrenamiento "directos" y costes "totalmente cargados" (asignando el coste completo del alineador a una sola ejecución) frente a costes amortizados a través de múltiples ejecuciones.
4. Estudios de Ablación: Evidencia empírica que muestra la necesidad del alineador, el impacto de los flujos de datos disjuntos y el compromiso entre calidad y eficiencia asociado al aumento del número de divisiones ($K$).

Resultados Experimentales
Los autores evaluaron MoT en un modelo de estilo Gemma de 12 capas y 1.3 mil millones de parámetros entrenado con el conjunto de datos C4.

• Paridad de Calidad: Un programa de "paridad de calidad", que extendió el entrenamiento del submodelo e incluyó una fase de adaptación de 30k pasos, logró la misma perplejidad (15.0) que la línea base monolítica.
• Compromisos de Cómputo:
  • El programa de paridad de calidad procesó más tokens agregados (47.1B frente a 33.6B) y tuvo una trayectoria crítica ideal más corta (potencial de aceleración de 1.7x), pero requirió más FLOPs totales (285.0 EFLOPs frente a 268.4 EFLOPs) cuando el coste del alineador se cargó totalmente a una sola ejecución.
  • Sin embargo, si el alineador se reutiliza en $R$ ejecuciones independientes, el coste efectivo por ejecución disminuye. El artículo señala que para $R \ge 3$, el coste efectivo cae por debajo de la línea base monolítica.
  • Un programa de menor cómputo (MoT + 15k de adaptación) alcanzó una perplejidad de 15.9 con 189.4 EFLOPs (totalmente cargados), lo cual está por debajo del presupuesto de la línea base, aunque no se proporcionó una línea base monolítica de igual cómputo para una comparación directa en este presupuesto específico.
• Ablaciones:
  • Necesidad del Alineador: Sin el alineador, la calidad de la composición en frío colapsó (PPL 38.9), confirmando que el alineador evita el desajuste de la interfaz.
  • Flujos de Datos: Los flujos de datos disjuntos para los submodelos mejoraron la calidad de la composición en frío (PPL 19.3) en comparación con los datos compartidos (PPL 20.3) cuando había un alineador presente.
  • Conteo de Divisiones: El aumento de las divisiones de $K=2$ a $K=4$ redujo el cómputo pero degradó la calidad de la composición en frío (PPL 24.8), resaltando un compromiso entre el paralelismo y la complejidad de la interfaz.

Significancia y Reivindicaciones
El artículo posiciona a MoT no como un reemplazo general para el preentrenamiento monolítico, sino como un marco a pequeña escala para estudiar si los sub-procesos andamiados pueden actuar como unidades de entrenamiento reutilizables.

Los autores sostienen tres conclusiones principales:

1. Los submodelos entrenados de forma independiente pueden recomponerse en un modelo de lenguaje coherente si se entrenan dentro de un andamio de alineador compartido; sin él, la composición en frío falla.
2. El desajuste entre los bloques entrenados de forma independiente es mayormente recuperable; una breve fase de adaptación cierra la mayor parte de la brecha, y un programa más largo puede alcanzar la paridad con el entrenamiento monolítico.
3. MoT altera la configuración de ingeniería del preentrenamiento al crear trabajos más pequeños, programables de forma independiente, reiniciables y ablables. Esto es particularmente adecuado para la investigación de entrenamiento a pequeña escala, lo que potencialmente la hace más barata, más reproducible y más fácil de iterar.

Los autores declaran explícitamente que la evidencia actual demuestra la componibilidad modular e identifica un régimen de programación de reutilización amortizada, en lugar de establecer una ventaja uniforme de cómputo o de hardware. El estudio está intencionadamente limitado en su alcance (una familia de modelos, un conjunto de datos, métricas de solo perplejidad) para servir como una prueba de mecanismo para un espacio de diseño concreto de sub-procesos andamiados y reutilizables.