Itai Gat Heli Ben-Hamu Marton Havasi Daniel Haziza Jeremy Reizenstein Gabriel Synnaeve David Lopez-Paz Brian Karrer Yaron Lipman

Résumé

Les modèles linguistiques à prédiction du jeton suivant autoregressifs offrent des capacités puissantes, mais rencontrent des défis importants dans leur déploiement pratique en raison des coûts computationnels et mémoire élevés associés à l’inférence, en particulier durant l’étape de décodage. Nous introduisons Set Block Decoding (SBD), un paradigme simple et flexible qui accélère la génération en intégrant, au sein d’une même architecture, la prédiction standard du jeton suivant (NTP) et la prédiction de jetons masqués (MATP). SBD permet à modèle d’échantillonner en parallèle plusieurs jetons futurs, non nécessairement consécutifs, ce qui constitue une différence clé par rapport aux méthodes d’accélération antérieures. Cette flexibilité ouvre la voie à l’utilisation de solveurs avancés issus de la littérature sur les diffusions discrètes, offrant des accélérations significatives sans compromettre la précision. SBD ne nécessite ni modification architecturale ni nouveaux hyperparamètres d’entraînement, préserve la compatibilité avec le cache exact des valeurs (KV-caching) et peut être mis en œuvre par une adaptation fine (fine-tuning) de modèles existants de prédiction du jeton suivant. En adaptant finement Llama-3.1 8B et Qwen-3 8B, nous démontrons que SBD permet de réduire de 3 à 5 fois le nombre de passages avant (forward passes) requis pour la génération, tout en atteignant une performance équivalente à celle obtenue avec un entraînement NTP classique.