La prédiction multi-tokens (MTP, qui amène le modèle à prédire également les quelques tokens suivants) est déjà une recette courante de pré-entraînement, améliorant à la fois les performances en aval et servant de tête de brouillon pour le décodage spéculatif. Les auteurs décomposent avec précision le gradient de l'entropie croisée de la MTP en deux parties : un terme d'auto-distillation KL qui attire la tête MTP vers la distribution du mot suivant du modèle lui-même (utile, améliorant le décodage spéculatif), et un résidu au rapport signal/bruit très faible, qui perturbe l'apprentissage et s'aggrave avec l'échelle. Soft-MTP utilise le « routage de gradient à travers » (straight-through gradient routing) pour ne conserver que le terme KL, réutilise le calcul d'entropie croisée existant, ne nécessite pas de matrice de logits sur l'ensemble du vocabulaire et est compatible avec les noyaux d'entropie croisée fusionnés.
J'aime 9 · Étoiles du code 6
Raison
- Modification minimale : une seule ligne de code pour la MTP standard, sans ajout de paramètres, les auteurs affirment qu'il n'y a pas de surcoût mesurable ; supporte à la fois les MTP de type DeepSeek et de type Meta.
- Les auteurs affirment qu'à chaque échelle, de 167M à 1.2B, la perplexité du mot suivant est plus faible, le décodage spéculatif plus rapide et la précision en zéro-shot plus élevée ; après un fine-tuning à 1.2B, le raisonnement mathématique et code est plus fort, et les performances ne sont pas inférieures à la minimisation directe de KL sur la plupart des métriques.
- Reproductible : quatre échelles × trois objectifs (NTP, MTP, Soft-MTP), soit 12 points de contrôle (checkpoints) hébergés sur Hugging Face, la configuration d'entraînement et les scripts d'évaluation du décodage spéculatif sont inclus ; les tests dans tests/ permettent de vérifier la cohérence entre le gradient et le KL explicite en float64 ; licence MIT.
Utilisation
- Code : git clone github.com/JaydenTeoh/Soft-MTP → pip install -r requirements.txt → python -m src.fetch_data --shards 5 (FineWeb-Edu pré-tokenisé, chaque configuration entraînée sur 31.5B tokens, au moins 4 shards, environ 120 Go d'espace disque requis lors de l'assemblage) ; l'optimiseur Muon nécessite PyTorch 2.9 ou supérieur.
- Entraînement : torchrun --standalone --nproc_per_node=4 -m src.train --config configs/167m_fineweb.yaml --objective soft_mtp (remplacer par hard_mtp / direct_kl / ntp pour les comparaisons).
- Voir les modèles prêts à l'emploi : AutoModelForCausalLM.from_pretrained("JaydenTeoh/soft-mtp-689m-model", trust_remote_code=True), utiliser le tokenizer original gpt2 ; pour l'évaluation du décodage spéculatif, utiliser python -m src.eval_spec --ckpts ….
- Pour l'intégrer dans votre propre code d'entraînement, consultez combined_mtp_loss() dans src/models/loss.py.
Risques
- Les modèles publiés sont tous des petits modèles de recherche de moins de 1.2B, entraînés uniquement sur FineWeb-Edu, et ne sont pas des modèles de dialogue prêts à l'emploi ; l'efficacité à plus grande échelle n'a pas été vérifiée.
- La valeur réside principalement dans la phase de pré-entraînement, il faut entraîner depuis le début pour en bénéficier ; l'évaluation de l'entraînement est connectée par défaut à Weights & Biases.
- Le résumé ne donne pas le facteur d'accélération spécifique du décodage spéculatif, les résultats sont rapportés par les auteurs ; le README rappelle que l'accélération en temps réel dépend du GPU, les comparaisons doivent être effectuées sur la même machine.