jaydenteoh/soft-mtp は「Multi-Token Prediction Done Right」論文の公式実装リポジトリである。Python で書かれ、MIT ライセンスで公開されている。リポジトリは標準のマルチトークン予測と Soft-MTP 目的関数を実装しており、1 つの損失関数でどちらかを選べる。したがって、このプロジェクトの中心は複数の目的関数を 1 つの実装の中で扱えるようにした点にある。論文を読む際は、関連する学習設定と評価手順をコードと併せて確認できる資料となる。
README には FineWeb-Edu トークンを用いた学習設定と事前学習済みモデル、学習および評価のコマンドが整理されている。設定のモデル規模は 167M から 1.2B パラメータまでだ。また、勾配復元(gradient recovery)技術を交差エントロピー計算と組み合わせている。テストコマンドは、計算された勾配が KL を直接最小化した結果と一致するかを確認できる構成になっている。この構成により、目的関数の説明だけでなく、その計算を確かめ方もリポジトリの中で見ることができる。
ただし、提供された情報には学習済みモデルの課題別性能や、標準のマルチトークン予測との比較結果はない。そのため、設定とテストコマンドだけでは実際の性能の優劣を判断できない。登録された最新リリースも存在しない。現在確認できる範囲は、論文の目的関数の実装と、その勾配計算を点検する手順までである。