Résumé Technique : BlockBatch

Énoncé du Problème

Les modèles de diffusion linguistique (dLLM) génèrent du texte en débruitant itérativement plusieurs positions de tokens en parallèle, offrant une alternative au décodage strictement autoregressif. Cependant, l'inférence efficace se heurte à un compromis critique de granularité. Les méthodes d'inférence par blocs existantes (par exemple, Fast-dLLM) sélectionnent une seule taille de bloc par requête. Les petits blocs préservent le conditionnement local mais nécessitent de nombreuses étapes de débruitage (NFE élevé), tandis que les grands blocs augmentent le parallélisme mais risquent des engagements prématurés et une accumulation d'erreurs de cache (dérive du KV-cache). Les méthodes d'accélération actuelles traitent la taille de bloc comme un hyperparamètre fixe, échouant à exploiter la nature complémentaire des différentes tailles de bloc pour différentes parties d'une même génération.

Méthodologie : BlockBatch

Les auteurs proposent BlockBatch, un cadre d'inférence en ligne sans entraînement qui traite la taille de bloc comme une dimension de branchement plutôt que comme un paramètre fixe. Au lieu de sélectionner une seule taille de bloc, BlockBatch exécute plusieurs branches avec différentes tailles de bloc (par exemple, 4, 8, 16, 32, 64, 128) pour la même requête au sein d'un seul passage avant groupé.

Le cadre repose sur trois mécanismes fondamentaux pour coordonner ces branches :

1. Fusion pilotée par la confiance : Les branches échangent des propositions de tokens uniquement lorsqu'elles sont compatibles (c'est-à-dire que leurs préfixes décodés correspondent exactement). Une branche destination n'accepte un token proposé par une branche source que si la distribution de probabilité de la branche destination elle-même attribue à ce token un score de confiance supérieur à un seuil fixe. Cela permet aux branches en retard de progresser sans imposer d'engagements de tokens incohérents.
2. Synchronisation basée sur le leader : Pour empêcher les branches en retard de gaspiller des calculs dans des régions obsolètes de l'espace du KV-cache, une règle de synchronisation est appliquée. Si une branche leader dépasse une branche en retard d'un nombre spécifique de tokens dépassant un seuil, la branche en retard copie l'état de la séquence et le cache KV du leader, réinitialisant ainsi efficacement sa trajectoire vers le chemin fiable le plus avancé.
3. Rafraîchissements périodiques de la séquence complète : Le débruitage local par bloc ne met à jour qu'un sous-ensemble de la séquence, ce qui peut entraîner une dérive accumulée par rapport à l'état KV globalement cohérent. BlockBatch recompute périodiquement le cache KV pour l'ensemble de la séquence des branches actives. Cela réancre les branches à un état globalement cohérent, corrigeant les erreurs accumulées lors des mises à jour locales.

Optimisation Système : Pour garantir l'efficacité, BlockBatch regroupe les positions de requête non masquées de toutes les branches actives dans un seul tampon de longueur variable. Cela permet au modèle de traiter plusieurs branches en un seul passage avant, en utilisant le noyau varlen de FlashAttention, assurant ainsi que la latence du débruitage multi-branches est comparable à celle d'une seule branche.

Caractérisations Clés

L'article fournit des caractérisations théoriques et empiriques expliquant pourquoi cette approche fonctionne :

• Trajectoire du KV-Cache : Différentes tailles de bloc induisent différentes trajectoires dans l'espace du KV-cache. Les petits blocs suivent des trajectoires conservatrices avec un conditionnement fin, tandis que les grands blocs effectuent des progrès locaux agressifs. Les rafraîchissements de séquence complète agissent comme des corrections globales qui empêchent ces trajectoires de trop diverger.
• Dynamique au niveau du Token : Les branches partagent souvent un préfixe initial, bifurquent à des positions sémantiquement décisives, et se réaccordent ultérieurement sur des tokens syntaxiquement légers (par exemple, espaces blancs, opérateurs, chiffres). Le cadre exploite cela en fusionnant les branches lorsqu'elles s'accordent sur des tokens à haute confiance et en les synchronisant lorsqu'elles divergent significativement.

Résultats Expérimentaux

BlockBatch a été évalué sur trois dLLM représentatifs (LLaDA-1.5-8B, LLaDA-Instruct-8B, Dream-Base-7B) à travers quatre benchmarks (GSM8K, MATH, HumanEval, MBPP).

• Efficacité : BlockBatch réduit le nombre moyen d'évaluations de fonctions de débruitage (NFE) de 26,6 % par rapport à Fast-dLLM (la référence de l'état de l'art).
• Vitesse : Il réalise une accélération globale moyenne de 1,33× par rapport à Fast-dLLM, avec des accélérations atteignant jusqu'à 5,8× dans des configurations spécifiques (par exemple, LLaDA-Instruct-8B sur MBPP).
• Précision : Malgré la réduction agressive des NFE et de la latence, BlockBatch préserve une précision comparable ou supérieure à celle des références pour toutes les tâches.
• Comparaison : Il surpasse systématiquement le décodage de diffusion Vanilla, Fast-dLLM et LocalLeap en termes d'efficacité des NFE et de latence sur les douze configurations modèle-tâche testées.

Importance et Revendications

L'article revendique que la diversité des tailles de bloc est un axe pratique et précédemment sous-exploré pour l'inférence de dLLM en branchement parallèle. En démontrant que différentes tailles de bloc offrent des compromis précision-efficacité complémentaires, les auteurs soutiennent que traiter la taille de bloc comme une dimension de branchement dynamique permet un décodage plus efficace que n'importe quelle configuration fixe unique.

L'importance de BlockBatch réside dans sa capacité à :

1. Exploiter la structure inhérente de l'inférence dLLM où différentes granularités sont utiles pour différentes parties d'une génération.
2. Coordonner plusieurs trajectoires de décodage sans entraînement, en utilisant uniquement des opérations au moment de l'inférence (fusion, synchronisation, rafraîchissement).
3. Fournir un cadre général qui améliore l'efficacité tout en maintenant la qualité du texte généré, identifiant la diversité des tailles de bloc comme une stratégie viable pour accélérer les modèles linguistiques basés sur la diffusion.