Pré-entraînement efficace avec superposition de tokens

Efficient Pre-Training with Token Superposition

May 7, 2026

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

cs.AI

Résumé

Le pré-entraînement des grands modèles de langage est souvent d’un coût prohibitif et inefficace à grande échelle, nécessitant des modifications complexes et invasives pour atteindre un débit de données élevé. Dans ce travail, nous présentons le Token-Superposition Training (TST), une méthode de remplacement simple qui améliore significativement le débit de données par FLOP lors du pré-entraînement, sans modifier le parallélisme, l’optimiseur, le tokeniseur, les données ou l’architecture du modèle. Le TST se déroule en deux phases : (i) une phase de superposition très efficace où nous combinons plusieurs tokens contigus en un seul ensemble et entraînons avec un objectif d’entropie croisée multi-classe (MCE), et (ii) une phase de récupération où nous revenons à l’entraînement standard. Nous évaluons extensivement le TST à l’échelle de 270M et 600M paramètres, et le validons sur un modèle de 3B et un modèle A1B à mélange d’experts de 10B, démontrant sa grande robustesse dans différents contextes. In fine, le TST surpasse systématiquement la perte de référence et les évaluations en aval, et, dans des configurations à perte égale, le TST permet une réduction allant jusqu’à 2,5 fois du temps total de pré-entraînement à l’échelle du modèle A1B de 10B.

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.