La predicción multi-tokeno (MTP) se ha convertido en un componente estándar en el preentrenamiento de modelos de lenguaje, mejorando el rendimiento en tareas posteriores y permitiendo una inferencia más rápida mediante la decodificación especulativa. Sin embargo, un estudio reciente identifica una deficiencia crítica en el entrenamiento estándar de MTP: un componente residual en el gradiente de la entropía cruzada interfiere con el aprendizaje tanto del siguiente token como del multi-tokeno, un problema que se agrava a medida que aumenta la escala del modelo. Este ruido residual enterraba la señal útil de la divergencia de KL, que es la clave para mejorar la decodificación especulativa.
Para resolver esto, los autores introducen la 'trampa de recuperación del gradiente de KL': al combinar los gradientes de entropía cruzada de dos distribuciones que comparten una cabeza de salida lineal, se puede recuperar exactamente el gradiente de la divergencia de KL. A diferencia de minimizar directamente la divergencia de KL, este método no requiere matrices de logits de tamaño vocabulario y es compatible con los kernels estándar de entropía cruzada fusionados. Esta técnica permite implementar Soft-MTP, un cambio de código de una sola línea que elimina el residuo ruidoso y backpropaga únicamente la señal de autodistilación de KL. Soft-MTP es compatible con los estilos de MTP de DeepSeek y Meta, no añade parámetros y no genera sobrecargas medibles. En pruebas desde 167M hasta 1.2B parámetros, Soft-MTP mejora la perplexidad de la predicción del siguiente token, acelera la decodificación especulativa y aumenta la precisión en tareas de razonamiento matemático y de código tras el ajuste fino, superando o igualando a la minimización directa de la divergencia de KL en la mayoría de las métricas.