Pre-entrenamiento eficiente con superposición de tokens

Efficient Pre-Training with Token Superposition

May 7, 2026

Autores: Bowen Peng, Théo Gigant, Jeffrey Quesnelle

cs.AI

Resumen

El preentrenamiento de modelos de lenguaje grandes suele ser prohibitivamente costoso e ineficiente a escala, requiriendo modificaciones complejas e invasivas para alcanzar un alto rendimiento de datos. En este trabajo presentamos el Entrenamiento por Superposición de Tokens (TST, por sus siglas en inglés), un método sustitutivo directo y sencillo que mejora significativamente el rendimiento de datos por FLOPs durante el preentrenamiento sin modificar el paralelismo, el optimizador, el tokenizador, los datos ni la arquitectura del modelo. TST se realiza en dos fases: (i) una fase de superposición altamente eficiente donde combinamos varios tokens contiguos en una bolsa y entrenamos usando un objetivo de entropía cruzada multi-caliente (MCE, por sus siglas en inglés), y (ii) una fase de recuperación donde volvemos al entrenamiento estándar. Evaluamos exhaustivamente TST a escalas de 270M y 600M parámetros, y lo validamos en un modelo de mezcla de expertos A1B de 3B y 10B, demostrando que es altamente robusto en diferentes entornos. En última instancia, TST supera consistentemente la pérdida base y las evaluaciones posteriores, y en escenarios de pérdida equivalente, TST logra una reducción de hasta 2.5 veces en el tiempo total de preentrenamiento a la escala de 10B A1B.

English

Pre-training of Large Language Models is often prohibitively expensive and inefficient at scale, requiring complex and invasive modifications in order to achieve high data throughput. In this work, we present Token-Superposition Training (TST), a simple drop-in method that significantly improves the data throughput per FLOPs during pre-training without modifying the parallelism, optimizer, tokenizer, data, or model architecture. TST is done in two phases: (i) A highly efficient superposition phase where we combine many contiguous tokens into one bag and train using a multi-hot cross-entropy (MCE) objective, and (ii) a recovery phase where we revert back to standard training. We extensively evaluate TST on the scale of 270M and 600M parameters and validate on 3B and a 10B A1B mixture of experts model, demonstrating that it is highly robust in different settings. Ultimately, TST consistently outperforms baseline loss and downstream evaluations, and under equal-loss settings, TST yields up to a 2.5x reduction in total pre-training time at the 10B A1B scale.