NITP: Predicción del Siguiente Token Implícito para el Preentrenamiento de LLM

Resumen

La predicción estándar del siguiente token (NTP) supervisa los modelos de lenguaje únicamente mediante etiquetas discretas en el espacio de logits de salida. Argumentamos que esta supervisión one-hot dispersa deja el espacio de representación latente subconstreñido, permitiendo que los estados ocultos deriven hacia configuraciones degeneradas y anisotrópicas que pueden limitar la generalización. Para abordar este problema, proponemos la Predicción del Siguiente Token Implícito (NITP), que complementa la predicción discreta con una supervisión continua densa directamente en el espacio de representación. NITP entrena al modelo para predecir el contenido semántico implícito del siguiente token, utilizando representaciones de capas superficiales del mismo modelo como objetivos autosupervisados estables. Proporcionamos un análisis teórico que muestra que NITP regulariza el panorama de optimización al mitigar los grados de libertad subconstreñidos y fomentar una geometría de representación compacta y estructurada. Empíricamente, en modelos densos y MoE que van desde 0.5B hasta 9B parámetros, NITP mejora sistemáticamente el rendimiento descendente con una carga computacional insignificante. En un modelo MoE de 9B, NITP logra una mejora absoluta del 5.7% en MMLU-Pro, junto con ganancias del 6.4% en C3 y del 4.3% en CommonsenseQA, con aproximadamente un 2% adicional de FLOPs de entrenamiento y ningún coste adicional de inferencia. Nuestra implementación está disponible en 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.