Polymathic AI et Princeton publient FlashMap, un outil qui accélère de 6,14 fois le décodage des LLM

Des chercheurs de Polymathic AI, de Princeton University, de l’USC ISI, de NYU et du Flatiron Institute ont publié FlashMap, un « drafter » destiné au décodage spéculatif des modèles de langage. Selon l’équipe, il peut accélérer jusqu’à 6,14 fois la génération sur Qwen3-8B, sans modifier le résultat produit. Les poids du drafter pour Qwen3-8B et Qwen3-4B sont disponibles sur Hugging Face.

  • FlashMap accélère jusqu’à 6,14 fois la génération de Qwen3-8B sans modifier le résultat.
  • Il propose un bloc de 16 tokens d’un coup, et ces tokens restent cohérents entre eux.
  • En mode deux étapes, il dépasse DFlash sur la longueur d’acceptation dans les 18 configurations testées.
  • Les poids sont publiés sous licence Apache-2.0, et le code sous licence MIT.

Le décodage spéculatif consiste à faire proposer plusieurs tokens (les unités de texte manipulées par un modèle) par un petit modèle auxiliaire, que le grand modèle vérifie ensuite en une seule passe. FlashMap propose un bloc complet de 16 tokens d’un coup, et selon les auteurs, les tokens de ce bloc restent cohérents entre eux. Le drafter compte 0,5 milliard de paramètres, soit 500 millions pour Qwen3-8B et 475 millions pour Qwen3-4B, et fonctionne sur le même GPU que le modèle cible.

Pourquoi les drafters en diffusion perdent-ils en précision ?

Le drafter DFlash prédit chaque position du bloc indépendamment, ce que les chercheurs appellent une approximation factorisée. La source prend l’exemple d’une suite de mots qui peut être « new york » ou « san diego » avec des probabilités égales. Un drafter indépendant propose alors « new diego » ou « san york » une fois sur deux, et le modèle cible rejette ces suites à la deuxième position.

FlashMap s’appuie sur les Flow Map Language Models. Chaque échantillon de bruit correspond à un bloc complet, et ce bruit détermine conjointement la version retenue pour chaque position. Le drafter est un transformeur de diffusion (DiT) à 6 blocs, avec une dimension d’embedding de 2048 et 16 têtes d’attention. Il lit les états cachés des couches 1,9, 17,25 et 33 du modèle cible.

Des résultats annoncés sur neuf benchmarks

Les chercheurs ont réentraîné EAGLE-3 et DFlash sur les mêmes données et dans les mêmes conditions, puis ont testé les drafters sur neuf benchmarks couvrant les mathématiques, le code et le chat, à des températures de 0 et 1. Selon la source, FlashMap à une étape est le drafter le plus rapide sur chaque tâche testée sur Qwen3-8B, aux deux températures. Il dépasse DFlash sur la longueur d’acceptation dans 15 configurations sur 18, avec environ la moitié de ses paramètres.

En mode deux étapes, la longueur d’acceptation dépasse DFlash dans 18 configurations sur 18, avec 10 % de mieux en moyenne. Les auteurs indiquent aussi que sur Qwen3-4B, à température 0, FlashMap est devancé par DFlash sur 4 des 9 tâches. Le cas le plus faible rapporté est un gain de 2,28 fois sur Alpaca à température 1. Le modèle a été entraîné sur 500 000 échantillons issus d’Open-PerfectBlend, pour un total d’environ 4 600 heures de calcul sur des H200.

Pour le déploiement, le paquet flashmap nécessite uniquement torch et transformers en version 4.56 ou supérieure, et les projets indiquent des chemins d’utilisation via vLLM 0.31 et SGLang 0.5.20. Les poids sont publiés sous licence Apache-2.0 et le code sous licence MIT.