jaydenteoh/soft-mtp is the official implementation repository for the “Multi-Token Prediction Done Right” paper. It is written in Python and released under the MIT license. The repository implements standard multi-token prediction and the Soft-MTP objective function so that either can be selected in a single loss function. The core of this project therefore lies in handling multiple objective functions within one implementation. It is useful material to consult alongside the paper, since the related training setup and evaluation procedure can be reviewed together with the code.
The README lays out the training setup that uses FineWeb-Edu tokens, the pretrained models, and the training and evaluation commands. The model scale in the setup ranges from 167M to 1.2B parameters. It also combines a gradient recovery technique with the cross-entropy computation. The test command is structured to check whether the computed gradient matches the result obtained by minimizing KL directly. This configuration makes it possible to look at not only the description of the objective functions but also how to verify those computations within the repository.
However, the provided information contains no per-task performance figures for the trained models, nor any comparison results against standard multi-token prediction. It is therefore not possible to judge which performs better from the setup and the test command alone. There is also no latest release registered. The verifiable scope currently ends at the paper's objective function implementation and the procedure for checking its gradient computation.