Pesquisadores eliminam o ruído do treinamento MTP em LLMs com uma única linha de código

Quase todos os grandes modelos de linguagem recentes resolvem, no pré-treinamento, uma tarefa um pouco mais difícil que a convencional. Eles preveem não só o próximo token, mas vários tokens à frente. A técnica se chama multi-token prediction, ou MTP. Ela serve a dois propósitos: o modelo lida melhor com tarefas de horizonte longo, como matemática e código, e as "cabeças" adicionais podem atuar como redatoras rápidas de rascunhos na decodificação especulativa. Um rascunho com vários tokens sai barato, e o modelo principal o verifica em uma única passada. Pesquisadores do projeto Soft-MTP descobriram que o treinamento padrão dessas cabeças carrega uma mistura ruidosa e propuseram uma forma de eliminá-la. Para isso, basta uma linha de código adicional.

O que há de errado com o treinamento MTP convencional

As cabeças MTP são treinadas com entropia cruzada padrão contra os dados reais. Os autores mostraram que o gradiente dessa função de perda se decompõe em duas partes, e de forma exata, sem aproximações.

• A primeira parte é a autodestilação por KL. A divergência KL mede o quanto uma distribuição de probabilidades difere de outra. Aqui, ela puxa cada cabeça MTP em direção à previsão do próprio modelo para o mesmo token. É essa parte que aumenta a taxa de rascunhos aceitos na decodificação especulativa: o token da cabeça só é aceito se coincidir com a previsão do modelo principal. • A segunda parte é o resíduo. Trata-se do gradiente da entropia cruzada comum do próximo token contra os dados, que chega por outro caminho. Ele tem baixa relação sinal-ruído e abafa o sinal útil da KL.

Por que o resíduo é ruidoso? O token-alvo nos dados é uma única amostra da distribuição real, um vetor "one-hot" com 1 em uma só posição. A variância dessa amostra não desaparece nem com o aumento do modelo, nem com treinamento prolongado. Os autores mediram isso diretamente. No resíduo, a relação sinal-ruído é bem menor que na componente KL, e a diferença cresce com o tamanho do modelo. Além disso, o resíduo aponta contra a atualização que a tarefa principal de prever o próximo token recebe. Ou seja, ele prejudica tanto o treinamento das cabeças MTP quanto o do próprio modelo.

Uma linha em vez de matrizes do tamanho do vocabulário

É possível obter o gradiente KL limpo pelo caminho direto: construir as distribuições do professor e do aluno e minimizar a divergência entre elas. Mas, com um vocabulário de dezenas de milhares de tokens, seria preciso manter na memória matrizes enormes de logits. Além disso, essa abordagem não funciona com os kernels "fundidos" (fused kernels) de entropia cruzada, criados justamente para evitar materializar essas matrizes. Por isso, os modelos de ponta treinam, na prática, as cabeças MTP com entropia cruzada comum.

Os autores propuseram o KL gradient-recovery trick. Se as duas distribuições compartilham uma cabeça de saída linear, o gradiente da KL entre elas pode ser recuperado com exatidão combinando dois gradientes de entropia cruzada já calculados. Tecnicamente, é um roteamento de gradiente straight-through: soma-se ao estado oculto da cabeça principal a expressão hk.detach() - hk. O valor não muda, então o treinamento da tarefa principal segue como antes. Em compensação, um gradiente adicional flui para a representação MTP e, somado ao gradiente usual da perda MTP, subtrai o resíduo ruidoso. No fim, apenas o sinal KL chega às cabeças MTP e ao tronco comum do modelo.

É assim que o Soft-MTP funciona. O método não adiciona parâmetros novos nem cria custo adicional perceptível. Funciona tanto com o MTP no estilo DeepSeek quanto com a variante da Meta, e é compatível com o paralelismo usado no pré-treinamento de modelos grandes. O código está publicado no GitHub: JaydenTeoh/Soft-MTP.

O que mostraram os experimentos de 167M a 1,2B

Os autores treinaram modelos de 167 milhões a 1,2 bilhão de parâmetros e compararam o Soft-MTP com o MTP padrão.

• A perplexidade na previsão do próximo token é menor em todos os tamanhos, e a diferença cresce com a escala. • A perplexidade das cabeças MTP é menor a partir do modelo de 689 milhões de parâmetros. • A vantagem de acurácia do Soft-MTP em tarefas sem ajuste fino (zero-shot) se mantém em todas as escalas. • A decodificação autoespeculativa fica mais rápida em todos os domínios testados. • Após o ajuste fino, o modelo de 1,2 bilhão de parâmetros resolve melhor tarefas de matemática e código.

Os autores também compararam o método com a minimização direta de KL, que exige o cálculo explícito dos logits. Na maioria das métricas, o Soft-MTP não perde para ela e, em alguns casos, a supera. A provável razão é que o Soft-MTP continua treinando a cabeça de saída compartilhada com entropia cruzada comum sobre os dados. Com KL pura, essa cabeça tenderia a se deslocar para as previsões do próprio modelo.

Custo e limitações

O método envolve um compromisso claro. Ao remover o resíduo ruidoso, o Soft-MTP introduz um viés: as cabeças MTP são puxadas para a distribuição do próprio modelo, e não para a distribuição real dos dados. Os autores consideram que, conforme a escala aumenta, essa troca fica cada vez mais vantajosa: o erro do modelo-professor diminui, enquanto o ruído da amostragem permanece. Além disso, as cabeças MTP servem principalmente para acelerar a inferência, e nela a concordância com a previsão do modelo principal é exatamente o que se quer.

Para equipes que pré-treinam seus próprios LLMs com MTP, a conclusão prática é simples. A mudança é minúscula, não custa nada em memória nem em velocidade e não quebra os kernels fundidos habituais. Os experimentos se limitam a modelos de até 1,2 bilhão de parâmetros, então o resultado ainda precisa ser verificado na escala do DeepSeek-V3. Mas a tendência de o ganho crescer junto com o modelo é animadora.