SparseD: Sparse Attention voor Diffusie Taalmodellen

Samenvatting

Hoewel diffusie-taalmmodellen (DLMs) een veelbelovend alternatief bieden voor autoregressieve modellen (ARs), lijden bestaande open-source DLMs onder hoge inferentielatentie. Dit knelpunt wordt voornamelijk veroorzaakt door de kwadratische complexiteit van de aandacht (attention) ten opzichte van de contextlengte bij het berekenen van alle query-key paren. Intuïtief is een natuurlijke strategie om deze complexiteit te verminderen, het beperken van de aandacht tot sparse patronen die alleen de meest relevante verbindingen behouden. Dergelijke benaderingen zijn goed ingeburgerd in ARs, waar de aandacht vaste en duidelijk gedefinieerde sparse patronen volgt. In DLMs observeren we echter verschillende sparse gedragingen: (1) aandachtspatronen variëren tussen heads, (2) aandachtspatronen in elke head blijven sterk vergelijkbaar over de denoising-stappen heen, en (3) vroege denoising-stappen zijn cruciaal voor de generatie. Deze bevindingen maken sparse aandachtmethoden die zijn ontworpen voor ARs grotendeels ongeschikt voor DLMs, omdat ze head-specifieke structuren niet vastleggen en het risico lopen de generatie te verslechteren wanneer ze worden toegepast in vroege denoising-stappen. Om deze uitdagingen aan te pakken, stellen we SparseD voor, een nieuwe sparse aandachtmethode voor DLMs. Gebruikmakend van de observaties, vereist SparseD slechts eenmalig het vooraf berekenen van head-specifieke sparse patronen, en hergebruikt deze over alle stappen heen. Dit voorkomt het opnieuw berekenen van sparse patronen bij elke denoising-stap. Tegelijkertijd gebruikt SparseD volledige aandacht in de vroege stappen, en schakelt later over naar sparse aandacht om de generatiekwaliteit te behouden. Samen maken deze SparseD tot een praktische en efficiënte oplossing voor het inzetten van DLMs in lang-context toepassingen. Experimentele resultaten tonen aan dat SparseD versnelling zonder verlies bereikt, met een snelheidsverbetering tot 1,50 keer ten opzichte van FlashAttention bij een contextlengte van 64k met 1.024 denoising-stappen.

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: Sparse Attention voor Diffusie Taalmodellen

SparseD: Sparse Attention for Diffusion Language Models

Samenvatting

Support