NITP: Predição do Próximo Token Implícito para o Pré-treinamento de LLMs

Resumo

A predição padrão do próximo token (NTP) supervisiona modelos de linguagem unicamente por meio de rótulos discretos no espaço de logits de saída. Argumentamos que essa supervisão esparsa de um vetor one-hot deixa o espaço de representação latente sub-condicionado, permitindo que estados ocultos derivem para configurações degeneradas e anisotrópicas que podem limitar a generalização. Para lidar com esse problema, propomos a Predição do Próximo Token Implícito (NITP), que complementa a predição discreta com supervisão densa contínua diretamente no espaço de representação. O NITP treina o modelo para prever o conteúdo semântico implícito do próximo token, utilizando representações de camadas superficiais do próprio modelo como alvos autossupervisionados estáveis. Fornecemos análise teórica mostrando que o NITP regulariza a paisagem de otimização ao mitigar graus de liberdade sub-condicionados e incentivar uma geometria de representação compacta e estruturada. Empiricamente, em modelos densos e MoE variando de 0,5B a 9B parâmetros, o NITP melhora consistentemente o desempenho downstream com custo computacional negligenciável. Em um modelo MoE de 9B, o NITP alcança uma melhoria absoluta de 5,7% no MMLU-Pro, junto com ganhos de 6,4% no C3 e 4,3% no CommonsenseQA, com aproximadamente 2% adicionais de FLOPs de treinamento e nenhum custo adicional de inferência. Nossa implementação está disponível em 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.