SparseD: Atenção Esparsa para Modelos de Linguagem de Difusão

Resumo

Embora os modelos de linguagem por difusão (DLMs) ofereçam uma alternativa promissora aos modelos autoregressivos (ARs), os DLMs de código aberto existentes sofrem com alta latência de inferência. Esse gargalo ocorre principalmente devido à complexidade quadrática da atenção em relação ao comprimento do contexto ao calcular todos os pares de consulta-chave. Intuitivamente, para reduzir essa complexidade, uma estratégia natural é restringir a atenção a padrões esparsos que retêm apenas as conexões mais relevantes. Essas abordagens são bem estabelecidas em ARs, onde a atenção segue padrões esparsos fixos e claramente definidos. No entanto, em DLMs, observamos comportamentos de esparsidade distintos: (1) os padrões de atenção variam entre as cabeças, (2) os padrões de atenção em cada cabeça permanecem altamente semelhantes entre as etapas de desruído, e (3) as etapas iniciais de desruído são críticas para a geração. Essas descobertas tornam os métodos de atenção esparsa projetados para ARs amplamente incompatíveis com DLMs, pois falham em capturar estruturas específicas das cabeças e correm o risco de degradar a geração quando aplicados nas etapas iniciais de desruído. Para enfrentar esses desafios, propomos o SparseD, um novo método de atenção esparsa para DLMs. Aproveitando as observações, o SparseD requer apenas o pré-cálculo de padrões esparsos específicos para cada cabeça uma única vez, reutilizando-os em todas as etapas. Isso evita o recálculo de padrões esparsos a cada etapa de desruído. Enquanto isso, o SparseD usa atenção completa nas etapas iniciais e, em seguida, alterna para atenção esparsa mais tarde para manter a qualidade da geração. Juntos, esses aspectos estabelecem o SparseD como uma solução prática e eficiente para a implantação de DLMs em aplicações de contexto longo. Resultados experimentais demonstram que o SparseD alcança aceleração sem perdas, proporcionando um aumento de velocidade de até 1,50 vezes em relação ao FlashAttention em um contexto de 64k com 1.024 etapas de desruído.

English

While diffusion language models (DLMs) offer a promising alternative to autoregressive models (ARs), existing open-source DLMs suffer from high inference latency. This bottleneck is mainly due to the attention's quadratic complexity with respect to context length in computing all query-key pairs. Intuitively, to reduce this complexity, a natural strategy is to restrict attention to sparse patterns that retain only the most relevant connections. Such approaches are well-established in ARs, where attention follows fixed and clearly defined sparse patterns. However, in DLMs, we observe distinct sparsity behaviors: (1) attention patterns vary across heads, (2) attention patterns in each head remain highly similar across denoising steps, and (3) early denoising steps are critical for generation. These findings render sparse attention methods designed for ARs largely incompatible with DLMs, as they fail to capture head-specific structures and risk degrading generation when applied in early denoising steps. To address these challenges, we propose SparseD, a novel sparse attention method for DLMs. Leveraging the observations, SparseD only requires pre-computing head-specific sparse patterns one time, and reuses them across all steps. This prevents recomputing sparse patterns at each denoising step. Meanwhile, SparseD uses full attention in the early steps, then switches to sparse attention later to maintain generation quality. Together, these establish SparseD as a practical and efficient solution for deploying DLMs in long-context applications. Experimental results demonstrate that SparseD achieves lossless acceleration, delivering up to 1.50times speedup over FlashAttention at a 64k context length with 1,024 denoising steps.

SparseD: Atenção Esparsa para Modelos de Linguagem de Difusão

SparseD: Sparse Attention for Diffusion Language Models

Resumo

Support