Investigadores eliminan el ruido del entrenamiento MTP en LLM con una sola línea de código

Casi todos los grandes modelos de lenguaje recientes resuelven durante el preentrenamiento una tarea algo más difícil que la habitual. No predicen solo el siguiente token, sino también varios tokens por delante. Esta técnica se llama multi-token prediction, MTP.

Sirve para dos cosas: el modelo rinde mejor en tareas de horizonte largo, como matemáticas y código, y las «cabezas» adicionales pueden usarse como redactores rápidos de borradores en la decodificación especulativa. El borrador de varios tokens se genera a bajo coste y el modelo principal lo verifica en una sola pasada. Los investigadores del proyecto Soft-MTP han descubierto que el entrenamiento estándar de estas cabezas arrastra una mezcla ruidosa y proponen una forma de eliminarla.

Para ello hace falta una línea de código adicional.

Qué falla en el entrenamiento MTP habitual

Las cabezas MTP se entrenan con la entropía cruzada estándar frente a los datos reales. Los autores demuestran que el gradiente de esta función de pérdida se descompone en dos partes, y que lo hace de forma exacta, sin aproximaciones.

• La primera parte es la autodestilación KL. La divergencia KL mide cuánto difiere una distribución de probabilidad de otra. Aquí empuja a cada cabeza MTP hacia la predicción del propio modelo para ese mismo token. Esta es la parte que aumenta la proporción de borradores aceptados en la decodificación especulativa: el token de una cabeza se acepta solo si coincide con la predicción del modelo principal. • La segunda parte es el residuo. Es el gradiente de la entropía cruzada habitual del siguiente token frente a los datos, que llega por otra vía. Tiene una relación señal-ruido baja y ahoga la señal KL útil.

¿Por qué es ruidoso el residuo? El token objetivo de los datos es una muestra concreta de la distribución real, un vector «one-hot» con un uno en una sola posición. La varianza de una muestra así no desaparece ni al aumentar el tamaño del modelo ni al entrenar durante más tiempo. Los autores lo midieron directamente. La relación señal-ruido del residuo es notablemente inferior a la del componente KL, y la brecha crece con el tamaño del modelo. Además, el residuo va en contra de la actualización que recibe la tarea principal de predicción del siguiente token. Es decir, perjudica tanto el entrenamiento de las cabezas MTP como el del propio modelo.

Una línea en lugar de matrices del tamaño del vocabulario

El gradiente KL limpio también puede obtenerse por la vía directa: construir las distribuciones del profesor y del alumno y minimizar la divergencia entre ellas. Pero con un vocabulario de decenas de miles de tokens habría que mantener en memoria matrices de logits enormes. Además, esta variante no funciona con los núcleos «fusionados» (fused kernels) de entropía cruzada estándar, diseñados precisamente para no materializar esas matrices. Por eso los modelos de frontera entrenan en la práctica las cabezas MTP con entropía cruzada convencional.

Los autores proponen una técnica llamada KL gradient-recovery trick. Si dos distribuciones comparten una cabeza de salida lineal, el gradiente de su KL puede recuperarse de forma exacta combinando dos gradientes de entropía cruzada ya calculados. Técnicamente se trata de un enrutamiento straight-through del gradiente: al estado oculto de la cabeza principal se le suma la expresión hk.detach() - hk. El valor no cambia, de modo que el entrenamiento de la tarea principal transcurre igual que antes. En cambio, a la representación MTP llega un gradiente adicional que, sumado al gradiente habitual de la pérdida MTP, resta el residuo ruidoso. El resultado es que a las cabezas MTP y al tronco común del modelo solo llega la señal KL.

Así funciona el método Soft-MTP. No añade parámetros nuevos ni genera una sobrecarga apreciable. Funciona tanto con MTP al estilo de DeepSeek como con la variante de Meta, y es compatible con el paralelismo que se usa en el preentrenamiento de modelos grandes. El código está publicado en GitHub: JaydenTeoh/Soft-MTP.

Qué mostraron los experimentos de 167M a 1,2B

Los autores entrenaron modelos de 167 millones a 1,2 mil millones de parámetros y compararon Soft-MTP con el MTP estándar.

• La perplejidad en la predicción del siguiente token es menor en todos los tamaños, y la brecha crece con la escala. • La perplejidad de las cabezas MTP es menor a partir del modelo de 689 millones de parámetros. • La ventaja de Soft-MTP en precisión en tareas sin ajuste adicional (zero-shot) se mantiene en cada escala. • La decodificación autoespeculativa se acelera en todos los dominios probados. • Tras el ajuste posterior, el modelo de 1,2 mil millones de parámetros resuelve mejor las tareas de matemáticas y código.

Aparte, compararon el método con la minimización KL directa, que requiere calcular los logits de forma explícita. En la mayoría de las métricas Soft-MTP no se queda por detrás de ella y en algunas la supera. La causa probable es que Soft-MTP sigue entrenando la cabeza de salida común con entropía cruzada convencional frente a los datos. Con un KL puro, esa cabeza se desplazaría hacia las propias predicciones del modelo.

Precio y limitaciones

El método implica un compromiso claro. Al eliminar el residuo ruidoso, Soft-MTP introduce un sesgo: las cabezas MTP se orientan hacia la distribución del propio modelo y no hacia la distribución real de los datos. Los autores consideran que el intercambio resulta cada vez más favorable al crecer la escala: el error del modelo profesor disminuye, mientras que el ruido de muestreo permanece. Además, las cabezas MTP sirven ante todo para acelerar la inferencia, donde la coincidencia con la predicción del modelo principal es justo lo que se busca.

Para los equipos que preentrenan sus propios LLM con MTP, la conclusión práctica es sencilla. El cambio es minúsculo, no cuesta nada en memoria ni en velocidad y no rompe los habituales núcleos fused. Los experimentos se limitan a modelos de hasta 1,2 mil millones de parámetros, así que el resultado aún está por comprobar a la escala de DeepSeek-V3. Pero la tendencia de que la ganancia crece junto con el modelo resulta alentadora.