NITP: Предсказание следующего неявного токена для предварительного обучения LLM

Аннотация

Стандартное предсказание следующего токена (NTP) управляет языковыми моделями исключительно с помощью дискретных меток в пространстве выходных логитов. Мы утверждаем, что такое разреженное one-hot управление оставляет пространство скрытых представлений недостаточно ограниченным, позволяя скрытым состояниям дрейфовать в вырожденные и анизотропные конфигурации, что может ограничивать обобщение. Для решения этой проблемы мы предлагаем предсказание следующего неявного токена (NITP), которое дополняет дискретное предсказание плотным непрерывным управлением непосредственно в пространстве представлений. NITP обучает модель предсказывать неявное семантическое содержание следующего токена, используя представления мелких слоев той же модели в качестве стабильных целей самоконтроля. Мы предоставляем теоретический анализ, показывающий, что NITP регулирует ландшафт оптимизации, уменьшая недостаточно ограниченные степени свободы и поощряя компактную, структурированную геометрию представлений. Эмпирически, на плотных моделях и моделях смеси экспертов (MoE) от 0,5 до 9 миллиардов параметров, NITP последовательно улучшает производительность на последующих задачах с пренебрежимо малыми вычислительными затратами. На модели MoE с 9 миллиардами параметров NITP достигает абсолютного улучшения на 5,7% на MMLU-Pro, а также прироста на 6,4% на C3 и 4,3% на CommonsenseQA, при дополнительных затратах FLOPs на обучение около 2% и без дополнительных затрат на вывод. Наша реализация доступна по адресу 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.