NITP : Prédiction implicite du prochain jeton pour le pré-entraînement des LLM

Résumé

La prédiction standard du prochain jeton (NTP) supervise les modèles de langage uniquement via des étiquettes discrètes dans l'espace des logits de sortie. Nous soutenons que cette supervision one-hot éparse laisse l'espace des représentations latentes sous-contraint, permettant aux états cachés de dériver vers des configurations dégénérées et anisotropes pouvant limiter la généralisation. Pour remédier à ce problème, nous proposons la prédiction implicite du prochain jeton (NITP), qui enrichit la prédiction discrète avec une supervision continue dense directement dans l'espace des représentations. NITP entraîne le modèle à prédire le contenu sémantique implicite du prochain jeton, en utilisant les représentations des couches superficielles du même modèle comme cibles auto-supervisées stables. Nous fournissons une analyse théorique montrant que NITP régularise le paysage d'optimisation en atténuant les degrés de liberté sous-contraints et en favorisant une géométrie de représentation compacte et structurée. Empiriquement, sur des modèles denses et MoE allant de 0,5B à 9B paramètres, NITP améliore systématiquement les performances en aval avec un surcoût de calcul négligeable. Sur un modèle MoE de 9B, NITP obtient une amélioration absolue de 5,7 % sur MMLU-Pro, ainsi que des gains de 6,4 % sur C3 et 4,3 % sur CommonsenseQA, avec environ 2 % de FLOPs d'entraînement supplémentaires et aucun coût d'inférence additionnel. Notre implémentation est disponible à l'adresse https://github.com/aHapBean/NITP.

English

Standard next-token prediction (NTP) supervises language models solely through discrete labels in the output logit space. We argue that this sparse one-hot supervision leaves the latent representation space under-constrained, allowing hidden states to drift into degenerate and anisotropic configurations that can limit generalization. To address this issue, we propose Next Implicit Token Prediction (NITP), which augments discrete prediction with dense continuous supervision directly in the representation space. NITP trains the model to predict the implicit semantic content of the next token, using shallow-layer representations from the same model as stable self-supervised targets. We provide theoretical analysis showing that NITP regularizes the optimization landscape by mitigating under-constrained degrees of freedom and encouraging a compact, structured representation geometry. Empirically, across dense and MoE models ranging from 0.5B to 9B parameters, NITP consistently improves downstream performance with negligible computational overhead. On a 9B MoE model, NITP achieves a 5.7% absolute improvement on MMLU-Pro, along with gains of 6.4% on C3 and 4.3% on CommonsenseQA, with approximately 2% additional training FLOPs and no additional inference cost. Our implementation is available at https://github.com/aHapBean/NITP.