Accélération des LLMs de diffusion via un décodage parallèle adaptatif

Accelerating Diffusion LLMs via Adaptive Parallel Decoding

May 31, 2025

Auteurs: Daniel Israel, Guy Van den Broeck, Aditya Grover

cs.AI

Résumé

La vitesse de génération des LLM est limitée par le décodage autorégressif, où les tokens sont prédits séquentiellement un par un. Alternativement, les modèles de langage de grande taille par diffusion (dLLM) permettent théoriquement une génération parallèle de tokens, mais en pratique peinent à atteindre la vitesse des modèles autorégressifs sans sacrifier significativement la qualité. Nous introduisons donc le décodage parallèle adaptatif (APD), une méthode novatrice qui ajuste dynamiquement le nombre de tokens échantillonnés en parallèle. Nous y parvenons en définissant un mélange multiplicatif entre les probabilités marginales du dLLM et la probabilité conjointe des séquences sous un petit modèle autorégressif auxiliaire. Cela inverse la configuration standard du décodage spéculatif, où l'objectif est d'échantillonner à partir d'un vérificateur autorégressif de grande taille en proposant des séquences à partir d'un modèle plus petit. Nous optimisons en outre l'APD en activant la mise en cache KV et en limitant la taille de l'entrée masquée. Au total, notre méthode propose trois paramètres ajustables pour équilibrer de manière flexible le débit et la qualité. Nous montrons que l'APD offre un débit nettement supérieur avec des dégradations minimes de la qualité sur les benchmarks en aval.

English

The generation speed of LLMs are bottlenecked by autoregressive decoding, where tokens are predicted sequentially one by one. Alternatively, diffusion large language models (dLLMs) theoretically allow for parallel token generation, but in practice struggle to achieve the speed of autoregressive models without significantly sacrificing quality. We therefore introduce adaptive parallel decoding (APD), a novel method that dynamically adjusts the number of tokens sampled in parallel. We achieve this by defining a multiplicative mixture between the dLLM marginal probabilities and the joint probability of sequences under a small auxiliary autoregressive model. This inverts the standard setup of speculative decoding, where the goal is to sample from a large autoregressive verifier by drafting from a smaller model. We further optimize APD by enabling KV caching and limiting the size of the masked input. Altogether, our method puts forward three tunable parameters to flexibly tradeoff throughput and quality. We show that APD provides markedly higher throughput with minimal quality degradations on downstream benchmarks.